import os
import sys
import linear_relax as LP_relax_file
import ip_model_whole as ip_model_whole_file
from ip_model_whole import IPOfunc
import numpy as np
import random
import pandas as pd
import math, time
import itertools
from sklearn import preprocessing
from sklearn.preprocessing import MinMaxScaler
import datetime
import torch
from torch import nn, optim
from torch.autograd import Variable
import torch.utils.data as data_utils
from torch.utils.data.dataset import Dataset
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error
import gurobipy as gp
import logging
import copy
from collections import defaultdict
import joblib
import gurobipy as gp
from gurobipy import GRB

total_month_num = LP_relax_file.month_num
month_num = LP_relax_file.month_num
x_num = LP_relax_file.x_num
y_num = LP_relax_file.y_num
var_num = LP_relax_file.var_num

featureNum = 4096
train_set_size = 70
target_num = 1
warm_start_stop_criterion = 10
stop_epoch_criterion = 20
log_regularizer = 1e-8
warm_start_value = 400

small_or_large = int(sys.argv[1])
startmark = int(sys.argv[2])
endmark = int(sys.argv[3])

dataset_path = os.path.abspath(os.path.dirname(os.getcwd()))
default_path = os.path.join(dataset_path, 'data/month_num=' + str(month_num) + '/')

if small_or_large == 0:
  store_file_path = os.path.join(dataset_path, 'data/month_num=' + str(month_num) + '/small/')
elif small_or_large == 1:
  store_file_path = os.path.join(dataset_path, 'data/month_num=' + str(month_num) + '/large/')

LP_relax_file.mkdir(store_file_path, 'MS (warm_start=' + str(warm_start_stop_criterion) + ', stop=' + str(stop_epoch_criterion) + ')')
    
def demand_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.demand)
        nn.init.constant_(m.bias, 0)

    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.demand, mode='fan_out', nonlinearity='relu')

    elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.demand, 1)
        nn.init.constant_(m.bias, 0)
        
def make_fc(num_layers, num_features, num_targets=target_num,
            activation_fn = nn.ReLU,intermediate_size=512, regularizers = True):
    net_layers = [nn.Linear(num_features, intermediate_size),activation_fn()]
    for hidden in range(num_layers-2):
        net_layers.append(nn.Linear(intermediate_size, intermediate_size))
        net_layers.append(activation_fn())
    net_layers.append(nn.Linear(intermediate_size, num_targets))
    net_layers.append(activation_fn())
    return nn.Sequential(*net_layers)
        

class MyCustomDataset():
    def __init__(self, feature, value):
        self.feature = feature
        self.value = value

    def __len__(self):
        return len(self.value)

    def __getitem__(self, idx):
        return self.feature[idx], self.value[idx]


class Intopt:
    def __init__(self, price, cost, n_features, num_layers=5, smoothing=False, thr=0.1, max_iter=None, method=1, mu0=None,
                 damping=0.5, target_size=target_num, epochs=8, optimizer=optim.Adam,
                 batch_size=month_num, **hyperparams):
        
        self.price = price
        self.cost = cost
        self.target_size = target_size
        self.n_features = n_features
        self.damping = damping
        self.num_layers = num_layers

        self.smoothing = smoothing
        self.thr = thr
        self.max_iter = max_iter
        self.method = method
        self.mu0 = mu0

        self.optimizer = optimizer
        self.batch_size = batch_size
        self.hyperparams = hyperparams
        self.epochs = epochs
        # print("embedding size {} n_features {}".format(embedding_size, n_features))

#        self.model = Net(n_features=n_features, target_size=target_size)
        self.model = make_fc(num_layers=self.num_layers,num_features=n_features)
        #self.model.apply(demand_init)
#        w1 = self.model[0].demand
#        print(w1)

        self.optimizer = optimizer(self.model.parameters(), **hyperparams)

    def fit(self, feature, value):
        logging.info("Intopt")
        train_df = MyCustomDataset(feature, value)

        criterion = nn.L1Loss(reduction='mean')  # nn.L1Loss(reduction='mean')
        grad_list = np.zeros(self.epochs)
        IP_grad_list = np.zeros(self.epochs)
        for i in range(self.epochs):
                IP_grad_list[i] = float("inf")
        for e in range(self.epochs):
            total_loss = 0
#          for parameters in self.model.parameters():
#            print(parameters)
            if e < warm_start_stop_criterion:
            #print('stage 1')
                train_dl = data_utils.DataLoader(train_df, batch_size=self.batch_size, shuffle=False)
                for feature, value in train_dl:
                    self.optimizer.zero_grad()
                    op = self.model(feature).squeeze()

                    loss = criterion(op, value)
                    total_loss += loss.item()
                    loss.backward()
                    self.optimizer.step()

                grad_list[e] = total_loss
                global stop_epoch
                stop_epoch = e
                # print("Epoch{} ::loss {} ->".format(e,total_loss))
                if e < warm_start_stop_criterion - 1:
                  print("{} ->".format(total_loss), end=" ")
                else:
                  print("{} ->".format(total_loss))
                
                global warm_start_value
                if e == warm_start_stop_criterion - 1 and grad_list[e] < warm_start_value:
                    self.model.eval()
                    criterion = nn.L1Loss(reduction='mean')  # nn.MSELoss(reduction='sum')
                    valid_df = MyCustomDataset(feature, value)
                    valid_dl = data_utils.DataLoader(train_df, batch_size=self.batch_size, shuffle=False)
                    corr_obj_list = []
                    true_obj_list = []
                    num = 0
                    for feature, value in valid_dl:
                        op = self.model(feature).squeeze()
                        # print(op)
                        loss = criterion(op, value)

                        true_demand = np.zeros(month_num)
                        pred_demand = np.zeros(month_num)
                        for i in range(month_num):
                            true_demand[i] = value[i]
                            pred_demand[i] = op[i]
                        
                        true_obj = LP_relax_file.actual_obj(self.price, self.cost, true_demand, n_instance=1)
                        true_obj_list.append(true_obj)
                        corrrlst = LP_relax_file.correction_single_obj(self.price, self.cost, pred_demand, true_demand)
                        corr_obj_list.append(corrrlst)
                        num = num + 1
                    
                    true_obj_np = np.array(true_obj_list)
                    corr_obj_np = np.array(corr_obj_list)
                    IP_grad_list[e] = np.mean(corr_obj_np)
#                    print(num)
                    print("TOV: ", np.mean(true_obj_np), "EOV: ", np.mean(corr_obj_np), "warm_start_PReg: ", np.mean(true_obj_np - corr_obj_np))
            
            else:
                if e == warm_start_stop_criterion:
                    lr = 1e-7
                    for param_group in self.optimizer.param_groups:
                        param_group['lr'] = lr
                # print(lr)
                train_dl = data_utils.DataLoader(train_df, batch_size=self.batch_size, shuffle=False)
                
                num = 0
                batchCnt = 0
                loss = Variable(torch.tensor(0.0, dtype=torch.double), requires_grad=True)
                for feature, value in train_dl:
                    self.optimizer.zero_grad()
                    op = self.model(feature).squeeze()
                    while torch.min(op) < 0 or torch.isnan(op).any() or torch.isinf(op).any():
                        self.optimizer.zero_grad()
    #                    self.model.__init__(self.n_features, self.target_size)
                        self.model = make_fc(num_layers=self.num_layers,num_features=self.n_features)
                        op = self.model(feature).squeeze()

#                    true_demand = value
                    
                    sol_cur = IPOfunc(price=self.price, cost=self.cost, true_demand=value, max_iter=self.max_iter, thr=self.thr, damping=self.damping,
                            smoothing=self.smoothing)(op)
                    
                    x_sol_cur = sol_cur[:x_num]
                    y_sol_cur = sol_cur[x_num:]
                    price_torch = torch.from_numpy(self.price).float()
                    cost_torch = torch.from_numpy(self.cost).float()
                    newLoss = - ((price_torch * y_sol_cur).sum() - (cost_torch * x_sol_cur).sum())
                    EOV_IP_value = newLoss.item()
                    total_loss += EOV_IP_value
                    newLoss.backward()
                    self.optimizer.step()
                    

                    batchCnt = batchCnt + 1
                    num = num + 1

                grad_list[e] = total_loss / num
                stop_epoch = e

                # compute IP_grad
                valid_dl = data_utils.DataLoader(train_df, batch_size=self.batch_size, shuffle=False)
                corr_obj_list = []
                true_obj_list = []
                num = 0
                for feature, value in valid_dl:
                    op = self.model(feature).squeeze()
                    #            print(op)
                    loss = criterion(op, value)

                    true_demand = np.zeros(month_num)
                    pred_demand = np.zeros(month_num)
                    for i in range(month_num):
                        true_demand[i] = value[i]
                        pred_demand[i] = op[i]

                    true_obj = LP_relax_file.actual_obj(self.price, self.cost, true_demand, n_instance=1)
                    true_obj_list.append(true_obj)
                    corrrlst = LP_relax_file.correction_single_obj(self.price, self.cost, pred_demand, true_demand)
                    corr_obj_list.append(corrrlst)
                    num = num + 1

                true_obj_np = np.array(true_obj_list)
                corr_obj_np = np.array(corr_obj_list)
                IP_grad_list[e] = np.mean(corr_obj_np)

            logging.info("EPOCH Ends")
            if e >= warm_start_stop_criterion:
                print("Epoch{} ::EOV {} ->".format(e,IP_grad_list[e]))
            #print("Epoch{}".format(e))
            #          for param_group in self.optimizer.param_groups:
            #            print(param_group['lr'])
            if grad_list[6] > warm_start_value:
                break
            if e >= 1 and abs(grad_list[e] - grad_list[e-1]) <= 0.001:
                break
            if e >= warm_start_stop_criterion and abs(IP_grad_list[e] - IP_grad_list[e-1]) <= 0.001:
                break
            if e >= warm_start_stop_criterion and abs(IP_grad_list[e]) < abs(IP_grad_list[e-1]):
                break

            

    def val_loss(self, feature, value):
        valueTemp = value.numpy()
#        test_instance = len(valueTemp) / self.batch_size
        instance_num = np.size(valueTemp, 0) / self.batch_size
#        print(valueTemp.shape, instance_num)
        true_demand_total = valueTemp
#        print(true_price.shape, true_weight.shape)
        true_obj = LP_relax_file.actual_obj(self.price, self.cost, true_demand_total, n_instance=int(instance_num))
#        print(np.sum(real_obj))

        self.model.eval()
        criterion = nn.L1Loss(reduction='mean')  # nn.MSELoss(reduction='sum')
        valid_df = MyCustomDataset(feature, value)
        valid_dl = data_utils.DataLoader(valid_df, batch_size=self.batch_size, shuffle=False)

        corr_obj_list = []
        len = np.size(valueTemp)
        predVal = torch.zeros(len)
        
        num = 0
        for feature, value in valid_dl:
            op = self.model(feature).squeeze()
#            print(op)
            loss = criterion(op, value)

            true_demand = np.zeros(month_num)
            pred_demand = np.zeros(month_num)
            for i in range(month_num):
                true_demand[i] = value[i]
                pred_demand[i] = op[i]
                predVal[i+num*month_num] = op[i]

            corrrlst = LP_relax_file.correction_single_obj(self.price, self.cost, pred_demand, true_demand)
            corr_obj_list.append(corrrlst)
            num = num + 1
            

        self.model.train()
        print("TOV: ", sum(true_obj)/instance_num, "EOV: ", sum(corr_obj_list)/instance_num, "PReg: ", sum(abs(true_obj - np.array(corr_obj_list)))/instance_num)
#        print(corr_obj_list)
#        print(corr_obj_list-real_obj)
#        print(np.sum(corr_obj_list))
#        return prediction_loss, abs(np.array(obj_list) - real_obj)
        return abs(np.array(corr_obj_list) - true_obj), predVal



print("*** Baseline ****")

simulation_time = 30
recordBest = np.zeros(simulation_time)

if small_or_large == 0:
  print("small price,", end=' ')
elif small_or_large == 1:
  print("large price,", end=' ')
print("month_num: ", month_num, "warm_start_stop_criterion: ", warm_start_stop_criterion, "stop_epoch_criterion: ", stop_epoch_criterion)


for testi in range(startmark, endmark):
    print(testi)
    stop_epoch = 0
    cost = np.loadtxt(os.path.join(dataset_path, 'data/month_num='+str(total_month_num)+'/cost/cost(' + str(testi) + ').txt'))
    if small_or_large == 0:
      price = np.loadtxt(os.path.join(dataset_path, 'data/month_num='+str(total_month_num)+'/small_price/price(' + str(testi) + ').txt'))
    elif small_or_large == 1:
      price = np.loadtxt(os.path.join(dataset_path, 'data/month_num='+str(total_month_num)+'/large_price/price(' + str(testi) + ').txt'))
    
    x_train = np.loadtxt(os.path.join(default_path, 'train_features/train_features(' + str(testi) + ').txt'))
    y_train = np.loadtxt(os.path.join(default_path, 'train_demands/train_demands(' + str(testi) + ').txt'))
    feature_train = torch.from_numpy(x_train).float()
    value_train = torch.from_numpy(y_train).float()
    
    x_test = np.loadtxt(os.path.join(default_path, 'test_features/test_features(' + str(testi) + ').txt'))
    y_test = np.loadtxt(os.path.join(default_path, 'test_demands/test_demands(' + str(testi) + ').txt'))
    feature_test = torch.from_numpy(x_test).float()
    value_test = torch.from_numpy(y_test).float()
    
    start = time.time()
    damping = 1e-2
    thr = 1e-3
    lr = 1e-5
    bestTrainCorrReg = float("inf")
    max_retrain_num = 10
    while stop_epoch < warm_start_stop_criterion and max_retrain_num > 0:
        max_retrain_num = max_retrain_num - 1
        clf = Intopt(price, cost, damping=damping, lr=lr, n_features=featureNum, thr=thr, epochs=stop_epoch_criterion)
        clf.fit(feature_train, value_train)

        if stop_epoch >= warm_start_stop_criterion or max_retrain_num <= 0:
            end = time.time()
            train_rslt, predTrainVal = clf.val_loss(feature_train, value_train)
            avgTrainCorrReg = np.mean(train_rslt)
            trainHSD_rslt = 'train: ' + str(np.mean(train_rslt))
            bestTrainCorrReg = avgTrainCorrReg
            torch.save(clf.model.state_dict(), 'MS_'+str(small_or_large)+'_'+str(month_num)+'month_model.pkl')
            print(trainHSD_rslt)


    clfBest = Intopt(price, cost, damping=damping, lr=lr, n_features=featureNum, thr=thr, epochs=stop_epoch_criterion)
    clfBest.model.load_state_dict(torch.load('MS_'+str(small_or_large)+'_'+str(month_num)+'month_model.pkl'))

    val_rslt, predTestVal = clfBest.val_loss(feature_test, value_test)

    predTestVal = predTestVal.detach().numpy()
    predTestDemand = np.zeros((predTestVal.size, 2))
    for i in range(predTestVal.size):
#        predValue[i][0] = int(i/itemNum)
        predTestDemand[i][0] = y_test[i]
        predTestDemand[i][1] = predTestVal[i]
    np.savetxt(os.path.join(store_file_path,'MS (warm_start=' + str(warm_start_stop_criterion) + ', stop=' + str(stop_epoch_criterion) + ')/MS_demands(' + str(testi) + ').txt'), predTestDemand, fmt="%.2f")
    
    HSD_rslt = 'test: ' + str(np.mean(val_rslt)) + ' MSE: ' + str(mean_squared_error(y_test, predTestVal))
    print(HSD_rslt, end=" ")
    print ('Elapsed time: ' + str(end-start))
    recordBest[testi] = np.sum(val_rslt)

print(recordBest)
