import numpy as np
from src.simu import *
import rsome as rso
from rsome import dro
import sys

print("Script Name:", sys.argv[0])
print("First argument:", sys.argv[1])  
print("Second argument:", sys.argv[2])


Config = {'K': 60, 'T': 8,
        'AreaLen': 40,
        'JobNum_Lb': 10, 'JobNum_Ub': 50,
        'dev': 0,
        'UtilityMode': 'dist', #'dist', #  #'random'
        'ErrorCorr': 0,
        'JobCorr': 0.0, 
        'ScheDri_Wage': 20, 'Dis_Coff': 0.9,
        'PriNum': 5, 'Prices': [20, 30, 40, 50, 60],
        'Sigma': [5,5,5,5,5,5,5,5],
        'L': 2, 'BreakPointMode': 'B',
        'InSamNo': int(sys.argv[1]), 'OutSamNo': 200}


def run(seed):

    np.random.seed(seed)
    data = optimization(Config)
    print('generating training samples...')
    data.InSamples = data.generate_Samples(data.InSamNo)
    print('generating test samples...')
    data.OutSamples = data.generate_Samples(data.OutSamNo, Config['dev'])
    data.BrePoints = np.zeros((data.L, data.T))


    data.z = data.generalized_assignment_problem(data.InSamples)
    data.est_moment_info()
    data.z_hat = data.est_z_hat()

    #empirical
    obj_emp, x_emp, firstCos_emp = data.generalized_empirical()
    eva_emp, _ = data.out_of_sample(x_emp)
    
    
    # RS
    print('Solving robust model')
    #data.z_hat = data.est_z_hat()
    data.kappa_bound = data.est_kappa_bound()
    print('data kappa up bound', data.kappa_bound)
    tau =  1.02 * obj_emp
    x_rs, kappa, firstCos_RS = data.bisection(tau)
    eva_rs, _ = data.out_of_sample(x_rs)

    # RSOME
    I = data.I
    T = data.T
    a, w = data.a, data.w
    price = data.prices
    price_copy = np.tile(price, T)
    PriNum = data.N

    mu = data.mu #(T, N)
    #print(mu)
    deviation = data.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() <=  data.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)
    
    x_moment = x.get()
    print(x_moment)
    eva_moment, _ = data.out_of_sample(x_moment)
    
    print('seed %s | moment %.3f | emp %.3f | %.3f%%' %
          (seed, eva_moment, eva_emp, 100*(eva_moment-eva_emp)/eva_emp ))
    
    
    return  eva_emp, eva_rs, eva_moment


SeedStart = int(sys.argv[2])
SeedNum = 10
name = 'RSOME/B_sample%s_seed%s-%s'%(Config['InSamNo'], SeedStart*SeedNum, (SeedStart+1)*SeedNum-1)

results = np.zeros((SeedNum, 3))
for s in range(SeedNum):
    rdseed = SeedStart*SeedNum+s
    results[s] = run(rdseed)

np.save('results/%s.npy'%name, results)
