

import numpy as np
import torch
import pickle
import os

import pandas as pd

from tools import loss_func_tools
from tools import data_generation_tools
from tools import iceo_framework
from tools import optimization_oracle_tools
from tools import baseline_tools

def newsvendor_test(problem_params):
    pred_model = problem_params['pred_model']
    n_train = problem_params['train_sample_size']
    n_test = problem_params['test_sample_size']
    dim_features = problem_params['dim_features']
    dim_sol = problem_params['dim_sol']
    dim_scenario = problem_params['dim_scenario']
    n_trails = problem_params['n_trails']
    epoch_num = problem_params['epoch_num']
    lr = problem_params['learning_rate']
    has_entropy = problem_params['entropy_loss']
    has_reg = problem_params['iceo_reg']
    # generate xz_samples
    # B_true = data_generation_tools.generate_true_hypothesis(problem_params, seed=456)
    B_true = data_generation_tools.generate_true_hypothesis_NN(problem_params)
    B_true_matrix = data_generation_tools.generate_true_hypothesis(problem_params, seed=456)
    print("B_true", B_true)
    print("B_true_matrix", B_true_matrix)

    if problem_params['baseline'] is None:
        test_results = pd.DataFrame(
        columns=['i', 'train_size', 'pred_model', 'approx_oracle_mse', 'train_loss',
                                   'test_loss_approx', 'test_loss'])
    else:
        test_results = pd.DataFrame(
            columns=['i', 'train_size', 'pred_model', 'approx_oracle_mse', 'train_loss',
                                    'test_loss_approx','test_loss']+problem_params['baseline'])
    if pred_model == 'linear':
        hidden_dim = problem_params.get('hidden_dim', 128)
        predict_model = iceo_framework.linear_pred(dim_features, dim_scenario,hidden_dim)
    elif pred_model == 'two_layers':
        hidden_dim = problem_params.get('hidden_dim', 128)
        predict_model = iceo_framework.two_layers_pred(dim_features, dim_scenario,hidden_dim)
    elif pred_model == 'two_layers_small':  
        hidden_dim = problem_params.get('hidden_dim', 32)
        predict_model = iceo_framework.two_layers_pred(dim_features, dim_scenario,hidden_dim)

    dim_p = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    pickle_dir = problem_params['oracle_pickle_dir']
    fullname = os.path.join(pickle_dir, 'oralce_output' + 'dim_p' + str(
            dim_p) + 'dim_w' + str(dim_w) + '.pkl')
    # for rho not equal to 0.01
    fullname = os.path.join(pickle_dir, 'oralce_output' + 'dim_p' + str(dim_p) + 'dim_w' + str(dim_w) +'sample'+str(100)+'.pkl')
    #
    oracle_output = pickle.load(open(fullname, "rb"))
    if problem_params['oracle_kernel'] == 'nn':
        oracle_model = oracle_output['oracle_model']
        mse = oracle_output['mse']
        for param in oracle_model.parameters():
            param.requires_grad = False
    elif problem_params['oracle_kernel']=='polynomial':
        p_train = oracle_output['p_train']
        p_train = torch.as_tensor(p_train, dtype=torch.float64)
        dual_coef = oracle_output['dual_coef']
        oracle_params = oracle_output['oracle_params']
        mse = oracle_output['mse']
    else:
        raise ValueError("Invalid oracle kernel!")
    print("approx oracle mse", mse)
# generate xz_samples
    for i in range(n_trails):
        if problem_params['data_generation_model'] == 'two_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        elif problem_params['data_generation_model'] == 'model mis-specification-deg':
            print('deg data',problem_params['deg_data'])
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true_matrix)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true_matrix)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true_matrix)
        elif problem_params['data_generation_model'] == 'multi_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        else:
            raise ValueError("ERROR! Invalid data_generation_model!!")

        x_train = torch.as_tensor(x_train0, dtype=torch.float32)
        x_test = torch.as_tensor(x_test0, dtype=torch.float32)
        z_train = torch.as_tensor(z_train0, dtype=torch.float32)
        z_test = torch.as_tensor(z_test0, dtype=torch.float32)

        x_valid = torch.as_tensor(x_valid0, dtype=torch.float32)
        z_valid = torch.as_tensor(z_valid0, dtype=torch.float32)

        def z_idx(problem_params, z_data):
            size = np.shape(z_data)[0]
            dim_scenario = problem_params['dim_scenario']
            scenario_list = problem_params['scenario_list']
            z_index = torch.zeros(size)
            for i in range(size):
                for k in range(dim_scenario):
                    # print(z_real[i].numpy())
                    # print(scenario_list[k])
                    if (z_data[i] == scenario_list[k, :]).all():
                        z_index[i] = k
            return z_index
        z_train_idx0 = z_idx(problem_params, z_train0)
        # print("z_test_idx shape",  z_train_idx0)
        z_test_idx0 = z_idx(problem_params, z_test0)
        z_valid_idx0 = z_idx(problem_params, z_valid0)

        z_train_idx = torch.as_tensor(z_train_idx0, dtype=torch.float32)
        z_test_idx = torch.as_tensor(z_test_idx0, dtype=torch.float32)
        z_valid_idx = torch.as_tensor(z_valid_idx0, dtype=torch.float32)
        
       
        ICEO_model = iceo_framework.ICEO_model(problem_params, predict_model,oracle_model)
        nv_loss = iceo_framework.nv_loss_func(problem_params)
        iceo_loss =  iceo_framework.iceo_loss_func(problem_params)
        entropy_loss = loss_func_tools.entropy_loss_func(problem_params)
        print("learning rate", lr)
        optimizer = torch.optim.Adam(ICEO_model.parameters(), lr=lr)
        # optimizer = torch.optim.Adagrad(ICEO_model.parameters(), lr=lr)   
        # optimizer = torch.optim.RMSprop(ICEO_model.parameters(), lr=lr, alpha=0.99, eps=1e-08, momentum=0.9, centered=False, weight_decay=0)

        best_valid_loss = 10000
        for epoch in range(epoch_num):
            batch_size = 64
            batch_index = np.random.choice(n_train, batch_size)
            x_train_batch = x_train[batch_index, :]
            z_train_batch = z_train[batch_index, :]
            z_train_idx_batch = z_train_idx[batch_index]

            optimizer.zero_grad()
            pred_w, pred_p= ICEO_model(x_train_batch)
            if has_entropy:
                loss = nv_loss(pred_w, z_train_batch) + 10*entropy_loss(pred_p, z_train_idx_batch)
            elif has_reg:
                loss = iceo_loss(pred_w, z_train_batch)
            else:
                loss = nv_loss(pred_w, z_train_batch) 
            loss.backward()
            optimizer.step()
            if epoch % 250 == 0:
                with torch.no_grad():
                    pred_w, _ = ICEO_model(x_train)
                    train_loss = nv_loss(pred_w, z_train).item()
                    print("epoch", epoch, "training loss overall", train_loss)
                # print("model params", list(ICEO_model.parameters()))
                train_loss = loss.item()
            if  epoch % 25 == 0:
                with torch.no_grad():
                    pred_w, _ = ICEO_model(x_valid)
                loss = nv_loss(pred_w, z_valid)
                valid_loss = loss.item()
                
                if valid_loss <  best_valid_loss:
                        best_valid_loss = valid_loss
                        best_model_params = ICEO_model.state_dict()
                        print("valid loss", valid_loss)
                        print("best parameter saved")
                        
       

        ICEO_model.load_state_dict(best_model_params)
        with torch.no_grad():
            pred_w, _ = ICEO_model(x_test)
        loss = nv_loss(pred_w, z_test)
        print("test loss iceo model", loss.item())

### use approx oracle 
        pred_w = pred_w.detach().numpy()
        test_loss = loss_func_tools.nv_loss_func_np(problem_params, pred_w, z_test)
        test_loss_approx = test_loss.item()
        pred_model_state_dict = ICEO_model.predict_model.state_dict()
        if pred_model == 'linear':
            new_predict_model = iceo_framework.linear_pred(dim_features, dim_scenario, hidden_dim)
        else:
            new_predict_model = iceo_framework.two_layers_pred(dim_features, dim_scenario,hidden_dim)
        new_predict_model.load_state_dict(pred_model_state_dict)

        p_test = new_predict_model(x_test).detach().numpy()
        w_test = [optimization_oracle_tools.opt_oracle(prob_vec, problem_params) for prob_vec in p_test]
        w_test = np.array(w_test)
        test_loss = loss_func_tools.nv_loss_func_np(problem_params, w_test, z_test)
        test_loss = test_loss.item()

        print("test iceo loss", test_loss)

        mse = mse.detach()
        
        print("test_results", test_results)
        print([i,
                    n_train, pred_model, mse.item(), train_loss, test_loss_approx, test_loss])
        
        test_results.loc[len(test_results.index)] =  [i,
                    n_train, pred_model, mse, train_loss,test_loss_approx,  test_loss]

    return test_results




def newsvendor_entropy_test(problem_params):
    pred_model = problem_params['pred_model']
    n_train = problem_params['train_sample_size']
    n_test = problem_params['test_sample_size']
    dim_features = problem_params['dim_features']
    dim_sol = problem_params['dim_sol']
    dim_scenario = problem_params['dim_scenario']
    n_trails = problem_params['n_trails']
    epoch_num = problem_params['epoch_num']
    lr = problem_params['learning_rate']
    # generate xz_samples
    # B_true = data_generation_tools.generate_true_hypothesis(problem_params, seed=456)
    B_true = data_generation_tools.generate_true_hypothesis_NN(problem_params)
    B_true_matrix = data_generation_tools.generate_true_hypothesis(problem_params, seed=456)
    print("B_true", B_true)
    print("B_true_matrix", B_true_matrix)


    # if problem_params['baseline'] is None:
    test_results = pd.DataFrame(
    columns=['i', 'train_size', 'pred_model', 'approx_oracle_mse', 'train_loss',
                                   'entropy_loss', 'test_loss'])
    # else:
    #     test_results = pd.DataFrame(
    #         columns=['i', 'train_size', 'pred_model', 'approx_oracle_mse', 'train_loss',
    #                                 'entropy_loss','test_loss']+problem_params['baseline'])
#load prediction model
    if pred_model == 'linear':
        hidden_dim = problem_params.get('hidden_dim', 128)
        predict_model = iceo_framework.linear_pred(dim_features, dim_scenario,hidden_dim)
    elif pred_model == 'two_layers':
        hidden_dim = problem_params.get('hidden_dim', 128)
        predict_model = iceo_framework.two_layers_pred(dim_features, dim_scenario,hidden_dim)
    elif pred_model == 'two_layers_small':  
        hidden_dim = problem_params.get('hidden_dim', 32)
        predict_model = iceo_framework.two_layers_pred(dim_features, dim_scenario,hidden_dim)


#load oracle
        # load approx_oracle pickle:
        #pickle_dir = 'oracle_data'
    pickle_dir = problem_params['oracle_pickle_dir']
    dim_p = problem_params['dim_scenario']
    dim_w = problem_params['dim_sol']
    fullname = os.path.join(pickle_dir, 'oralce_output' + 'dim_p' + str(
            dim_p) + 'dim_w' + str(dim_w) + '.pkl')
    oracle_output = pickle.load(open(fullname, "rb"))
    if problem_params['oracle_kernel'] == 'nn':
        oracle_model = oracle_output['oracle_model']
        mse = oracle_output['mse']
        ## freeze oracle model
        for param in oracle_model.parameters():
            param.requires_grad = False
    elif problem_params['oracle_kernel']=='polynomial':
        p_train = oracle_output['p_train']
        p_train = torch.as_tensor(p_train, dtype=torch.float64)
        dual_coef = oracle_output['dual_coef']
        oracle_params = oracle_output['oracle_params']
        mse = oracle_output['mse']
    else:
        raise ValueError("Invalid oracle kernel! Use nn or polynomial instead.")
    print("approx oracle mse", mse)
# generate xz_samples
    for i in range(n_trails):
        if problem_params['data_generation_model'] == 'two_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        elif problem_params['data_generation_model'] == 'model mis-specification-deg':
            print('deg data',problem_params['deg_data'])
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true_matrix)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true_matrix)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true_matrix)
        elif problem_params['data_generation_model'] == 'multi_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        else:
            raise ValueError("ERROR! Invalid data_generation_model!!")


        def z_idx(problem_params, z_data):
            size = np.shape(z_data)[0]
            dim_scenario = problem_params['dim_scenario']
            scenario_list = problem_params['scenario_list']
            z_index = torch.zeros(size)
            ###
            for i in range(size):
                for k in range(dim_scenario):
                    # print(z_real[i].numpy())
                    # print(scenario_list[k])
                    if (z_data[i] == scenario_list[k, :]).all():
                        z_index[i] = k
            return z_index
        z_train_idx0 = z_idx(problem_params, z_train0)
        # print("z_test_idx shape",  z_train_idx0)
        z_test_idx0 = z_idx(problem_params, z_test0)
        z_valid_idx0 = z_idx(problem_params, z_valid0)

        z_train_idx = torch.as_tensor(z_train_idx0, dtype=torch.float32)
        z_test_idx = torch.as_tensor(z_test_idx0, dtype=torch.float32)
        z_valid_idx = torch.as_tensor(z_valid_idx0, dtype=torch.float32)


        x_train = torch.as_tensor(x_train0, dtype=torch.float32)
        x_test = torch.as_tensor(x_test0, dtype=torch.float32)
        # z_train = torch.as_tensor(z_train0, dtype=torch.float64)
        z_test = torch.as_tensor(z_test0, dtype=torch.float32)

        x_valid = torch.as_tensor(x_valid0, dtype=torch.float32)
        # z_valid = torch.as_tensor(z_valid0, dtype=torch.float64)

       
      
        entropy_loss = loss_func_tools.entropy_loss_func(problem_params)
        # iceo_loss = iceo_framework.iceo_loss_func(problem_params)
        print("learning rate", lr)
        optimizer = torch.optim.Adam(predict_model.parameters(), lr=lr)
        # optimizer = torch.optim.Adagrad(ICEO_model.parameters(), lr=lr)   
        # optimizer = torch.optim.RMSprop(ICEO_model.parameters(), lr=lr, alpha=0.99, eps=1e-08, momentum=0.9, centered=False, weight_decay=0)

# Train your m
        best_valid_loss = 100000
        for epoch in range(epoch_num):
            batch_size = 32
            batch_index = np.random.choice(n_train, batch_size)
            x_train_batch = x_train[batch_index, :]
            z_train_idx_batch = z_train_idx[batch_index]

            optimizer.zero_grad()
            pred_p = predict_model(x_train_batch)
            loss = entropy_loss(pred_p, z_train_idx_batch)
            loss.backward()
            optimizer.step()
            if epoch % 250 == 0:
                with torch.no_grad():
                    pred_p = predict_model(x_train)
                train_loss = entropy_loss(pred_p, z_train_idx).item()
                print("epoch", epoch, "training loss overall", train_loss)
                # print("model params", list(ICEO_model.parameters()))
                train_loss = loss.item()
            if  epoch % 25 == 0:
                # batch_size = n_test
                # batch_index = np.random.choice(n_test, batch_size)
                # x_valid_batch = x_valid[batch_index, :]
                # z_valid_batch = z_valid[batch_index, :]
                with torch.no_grad():
                    pred_p = predict_model(x_valid)
                loss = entropy_loss(pred_p, z_valid_idx)
                valid_loss = loss.item()
                
                if valid_loss <  best_valid_loss:
                        best_valid_loss = valid_loss
                        best_model_params = predict_model.state_dict()
                        print("valid loss", valid_loss)
                        print("best parameter saved")
                        
        pickle_dir = 'Path/entropy_initialization/'
        fullname = os.path.join(pickle_dir, 'best_model_params' + 'n_train'+ str(n_train) + 'dim_p' + str(
            dim_p) + 'dim_w' + str(dim_w) + '.pkl')
        pickle.dump(best_model_params, open(fullname, "wb"))
        predict_model.load_state_dict(best_model_params)
        print("Pickle dumped!")
        with torch.no_grad():
            pred_p = predict_model(x_test)
        entropy_loss = entropy_loss(pred_p, z_test_idx)
        print("test loss entropy model", loss.item())

        p_test = predict_model(x_test).detach().numpy()
        w_test = [optimization_oracle_tools.opt_oracle(prob_vec, problem_params) for prob_vec in p_test]
        w_test = np.array(w_test)
        test_loss = loss_func_tools.nv_loss_func_np(problem_params, w_test, z_test)
        test_loss = test_loss.item()



        #loss = iceo_loss(pred_w, z_test)
        print("test entropy nv loss", test_loss)
        #test_loss = loss.item()


        mse = mse.detach()
        


        print("test_results", test_results)
        print([i,
                    n_train, pred_model, mse.item(), train_loss, entropy_loss, test_loss])
        
        test_results.loc[len(test_results.index)] =  [i,
                    n_train, pred_model, mse, train_loss,entropy_loss,  test_loss]

    return test_results






def newsvendor_test_baseline_only(problem_params):
    # pred_model = problem_params['pred_model']
    # pred_model_back = prediction_tools.two_layers_model_back
    n_train = problem_params['train_sample_size']
    n_test = problem_params['test_sample_size']
    dim_features = problem_params['dim_features']
    dim_sol = problem_params['dim_sol']
    dim_scenario = problem_params['dim_scenario']
    n_trails = problem_params['n_trails']
    # generate xz_samples
    # B_true = data_generation_tools.generate_true_hypothesis(problem_params, seed=456)
    B_true = data_generation_tools.generate_true_hypothesis_NN(problem_params)
    print("B_true", B_true)


    test_results = pd.DataFrame(
        columns=['i', 'train_size']+problem_params['baseline'])
    # print(data_param_value)
    for i in range(n_trails):
        # load approx_oracle pickle:

        if problem_params['data_generation_model'] == 'two_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        elif problem_params['data_generation_model'] == 'context_shift':
            print('deg data', problem_params['deg_data'])
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true, low=-2, high=2)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true,
                                                                low=-2 + problem_params['deg_data'],
                                                                high=2 + problem_params['deg_data'])
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true, low=-2, high=2)
        elif problem_params['data_generation_model'] == 'multi_layer':
            x_train0, z_train0 = data_generation_tools.xz_dataset(n_train, problem_params, B_true)
            x_test0, z_test0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
            x_valid0, z_valid0 = data_generation_tools.xz_dataset(n_test, problem_params, B_true)
        else:
            raise ValueError("ERROR! Invalid data_generation_model!!")

        x_train = torch.as_tensor(x_train0, dtype=torch.float64)
        x_test = torch.as_tensor(x_test0, dtype=torch.float64)
        z_train = torch.as_tensor(z_train0, dtype=torch.float64)
        z_test = torch.as_tensor(z_test0, dtype=torch.float64)

        x_valid = torch.as_tensor(x_valid0, dtype=torch.float64)
        z_valid = torch.as_tensor(z_valid0, dtype=torch.float64)


        if  problem_params['baseline'] is not None:
            print(" problem_params['baseline']")
            baseline_loss_list = baseline_tools.get_baseline_loss(problem_params, x_train0, z_train0, x_test0, z_test0)
            test_results.loc[len(test_results.index)] = [i,
                n_train]+ baseline_loss_list

        else:
            baseline_iceo = None
            test_results.loc[len(test_results.index)] =  [i,
                    n_train, baseline_iceo ]

    return test_results


