import numpy as np
from gurobipy import *
import numpy as np
import datetime
from rsome import dro
import rsome.grb_solver as grb
import rsome as rso

# simulate the MDP process of a single courier
class COURIER:  # a single courier
    def __init__(self, conf):
        #self.Q = np.random.uniform(conf['Qmin'], conf['Qmax']) # capacity of the courier
        self.AreaLen = conf['AreaLen'] # the length of the area
        self.T = conf['T'] # T period 
        self.mode = conf['UtilityMode']
        #self.ProbOccu = conf['ProbOccu'] # whether it is occupied or not
        self.ErrorCorr = conf['ErrorCorr'] # the correlation of error terms
        
        self.PriceMin = conf['Prices'][0]
        self.PriceMax = conf['Prices'][-1]
        self.prices = conf['Prices']

        # the following are the parameters of the MDP
        self.sigma = conf['Sigma'] # the standard deviation of the utility function
        self._ErrorCov()
        self.error = np.random.multivariate_normal(np.zeros(self.T), self.cov)
        
        # States that will be updated as t goes
        self.BidPrices = None
        #self.capacity = self.Q
        self.loc = (np.random.uniform(0, self.AreaLen), np.random.uniform(0, self.AreaLen)) #(0, 0)

        # constant, distant, whether it is remote or not
        self.util_weights = np.array([np.random.uniform(0.5, 1.5)]) #np.array([np.random.uniform(0,1), np.random.uniform(0,1)]) #np.random.uniform(-5, 10)])
        #self.util_weights_remote = np.random.uniform(-5, 10)

    def _ErrorCov(self,):
        T = self.T
        cov = np.zeros((T, T))
        for t1 in range(T):
            for t2 in range(t1):
                cov[t1,t2] = self.ErrorCorr*(self.sigma[t1]*self.sigma[t2])
                cov[t2,t1] = cov[t1,t2]
            cov[t1,t1] = self.sigma[t1]**2
        self.cov = cov
    
    def utility_dist(self, loc, ori, des, t):
        pickup = np.sqrt((ori[0]-loc[0])**2 + (ori[1]-loc[1])**2)
        dist = np.sqrt((ori[0]-des[0])**2 + (ori[1]-des[1])**2)
        #direct_dist = np.sqrt((loc[0]-des[0])**2 + (loc[1]-des[1])**2)
        #return self.PriceMin + (dist-self.PriceMin)/(self.PriceMax-self.PriceMin) + self.error[t]
        
        return pickup + dist 
    
    def utility_random(self, t):
        return np.random.uniform(self.PriceMin, self.PriceMax) 
    
    def utility_random2(self, t):
        if np.random.uniform(0,1) < 0.5:
            return self.PriceMin 
        else:
            return self.PriceMax 

    def utility_complex(self, loc, ori, des, t):
        pickup = np.sqrt((ori[0]-loc[0])**2 + (ori[1]-loc[1])**2)
        dist = np.sqrt((ori[0]-des[0])**2 + (ori[1]-des[1])**2)
        #remote = 0
        #if des[0] > 0.75*self.AreaLen and des[1] > 0.75*self.AreaLen:
        #    remote = 1
        
        util = self.util_weights[0]*(pickup + dist) 
        
        return util

    def _CourierActions(self, dems, t):

        if True:
            calls_loc = dems.calls_loc
            J = dems.JobNums[t]
            bps = np.zeros(J)
            for j in range(J):
                courier_loc = self.loc
                ori, des = calls_loc[t, j]['ori'], calls_loc[t, j]['des']
                if self.mode == 'complex':
                    utility = self.utility_complex(courier_loc, ori, des, t) 
                elif self.mode == 'dist':
                    utility = self.utility_dist(courier_loc, ori, des, t)
                elif self.mode == 'random':
                    utility = self.utility_random(t)
                elif self.mode == 'random2':
                    utility = self.utility_random2(t)
            # bid the lowest price that is not less than utility

                if t in [0,7]:
                    utility *= 1.2
                elif t in [3,4]:
                    utility *= 0.8

                utility += self.error[t]

                for p in self.prices:
                    if p >= utility:
                        bps[j] = p
                        break
                if bps[j] == 0:
                    bps[j] = self.prices[-1]
            self.BidPrices = bps

    def _update_loc(self, bid):
        # update location based on the accepted bid
        if bid is not None: # didn't get the job 
            self.loc = bid['des'] #(des[0] ,des[1])
        # otherwise the loc didn't change


class DEMAND: # independently generated demand at T period
    def __init__(self, conf, JobNum_Dev):
        self.AreaLen = conf['AreaLen'] # the length of the area'
        self.T = conf['T']
        self.loc = [None] * self.T
        self.K = conf['K'] # number of different jobs
        
        self.JobNum_Corr = conf['JobCorr'] # correlation of job number
        self.JobNum_Lb = int(conf['JobNum_Lb']*(1- JobNum_Dev))
        self.JobNum_Ub = int(conf['JobNum_Ub']*(1+ JobNum_Dev))
        self.JobNum_Mid = (self.JobNum_Lb + self.JobNum_Ub)/2

        self._JobNum_cov() 
        self._JobNums()
        self._calls()

    def trunc(self,a,b,x):
        return min(b, max(a, x))
    
    def _JobNum_cov(self,):
        T = self.T
        std =  self.JobNum_Mid
        cov = np.zeros((T,T))
        for t1 in range(T):
            for t2 in range(t1):
                cov[t1,t2] = self.JobNum_Corr*(std**2)
                cov[t2,t1] = cov[t1,t2]
            cov[t1,t1] = std**2
        self.JobNum_cov = cov

    def _JobNums(self,): # normal truncated, mean at midpoint
        T = self.T
        temp = np.random.multivariate_normal(self.JobNum_Mid*np.ones(T), self.JobNum_cov) 
        self.JobNums = np.array([self.trunc(self.JobNum_Lb, self.JobNum_Ub, int(temp[t])) for t in range(T)]) 


    def _calls(self,):
        self.calls_loc = {}
        for t in range(self.T):
            for j in range(self.JobNums[t]):
                self.calls_loc[t, j] = {'ori': (np.random.uniform(0, self.AreaLen), np.random.uniform(0, self.AreaLen)),
                                        'des': (np.random.uniform(0, self.AreaLen), np.random.uniform(0, self.AreaLen))}
                                    #'capacity': np.random.randint(0, self.p)}

class simulation:
    def __init__(self, conf):
        # initialize K couriers
        self.conf = conf
        self.T = conf['T']
        self.K = conf['K']
        self.L = conf['L']
        self.prices = conf['Prices']
    
    def bid_assignment(self, CourJobPrice): # re-index the courier for those who are available
        K_t, J_t = CourJobPrice.shape

        model = Model('bid assginment')
        s = model.addVars(K_t, J_t, lb=0, vtype=GRB.CONTINUOUS)
        obj = LinExpr(0)
        for k in range(K_t):
            for j in range(J_t):
                obj.addTerms(CourJobPrice[k,j], s[k,j])
        for j in range(J_t):
            model.addConstr(quicksum(s[k,j] for k in range(K_t)) == 1)
        for k in range(K_t):
            model.addConstr(quicksum(s[k,j] for j in range(J_t)) <= 1)
        
        #model.addConstr(quicksum(s[k,j] for k in range(K_t) for j in range(J_t)) >= min(K_t, J_t))
        model.setObjective(obj, GRB.MINIMIZE)
        obj.clear()
        model.setParam("OutputFlag", 0)
        model.optimize()
        
        # obtain solution and z
        z = np.zeros(self.N)
        ss = np.zeros((K_t, J_t))
        for k in range(K_t):
            for j in range(J_t):
                ss[k,j] = np.round(s[k,j].X)
                for n in range(self.N):
                    if self.prices[n] == CourJobPrice[k,j]:
                        z[n] += s[k,j].X
        return z, ss
    
    def trunc(self,a,b,x):
        return min(b, max(a, x))

    def _courier_avail_list(self,):
        mean = self.K/2
        avail_courier_num = self.trunc(0, self.K, int(np.random.normal(mean, mean)))
        courier_avail_list = np.random.choice(self.K, avail_courier_num, replace=False)
        return courier_avail_list
     

    def generate_Samples(self, Samp_No, JobNum_dev=0):
        #print(Samp_No)
        samples = {}
        for omega in range(Samp_No):    
            # generate a sample path of calls and couriers
            dems = DEMAND(self.conf, JobNum_dev)
            samples[omega] = {'JobNum':dems.JobNums}
            CourNum = np.zeros(self.T) # number of available couriers
            
            # initiate couriers
            couriers = {}
            CourJobPrice = {}
            z = {}
            for k in range(self.K):
                couriers[k] = COURIER(self.conf)
            
            # simulate in each period t
            for t in range(self.T):
                courier_avail_list = self._courier_avail_list()
                for k in range(self.K):
                    if k in courier_avail_list:
                        couriers[k]._CourierActions(dems, t) 
                    else:
                        couriers[k].loc = (np.random.uniform(0, self.conf['AreaLen']), np.random.uniform(0, self.conf['AreaLen']))
                        couriers[k].BidPrices = None

                #for k in range(self.K):
                    # give calls for each courier to bid
                #    couriers[k]._CourierActions(dems, t) 
                #    if couriers[k].BidPrices is not None:
                #        courier_avail_list.append(k)

                # generate bidding set (couriers, job, price)
                K_t, J_t = len(courier_avail_list), dems.JobNums[t]
                #print('available couriers %s, demand is %s'%(K_t, J_t))
                CourJobPrice[t] = np.zeros((K_t, J_t))
                for kk in range(len(courier_avail_list)):
                    CourJobPrice[t][kk,:] = couriers[courier_avail_list[kk]].BidPrices
                if K_t < J_t:  
                    # add third party bidder with highest price
                    CourJobPrice[t] = np.concatenate((CourJobPrice[t], self.prices[-1]* np.ones((J_t-K_t, J_t))), axis=0)
                    K_t = J_t
                CourNum[t] = K_t

                # assign jobs
                z[t], ss = self.bid_assignment(CourJobPrice[t])
                # update courier location based on assignment
                for kk in range(len(courier_avail_list)):
                    if sum(ss[kk]) == 1: # courier kk is assigned with jobs
                        temp = np.nonzero(ss[kk])
                        j = temp[0][0]
                        bid = dems.calls_loc[t,j]
                        couriers[courier_avail_list[kk]]._update_loc(bid)        
            samples[omega].update({'CourNum':CourNum, 'z': z, 'CourJobPrice': CourJobPrice})
        #print(samples[0])
        return samples
           
class optimization(simulation):
    def __init__(self, conf):  # 先继承，在重构
        simulation.__init__(self, conf)  # 继承父类的构造
        self.InSamNo = conf['InSamNo']         # number of in-samples
        self.OutSamNo = conf['OutSamNo']       # number of out-of-samples
        self.N = int(conf['PriNum'])
        self.a = np.array([[1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0, 0, 0], [1, 1, 1, 1, 0, 0, 0, 0],
                           [0, 1, 0, 0, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 0, 0], [0, 1, 1, 1, 0, 0, 0, 0], [0, 1, 1, 1, 1, 0, 0, 0],
                           [0, 0, 1, 0, 0, 0, 0, 0], [0, 0, 1, 1, 0, 0, 0, 0], [0, 0, 1, 1, 1, 0, 0, 0], [0, 0, 1, 1, 1, 1, 0, 0],
                           [0, 0, 0, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 0, 0, 0], [0, 0, 0, 1, 1, 1, 0, 0], [0, 0, 0, 1, 1, 1, 1, 0],
                           [0, 0, 0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 1, 1, 0, 0], [0, 0, 0, 0, 1, 1, 1, 0], [0, 0, 0, 0, 1, 1, 1, 1],
                           [0, 0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1, 1, 0], [0, 0, 0, 0, 0, 1, 1, 1],
                           [0, 0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 1, 1],
                           [0, 0, 0, 0, 0, 0, 0, 1]]) # 26 x 8
                           
        self.I, self.T = 26, 8 #self.a.shape
        self.w = np.zeros(self.I)
        for i in range(self.I):
            NumPer = np.sum(self.a, axis=1)[i]     # the number of periods covered by a shift
            self.w[i] =  self.conf['ScheDri_Wage']*NumPer * (self.conf["Dis_Coff"] ** (NumPer-1))

        # generate samples, breakpoints, and upperbound of z ... after sample realization
        self.InSamples = None  #generate_Samples()
        self.OutSamples = None #generate_Samples()   
        self.BrePoints = None     
        self.z = None         # reduced information vector obtained from solving assignment problem
        self.z_hat = None     # size = (L, T), used for robust model
        self.kappa_bound = None

        # robust
        self.q = np.zeros((self.N,self.N))
        for n in range(self.N):
            for m in range(self.N):
                if m < n:
                    self.q[n,m] = self.prices[m]
                else:
                    self.q[n,m] = self.prices[n]
        self.search_epsilon = 1e-4   # convergence gap for bisection search

    
    def bisection(self, tau):
        epsilon = self.search_epsilon
        kappa_lb = 0
        kappa_ub = self.kappa_bound
        w = self.w
            
        k = kappa_ub
        x, rho_val = self.rho(k)
        print('k  %.4f,'%k,'rho %.4f'%rho_val)
        print('tau',tau)
        print('rho', rho_val)
        if rho_val - tau > 1e-4:
            print('upper bound is wrong')
            return None, -1, None
        else:
            k = (kappa_lb + kappa_ub)/2
        feasibility = 0
        
        iters = 0
        while(kappa_ub-kappa_lb > epsilon):
            x_try, rho_val = self.rho(k)   # rho is decreasing in k
            print('iteration %s'%iters)
            print('k  %.4f,'%k,'rho %.4f'%rho_val)
            if rho_val <= tau:
                feasibility = 1
                x = x_try
                kappa_ub = k
                k = (kappa_lb + kappa_ub) / 2
            else:
                kappa_lb = k
                k = (kappa_lb + kappa_ub) / 2
            print('lb %.4f, up %.4f'%(kappa_lb, kappa_ub))
            iters += 1

        if feasibility == 0:
            return None, -1, None

        return x, kappa_ub, w@x

    def est_z_hat(self, ):
        # z  L x SamNo x T x N
        z_hat = np.max(np.sum(self.z, axis = 3), axis = 1)
        #print(z_hat)
        
        # to make sure it is not zero
        z_hat[z_hat == 0] = 0.001
        return  z_hat # L x T
    
    def est_kappa_bound(self,):
        ub_bound = np.sum(self.prices)*np.max(self.z_hat[0,:]) # \hat z_{tl} has maximum when l=0 since 1^\top z_{tl} = J_t - u_{tl}
        return ub_bound

    
    def alpha(self, k, l, omega, t, n):
        model = Model('alpha')
        z = model.addVars(self.N, lb = 0, vtype = GRB.CONTINUOUS) 
        h = model.addVar(lb = 0, vtype = GRB.CONTINUOUS)       # an auxiliary variable for the inf-norm term
        for m in range(self.N):
            model.addConstr( h >= z[m] - self.z[l, omega, t, m])
            model.addConstr( h >= -z[m] + self.z[l, omega, t, m])
        model.addConstr(quicksum(z) <= self.z_hat[l, t] )

        #print('z', self.z_hat[l,t])
        model.setObjective(quicksum(self.q[n,m]*z[m] for m in range(self.N))-k*(1/self.z_hat[l,t])*h, GRB.MAXIMIZE)
        model.setParam("OutputFlag", 0)
        model.optimize()

        return model.ObjVal

    def rho(self, k):
        L = self.L
        I = self.I
        SamNo = self.InSamNo
        T = self.T
        w = self.w
        a = self.a
        PriNum = self.N
        BrePoints = self.BrePoints
        price = self.prices

        alpha_val = np.zeros((L, SamNo, T, PriNum))
        for l in range(L):
            for omega in range(SamNo):
                for t in range(T):
                    for n in range(PriNum):
                        alpha_val[l, omega, t, n] = self.alpha(k, l, omega, t, n)
        
        model = Model('rho')
        x = model.addVars(I, lb = 0, vtype = GRB.CONTINUOUS, name='x')
        gamma = model.addVars(L, T, lb = 0, vtype = GRB.CONTINUOUS, name='gamma')
        eta = model.addVars(L, T, lb = 0, vtype = GRB.CONTINUOUS, name='eta')
        r = model.addVars(L, SamNo, T, lb = 0, vtype = GRB.CONTINUOUS, name = 'r')

        # define objective function
        obj = LinExpr(0)
        for i in range(I):
            obj.addTerms(w[i], x[i])
        for l in range(L):
            for omega in range(SamNo):
                for t in range(T):
                        obj.addTerms(1/SamNo, r[l, omega, t])
        model.setObjective(obj, GRB.MINIMIZE)
        obj.clear()
        # define constraints
        for t in range(T):
            model.addConstr(quicksum(eta[l,t] for l in range(L)) == quicksum(a[i,t]*x[i] for i in range(I)))
            model.addConstr(quicksum(gamma[l,t] for l in range(L)) == 1)
            for l in range(L):
                for omega in range(SamNo):
                    for n in range(PriNum):
                        model.addConstr(r[l,omega,t] >= alpha_val[l,omega,t,n]*gamma[l,t] - (eta[l,t] - BrePoints[l,t]*gamma[l,t])*price[n])
            
        model.setParam("OutputFlag", 0)
        model.optimize()
        obj_val = model.ObjVal
        x_val = np.zeros(I)
        for i in range(I):
            x_val[i] = x[i].X    # need to check whether it can be relaxed to continuous

        return x_val, obj_val

    def generate_BrePoints(self,):
         # breakpoints
        L = self.L
        T = self.T  
        SamNo = self.InSamNo
        Samples = self.InSamples

        BrePoints = np.zeros((L,T))       # the first breakpoint is 0
        for t in range(T):
            J_t = 0
            for omega in range(SamNo):
                J_t += Samples[omega]["JobNum"][t]
            if L > 1:
                BrePoints[L-1,t] = int((J_t/SamNo)*0.5)

        return BrePoints
    
    def generate_BrePoints_exact(self, x_oriemp):
        L = self.L
        T = self.T
        I = self.I  
        a = self.a
        BrePoints = np.zeros((L,T))
        for t in range(T):
            BrePoints[1, t] = np.sum(a[i][t]*x_oriemp[i] for i in range(I))
        return BrePoints

    def generalized_assignment_problem(self, Samples):
        T = self.T 
        L = self.L    # number of break points is 2 for each t
        price = self.prices
        PriNum = self.N
        BrePoints = self.BrePoints
        SamNo = len(Samples)
        z_val = np.zeros((L, SamNo, T, PriNum)) 

        """
        Solve the model on page 11 for each sample and each period to get the value of reduced information vector
        """
        for l in range(L):
            for omega in range(SamNo):
                for t in range(T):
                    K_t = int(Samples[omega]["CourNum"][t])
                    J_t = int(Samples[omega]["JobNum"][t])
                    driver_job_price = Samples[omega]["CourJobPrice"][t]

                    model = Model('Gen_Model')
                    # s = model.addMVar((K_t,J_t), lb=0, vtype= GRB.CONTINUOUS)
                    s = model.addVars(K_t, J_t, lb=0, vtype=GRB.CONTINUOUS)
                    obj = LinExpr(0)
                    for k in range(K_t):
                        for j in range(J_t):
                            obj.addTerms(driver_job_price[k,j], s[k,j])
                
                    for k in range(K_t):
                        model.addConstr(quicksum(s[k,j] for j in range(J_t)) <= 1)
                    for j in range(J_t):
                        model.addConstr(quicksum(s[k,j] for k in range(K_t)) <= 1)
                    model.addConstr(quicksum(s[k,j] for k in range(K_t) for j in range(J_t)) >= J_t - BrePoints[l,t])
                    
                    model.setObjective(obj, GRB.MINIMIZE)
                    obj.clear()
                    model.setParam("OutputFlag", 0)
                    model.optimize()
                # compute z
                    for p in range(PriNum):
                        for k in range(K_t):
                            for j in range(J_t):
                                if price[p] == driver_job_price[k,j]:
                                    z_val[l,omega,t,p] += s[k,j].X
        return z_val  # L x Omega x T x N
    
    def generalized_empirical(self,):
        L = self.L
        I = self.I
        T = self.T
        a, w = self.a, self.w
        price = self.prices
        z = self.z
        SamNo = self.InSamNo
        PriNum = self.N
        BrePoints = self.BrePoints
        
        model = Model('Model15')
        # define variables
        x = model.addVars(I, lb=0, vtype=GRB.CONTINUOUS, name='x')     # check it can be relaxed to continuous
        y = model.addVars(L, SamNo, T, PriNum, lb=0, vtype=GRB.CONTINUOUS, name='y')
        gamma = model.addVars(L, T, lb=0, vtype=GRB.CONTINUOUS, name='gamma')
        eta = model.addVars(L, T, lb=0, vtype=GRB.CONTINUOUS, name='eta')

        # define objective function
        obj = LinExpr(0)
        for i in range(I):
            obj.addTerms(w[i], x[i])
        for l in range(L):
            for omega in range(SamNo):
                for t in range(T):
                    for p in range(PriNum):
                        obj.addTerms(1/SamNo * price[p], y[l,omega,t,p])
        model.setObjective(obj, GRB.MINIMIZE)
        obj.clear()

        # define constraints
        for t in range(T):
            for l in range(L):
                for omega in range(SamNo):
                    model.addConstr(quicksum(y[l,omega,t,p] for p in range(PriNum)) >= sum(z[l,omega,t,:])*gamma[l,t]
                                    - eta[l,t] + BrePoints[l,t]*gamma[l,t])
                    for p in range(PriNum):
                        model.addConstr(y[l,omega,t,p] <= z[l,omega,t,p]*gamma[l,t])

            model.addConstr(quicksum(eta[l,t] for l in range(L)) == quicksum(a[i][t]*x[i] for i in range(I)))
            model.addConstr(quicksum(gamma[l,t] for l in range(L)) == 1)

        # optimize
        model.setParam("OutputFlag", 0)
        model.optimize()
        obj_val = model.ObjVal
        x_val = [x[i].X for i in range(I)]

        return obj_val, x_val, w@x_val

    def origin_empirical(self, ):
        L = self.L
        I = self.I
        T = self.T
        a, w = self.a, self.w
        SamNo = self.InSamNo
        Samples = self.InSamples
        
        model = Model('Model4')
        # define variables
        x = model.addVars(I, lb=0, vtype=GRB.CONTINUOUS, name='x')     # check it can be relaxed to continuous
        s = {}

        # define objective function
        obj = LinExpr(0)
        for i in range(I):
            obj.addTerms(w[i], x[i])
        
        for omega in range(SamNo):
            for t in range(T):
                K_t = int(Samples[omega]["CourNum"][t])
                J_t = int(Samples[omega]["JobNum"][t])
                s[omega, t] = model.addVars(K_t, J_t, lb=0, vtype=GRB.CONTINUOUS)
                driver_job_price = Samples[omega]["CourJobPrice"][t]
                for k in range(K_t):
                    for j in range(J_t):
                        #print(s[omega,t][k,j])
                        obj.addTerms((1./SamNo)*driver_job_price[k,j], s[omega, t][k ,j])
        
                for k in range(K_t):
                    model.addConstr(quicksum(s[omega,t][k,j] for j in range(J_t)) <= 1)
                for j in range(J_t):
                    model.addConstr(quicksum(s[omega,t][k,j] for k in range(K_t)) <= 1)
                
                model.addConstr(quicksum(s[omega,t][k,j] for k in range(K_t) for j in range(J_t)) >= J_t - quicksum(a[i][t]*x[i] for i in range(I)))

        model.setObjective(obj, GRB.MINIMIZE)

        # optimize
        model.setParam("OutputFlag", 0)
        model.optimize()
        obj_val = model.ObjVal
        x_val = [x[i].X for i in range(I)]

        return obj_val, x_val, w@x_val
    
    
    def est_moment_info(self,):

        mu = np.mean(self.z[0, :, :, :], axis = 0) # (T, N)

        #deviation = np.max( np.abs(self.z[0, :, :, :] - mu), axis = 0) # (T, N)
        
        #deviation = np.std(self.z[0, :, :, :], axis = 0 ) # standard deivation
        deviation = np.mean( np.abs(self.z[0, :, :, :] - mu), axis = 0)

        self.mu = np.array([mu[t,n] for t in range(self.T) for n in range(self.N)])
        self.deviation = np.array([deviation[t,n] for t in range(self.T) for n in range(self.N)])


    def moment_dro(self, ):
        L = self.L
        I = self.I
        T = self.T
        a, w = self.a, self.w
        price = self.prices
        price_copy = np.tile(price, T)
        PriNum = self.N

        
        self.est_moment_info()
        zhat = self.est_z_hat()
        mu = self.mu #(T, N)
        deviation = self.deviation #(T, N)
        
        # modeling with RSOME
        model = dro.Model()                                  # define a DRO model
        
        z = model.rvar(T*PriNum)                            # random variable z = (z_1, z_2, ..., z_T)
        u = model.rvar(T*PriNum)                            # auxiliary variable u

        fset = model.ambiguity()   
        fset.exptset(rso.E(z) == mu, rso.E(u) == deviation)
        fset.suppset(z >=0, u >= 0, z-mu <= u, mu - z <= u, [z[t*PriNum:(t+1)*PriNum].sum() <=  z_hat[0, t] for t in range(T)])


        x = model.dvar(I)
        y = model.dvar(T*PriNum)
        y.adapt(z)
        y.adapt(u)

        model.minsup(w@x + rso.E(price_copy@y), fset)
        for t in range(T):
            model.st( y[t*PriNum:(t+1)*PriNum].sum() >= z[t*PriNum:(t+1)*PriNum].sum()  - (a[:,t]*x).sum() )
        model.st( 0<= y, y<= z)

        model.st(x >=0)
        model.solve(grb)

        model.solve(grb)
        x_val = x.get()
        return x_val, w@x



    def out_of_sample(self, x_in):
        start = datetime.datetime.now()
        # only need to evaluate model (1) for each omega
        I = self.I
        T = self.T
        a, w = self.a, self.w
        SamNo = self.OutSamNo
        Samples = self.OutSamples
        
        evaluation = w@x_in      # first-stage cost
        for omega in range(SamNo):
            model = Model('f(x,B_t)^\omega-test-%s'%omega)
            obj = LinExpr(0)
            s = {}
            for t in range(T):
                K_t = int(Samples[omega]["CourNum"][t])
                J_t = int(Samples[omega]["JobNum"][t])
                s[t] = model.addVars(K_t, J_t, lb=0.0, vtype= GRB.CONTINUOUS, name='s')
                driver_job_price = Samples[omega]["CourJobPrice"][t]
                for k in range(K_t):
                    for j in range(J_t):
                        obj.addTerms(driver_job_price[k,j], s[t][k,j])
                        
                for k in range(K_t):
                    model.addConstr(quicksum(s[t][k,j] for j in range(J_t)) <= 1)
                for j in range(J_t):
                    model.addConstr(quicksum(s[t][k,j] for k in range(K_t)) <= 1)
                model.addConstr(quicksum(s[t][k,j] for k in range(K_t) for j in range(J_t)) >= J_t - a[:,t]@x_in)

            model.setObjective(obj, GRB.MINIMIZE)
            obj.clear()
            model.setParam("OutputFlag", 0)
            model.optimize()
            evaluation += (1/SamNo)*model.ObjVal
        #print(evaluation)
        end = datetime.datetime.now()
        return evaluation, end-start

