from rsome import ro
import numpy as np
from rsome import grb_solver as grb
from gurobipy import*
import math

def compute_function_f(K_t, JobNum, driver_job_price, SchNum):

    model = ro.Model('functionF')
    s = model.dvar((K_t, JobNum))

    model.min((driver_job_price*s).sum())

    model.st(s.sum(axis=0) <= 1)        # each job can be assigned to at most one driver
    model.st(s.sum(axis=1) <= 1)        # each driver can be assigned at most one job
    model.st(s.sum() >= JobNum - SchNum)   # the total number of assigned jobs
    model.st(s>=0)

    model.solve(grb)
    obj_val = model.get()

    return obj_val

def compute_function_g(K_t, JobNum, driver_job_price, SchNum, PriNum, price):

    model = ro.Model('Model5')
    s = model.dvar((K_t, JobNum))

    model.min((driver_job_price*s).sum())

    model.st(s.sum(axis=0) == 1)       # each job should be assigned to one driver
    model.st(s.sum(axis=1) <= 1)       # each driver can be assigned at most one job
    model.st(s >= 0)

    model.solve(grb)
    obj_val = model.get()
    s_val = s.get()

    # compute z_{tn}
    z = np.zeros(PriNum)
    for p in range(PriNum):
        for k in range(K_t):
            for j in range(JobNum):
                if price[p] == driver_job_price[k,j]:
                    z[p] += s_val[k,j]
                    # if s_val[k,j] > 0:
                    #     print("driver ", k, " ---> ", j, ": ", price[p])


    model_g = ro.Model('g_function')
    y = model_g.dvar(PriNum)

    model_g.min(y@np.transpose(price))

    model_g.st(y.sum()>=z.sum()-SchNum)
    model_g.st(y<=z)
    model_g.st(y>=0)

    model_g.solve(grb)
    obj_val_g = model_g.get()
    y_value = y.get()
    # print("obj with s: ", obj_val)
    # print("z value: ", z)
    # print("basic solutions: ", y_value, obj_val_g)

    return obj_val_g, z

def compute_function_h(z_value, K_t, JobNum, driver_job_price, SchNum, PriNum, price):
    """
    We consider L=2 breakpoints, u1=0, u2=SchNum/2. For each l \in [L], we solve a model to find the value of s.
    """
    L = 2
    BreakPoints = [0, int(JobNum/2)]
    z = np.zeros((L,PriNum))

    z[0,:] = z_value[:]
    # solve the model for breakpoint 1
    model = ro.Model('Model10')
    s = model.dvar((K_t, JobNum))

    model.min((driver_job_price*s).sum())

    model.st(s.sum(axis=0) <= 1)       # each job should be assigned to at most driver
    model.st(s.sum(axis=1) <= 1)       # each driver can be assigned at most one job
    model.st(s.sum() >= JobNum - BreakPoints[L-1])  # the total number of assigned jobs
    model.st(s >= 0)

    model.solve(grb)
    obj_val = model.get()
    s_val = s.get()
    # compute z_{tn} for breakpoint 1
    for p in range(PriNum):
        for k in range(K_t):
            for j in range(JobNum):
                if price[p] == driver_job_price[k,j]:
                    z[L-1,p] += s_val[k,j]

    model_h = ro.Model('h_function')
    y = model_h.dvar((L,PriNum))
    gamma = model_h.dvar(L)
    eta = model_h.dvar(L)

    model_h.min((y@np.transpose(price)).sum())

    for l in range(L):
        model_h.st(y[l,:].sum()>=(z[l,:].sum())*gamma[l]-(eta[l]-BreakPoints[l]*gamma[l]))
        model_h.st(y[l,:]<=z[l,:]*gamma[l])

    model_h.st(eta.sum() == SchNum)
    model_h.st(gamma.sum() == 1)
    model_h.st(y>=0, gamma>=0, eta>=0)

    model_h.solve(grb)
    obj_val_h = model_h.get()

    y_value = y.get()
    gamma_value = gamma.get()
    eta_value = eta.get()
    # print(z)
    # print("generalized of y: ", y_value)
    # print("generalized of gamma: ", gamma_value)
    # print("generalized of eta: ", eta_value)

    return obj_val_h


