import numpy as np
from src.simu import *
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': sys.argv[1],#'B',
        'InSamNo': 7, 'OutSamNo': 200}

SeedStart = int(sys.argv[2])
SeedNum = 10

lenTar = 16
Target = [1+0.01*i for i in range(lenTar)]
test_dev = [0, -0.1, 0.1, -0.2, 0.2]

name = 'main/%s_sample%s_seed%s-%s_tar%s'%(Config['BreakPointMode'], Config['InSamNo'], SeedStart*SeedNum, (SeedStart+1)*SeedNum-1, lenTar)

Emp_x = {}
RS_x = {}
Emp_target = np.zeros(SeedNum)
RS_target = np.zeros((SeedNum, lenTar))
Emp_test = np.zeros(SeedNum)
RS_test = np.zeros((SeedNum, lenTar))

RS_dev_test = np.zeros((SeedNum, len(test_dev), lenTar))

for s in range(SeedNum):
    rdseed = SeedStart*SeedNum+s
    print('====================================== random instance %s ======================================'%rdseed)
    np.random.seed(rdseed)
    print(Config)
    data = optimization(Config)
        
    # generating samples
    print('generating training samples...')
    data.InSamples = data.generate_Samples(data.InSamNo)
    print('generating test samples...')
    data.OutSamples = data.generate_Samples(data.OutSamNo)
    
    # generate break points and z
    if Config['BreakPointMode'] == 'B':
        data.BrePoints = np.zeros((data.L, data.T))
    elif Config['BreakPointMode'] == 'G1':
        data.BrePoints = data.generate_BrePoints() 
    elif Config['BreakPointMode'] == 'G2':
        print('solving exact empirical model...')
        _, x_exact_emp, firstCos_exact_emp = data.origin_empirical()
        data.BrePoints = data.generate_BrePoints_exact(x_exact_emp)

    print('generating reduced information vectors...')
    data.z = data.generalized_assignment_problem(data.InSamples) 

    print('solving generalized empirical model...')
    Emp_target[s], x_emp, firstCos_emp = data.generalized_empirical()
    print('out-of-sample test for EMP...')
    Emp_test[s], _ = data.out_of_sample(x_emp)

    print('EMP solution for seed %s'%(rdseed))
    Emp_x[s] = x_emp
    print(x_emp)
    
    # robust 
    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)

    RS_target[s, 0] = Emp_target[s]
    RS_x[s, 0] = x_emp
    RS_test[s, 0] = Emp_test[s]

    for i in range(1, lenTar):
        tau =  Target[i] * Emp_target[s]
        RS_target[s, i] = tau
        x_rs, kappa, firstCos_RS = data.bisection(tau)
        print('RS solution for seed %s target %.2f'%(rdseed, Target[i]))
        print(x_rs)
        RS_x[s, i] = x_rs
        print('out-of-sample test for RS...')
        RS_test[s, i], _ = data.out_of_sample(x_rs)
    
    if Config['BreakPointMode'] == 'G2':
        for i in range(lenTar):
            RS_dev_test[s, 0, i] = RS_test[s, i]

        for j in range(1, len(test_dev)):
            print('generating test samples for deviation %s...'%test_dev[j])
            data.OutSamples = data.generate_Samples(data.OutSamNo, test_dev[j])
            #Emp_dev_test[s, j], _ = data.out_of_sample(x_emp)
            #print(x_emp)
            for i in range(lenTar):
                RS_dev_test[s, j, i], _ = data.out_of_sample(RS_x[s, i])


            
    np.save('results/%s_RS_sol.npy'%name, [[RS_x[ss, ii] for ii in range(lenTar)] for ss in range(s+1)])
    np.save('results/%s_RS_target.npy'%name, RS_target)
    np.save('results/%s_RS_test.npy'%name, RS_test)
    if Config['BreakPointMode'] == 'G2':
        np.save('results/%s_RS_dev_test.npy'%name, RS_dev_test)