# -*- coding: utf-8 -*-
"""
Created on Sat Oct  8 18:40:05 2022

@author: amhijazi
"""

import random
import torch
import torch.nn as nn
import time
import argparse
import os
import threading
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
import pandas as pd 
import numpy as np
from matplotlib import pyplot as plt
import warnings
from gurobipy import *
import pickle
import multiprocessing
from joblib import Parallel, delayed
from PointerTransformeModelDefinition import PointerTransformer, test
from PointerTransformeModelDefinition import CGDataset as PointerCGDataset
from CGModulesDefinitions import generate_input_dataUniform, generate_input_dataWeibull, generateBasicInitialSchedules, OriginalModel,  MasterProblem, Subproblem, SubproblemLS_SingleOO
warnings.filterwarnings("ignore", category=UserWarning)
#torch.random.manual_seed(231)
from concurrent.futures import ThreadPoolExecutor
torch.manual_seed(231)
torch.cuda.manual_seed(231)


#%%
###################
# Hyper-parameters
###################
class DotDict(dict):
    def __init__(self, **kwds):
        self.update(kwds)
        self.__dict__ = self



#%%
def experiment(seed, M, N, iteration):
    np.random.seed(seed)
    random.seed(seed)
    InputDf = generate_input_dataUniform(seed, M,N)

    Schedule = generateBasicInitialSchedules(seed, M,N, InputDf, iteration)
    original = False
    if original:
        start_timetotal = time.time()
        m = OriginalModel(M, N, InputDf)
        m.buildModel()
        status = m.solveModel(3600, 0)
        GurobiObj = m.getObj()
        print('Gurobi Obj ', GurobiObj)
        gurobiTime = time.time() - start_timetotal
        print('gruobi time ', gurobiTime)
        print('status ', status)
        print('gap ', m.getGap())

        if status == True:
            print('optimal')
            
    return InputDf, Schedule

def run(method,modelPointer,  M, N, InputDf, Schedule,timelimit, DPNumCols = None):
    print('--------------------------')    
    print ('method ', method)
    
    master = MasterProblem(M, N, Schedule, InputDf)
    master.buildModel()
 
    ScheduleDf = Schedule.reset_index(drop = False)
    
    Input_NNData = torch.from_numpy(np.array([np.array([np.concatenate(([0.0], InputDf.loc[m, :].ProcessingTime)), np.concatenate(([0.0],InputDf.loc[m, :].Weight))]).T for m in range(1, M+1)]))
    modelImprovable = True
    count = [max(ScheduleDf['ScheduleID'][ScheduleDf['Machine'] == k]) for k in range (1, M+1)]
    NNCount = np.zeros(M)
    DPCount = np.zeros(M)
    LSCount = np.zeros(M)
    master.SolveRelaxedModel()   
    start_timetotal = time.time()
    ObjValTracking = pd.DataFrame(columns=['time', 'obj_val'])
    
    ObjValTracking = ObjValTracking._append({'time': time.time() - start_timetotal, 'obj_val': master.getObj()}, ignore_index = True)
    if method == 'Pointer Greedy DP':
        cols = 1
        TOKENS = {'<sos>': 0, '<eos>': 1}
        
        n_tokens = len(TOKENS)
        max_jobs = N + 1
        padded_data = torch.zeros((M, max_jobs + n_tokens, 4 + n_tokens), dtype=torch.float32)

        padded_data[:, TOKENS['<sos>'], 3] = 1.0
        padded_data[:, TOKENS['<eos>'], 4] = 1.0
        padded_data[:, n_tokens + max_jobs:, 5] = 1.0
        
        padded_data[:, n_tokens:n_tokens + max_jobs, 0] = Input_NNData[:, :, 0]
        padded_data[:, n_tokens:n_tokens + max_jobs, 1] = Input_NNData[:, :, 1]

        batch_lengths = torch.full((M,),max_jobs + n_tokens, dtype = torch.int64)
        
        while modelImprovable and (time.time() - start_timetotal) < timelimit:
            #duals_tensor = torch.tensor(master.getDuals(), dtype=torch.float32)
            duals = master.getDuals()
            
            job_duals = torch.tensor(duals[:N], dtype=torch.float32)
            machine_duals = torch.tensor(duals[N:N+M], dtype=torch.float32)
            padded_data[:, n_tokens:n_tokens + max_jobs, 2] = job_duals.unsqueeze(0)
            padded_data[:, 0, 2] = machine_duals
   
            #machine_duals = duals_tensor[N:N+M].unsqueeze(1)
            #static_duals = duals_tensor[:len(duals_tensor) - M].unsqueeze(0).repeat(M, 1)
            #duals_NNData = torch.cat([machine_duals, static_duals], dim=1).unsqueeze(2)
            # Combine with input data
            #NNData = torch.cat((Input_NNData, duals_NNData), dim=2)

            #padded_data[:, n_tokens:n_tokens + max_jobs, :3] = NNData.squeeze(2)
            # Run model inference
            ReducedCost, Cost, Schedules = test(modelPointer, padded_data, batch_lengths, M)

            ReducedCost = ReducedCost.cpu().numpy()
            Cost = Cost.cpu().numpy()
            duals = pd.Series(master.getDuals())  
            duals.index +=1
            modelImprovable = False
            for k in range(1, M+1):
                if ReducedCost[k-1] < -0.1:
                    modelImprovable = True
                    NNCount[k-1] += 1
                    output_schedule = Schedules[k-1]
                    negSchedule = {
                        'Machine': k,
                        'JobsInSchedule': output_schedule,
                        'ScheduleCost': Cost[k-1]
                    }
                    count[k-1] += 1

                    # Build new column
                    jobs = np.zeros(N+1)
                    jobs[negSchedule['JobsInSchedule']] = 1
                    machines = np.zeros(M+1)
                    machines[k] = 1
                    newColumn = np.concatenate((jobs[1:], machines[1:]))
                    master.addColumn(negSchedule['ScheduleCost'], newColumn, f'y{[k, count[k-1]]}')
            if not modelImprovable:
                for k in range(1, M+1):        
                        subproblem = SubproblemLS_SingleOO(k, N, InputDf, duals, cols)
                        Schedules, obj, reducedCost = subproblem.local_search()  
                        for i in range(len(Schedules)):
                            if reducedCost[i] < -0.1:
                                LSCount[k-1] += 1
                                modelImprovable = True    
                                count[k-1] += 1
                                negSchedule = {'Machine':k, 'JobsInSchedule': Schedules[i], 
                                            'ScheduleCost': obj[i]}
                                jobs = np.zeros(N+1)
                                jobs[negSchedule['JobsInSchedule']] = 1
                                
                                machines = np.zeros(M+1)
                                machines[k] = 1            
                                newColumn = np.concatenate((jobs[1:N+1], machines[1:M+1]))    
                
                                master.addColumn(negSchedule['ScheduleCost'], newColumn, 'y'+str([k,count[k-1]]))  
            master.SolveRelaxedModel()  
            ObjValTracking = ObjValTracking._append({'time': time.time() - start_timetotal, 'obj_val': master.getObj()}, ignore_index = True)
           
    elif method == 'Pointer DP':
        TOKENS = {'<sos>': 0, '<eos>': 1}
        n_tokens = len(TOKENS)
        max_jobs = N + 1
        padded_data = torch.zeros((M, max_jobs + n_tokens, 4 + n_tokens), dtype=torch.float32)

        padded_data[:, TOKENS['<sos>'], 3] = 1.0
        padded_data[:, TOKENS['<eos>'], 4] = 1.0
        padded_data[:, n_tokens + max_jobs:, 5] = 1.0

        batch_lengths = torch.full((M,),max_jobs + n_tokens, dtype = torch.int64)
        
        while modelImprovable and (time.time() - start_timetotal) < timelimit:
            duals_tensor = torch.tensor(master.getDuals(), dtype=torch.float32)
            
            machine_duals = duals_tensor[N:N+M].unsqueeze(1)
            static_duals = duals_tensor[:len(duals_tensor) - M].unsqueeze(0).repeat(M, 1)
            duals_NNData = torch.cat([machine_duals, static_duals], dim=1).unsqueeze(2)
            # Combine with input data
            NNData = torch.cat((Input_NNData, duals_NNData), dim=2)

            padded_data[:, n_tokens:n_tokens + max_jobs, :3] = NNData.squeeze(2)
            # Run model inference
            ReducedCost, Cost, Schedules = test(modelPointer, padded_data, batch_lengths, M)

            ReducedCost = ReducedCost.cpu().numpy()
            Cost = Cost.cpu().numpy()

            modelImprovable = False
            for k in range(1, M+1):
                if ReducedCost[k-1] < -0.1:
                    modelImprovable = True
                    NNCount[k-1] += 1
                    output_schedule = Schedules[k-1]
                    negSchedule = {
                        'Machine': k,
                        'JobsInSchedule': output_schedule,
                        'ScheduleCost': Cost[k-1]
                    }
                    count[k-1] += 1

                    # Build new column
                    jobs = np.zeros(N+1)
                    jobs[negSchedule['JobsInSchedule']] = 1
                    machines = np.zeros(M+1)
                    machines[k] = 1
                    newColumn = np.concatenate((jobs[1:], machines[1:]))
                    master.addColumn(negSchedule['ScheduleCost'], newColumn, f'y{[k, count[k-1]]}')
            if not modelImprovable:
                duals = pd.Series(master.getDuals())  
                duals.index +=1
                for k in range(1, M+1):
                    subproblem = Subproblem(k, N, InputDf, duals, 1)
                    output = subproblem.DP()  
                   #check if new schedules improve master problem
                    for negSchedule in  output:
                          DPCount[k-1] += 1
                          modelImprovable = True    
                          count[k-1] += 1
               
                          jobs = np.zeros(N+1)
                          jobs[negSchedule['JobsInSchedule']] = 1
                          machines = np.zeros(M+1)
                          machines[k] = 1            
                          newColumn = np.concatenate((jobs[1:N+1], machines[1:M+1]))    
               
                          master.addColumn(negSchedule['ScheduleCost'], newColumn, 'y'+str([k,count[k-1]]))  

            master.SolveRelaxedModel()  
            ObjValTracking = ObjValTracking._append({'time': time.time() - start_timetotal, 'obj_val': master.getObj()}, ignore_index = True)

           
    elif method == 'DP' or method == 'DP 5' or method == 'DP 20':
            
            while (modelImprovable) and time.time() - start_timetotal < timelimit:
                #solve relaxed master
                duals = pd.Series(master.getDuals())  
                #build subproblem
                duals.index +=1
                modelImprovable = False
                for k in range(1, M+1):
                    subproblem = Subproblem(k, N, InputDf, duals, DPNumCols)
                    output = subproblem.DP()  
                    #check if new schedules improve master problem
                    for negSchedule in  output:
                
                            DPCount[k-1] += 1
                            modelImprovable = True    
                            count[k-1] += 1
            
                            jobs = np.zeros(N+1)
                            jobs[negSchedule['JobsInSchedule']] = 1
                            
                            machines = np.zeros(M+1)
                            machines[k] = 1            
                            newColumn = np.concatenate((jobs[1:N+1], machines[1:M+1]))    
            
                            master.addColumn(negSchedule['ScheduleCost'], newColumn, 'y'+str([k,count[k-1]]))  
                master.SolveRelaxedModel()    
                ObjValTracking = ObjValTracking._append({'time': time.time() - start_timetotal, 'obj_val': master.getObj()}, ignore_index = True)

    elif method == 'Greedy DP':
        cols = 1
   
        while (modelImprovable) and time.time() - start_timetotal <  timelimit:
            duals = pd.Series(master.getDuals())  
            #build subproblem
            duals.index +=1
            modelImprovable = False
            for k in range(1, M+1):
                subproblem = SubproblemLS_SingleOO(k, N, InputDf, duals, cols)
                Schedules, obj, reducedCost = subproblem.local_search()  
                #check if new schedules improve master problem
                for i in range(len(Schedules)):
                    if reducedCost[i] < -0.1:
                        
                            LSCount[k-1] += 1
                            modelImprovable = True    
                            count[k-1] += 1
                            negSchedule = {'Machine':k, 'JobsInSchedule': Schedules[i], 
                                        'ScheduleCost': obj[i]}
                            jobs = np.zeros(N+1)
                            jobs[negSchedule['JobsInSchedule']] = 1
                            
                            machines = np.zeros(M+1)
                            machines[k] = 1            
                            newColumn = np.concatenate((jobs[1:N+1], machines[1:M+1]))    
            
                            master.addColumn(negSchedule['ScheduleCost'], newColumn, 'y'+str([k,count[k-1]]))  
            if not modelImprovable:
                for k in range(1, M+1):
                    subproblem = Subproblem(k, N, InputDf, duals, 1)
                    output = subproblem.DP()  
                    
                    #check if new schedules improve master problem
                    for negSchedule in  output:
                            DPCount[k-1] += 1
                            modelImprovable = True    
                            count[k-1] += 1
                            jobs = np.zeros(N+1)
                            jobs[negSchedule['JobsInSchedule']] = 1 
                            machines = np.zeros(M+1)
                            machines[k] = 1            
                            newColumn = np.concatenate((jobs[1:N+1], machines[1:M+1]))    
                            master.addColumn(negSchedule['ScheduleCost'], newColumn, 'y'+str([k,count[k-1]]))  
            master.SolveRelaxedModel()    
            ObjValTracking = ObjValTracking._append({'time': time.time() - start_timetotal, 'obj_val': master.getObj()}, ignore_index = True)
     
        #solve the integer master problem with all the patterns with all the patterns introduced previously
    Time = time.time() - start_timetotal
        
    print('Time ', Time)
    print('NNCount ', sum(NNCount))
    print('DPCount ', sum(DPCount))
    print('LSCount ', sum(LSCount))
    print('Total cols ', sum(DPCount) + sum(NNCount) + sum(LSCount))
    master.SolveRelaxedModel()
    print('Obj ' , master.getObj())
    print('--------------------------')   
    return ObjValTracking 
#%%   

def run_exp(seed, M, N, iteration, modelPointer, timelimit):
    print('seed ', seed)
    InputDf, Schedule = experiment(seed, M, N, iteration)
    try:
        results = {
            'Pointer+DP': run('Pointer DP', modelPointer, M, N, InputDf, Schedule, timelimit,  DPNumCols=None),
            'Greedy+DP': run('Greedy DP', None, M, N, InputDf, Schedule,timelimit, DPNumCols=None),
                    }

        return seed, results
    except Exception as e:
        print(f"Error in seed {seed}: {e}")
        return seed, None


if __name__ == "__main__":
    # Create result dictionaries
    multiprocessing.set_start_method("spawn", force=True)
    seeds =  [550660, 551660, 550661, 551661, 552660, 987543, 787542, 7875429, 900000, 83735228]
    args = DotDict()        
    args.c_inputs = 6
    args.c_embed = 64
    args.c_hidden = 64
    args.n_heads = 8
    args.n_layers_encoder = 2
    args.n_layers_decoder = 2
    args.dropout = 0.0
    args.use_cuda = True
    args.n_workers = 2

    args.n_epochs = 100
    args.batch_size = 16
    args.lr = 0.0001
    args.weight_decay = 1e-06
    log_interval = 500

    ###################
    # Hardware : CPU / GPU(s)
    ###################

    device = torch.device("cpu"); gpu_id = -1 # select CPU

    modelPointer = PointerTransformer(c_inputs=args.c_inputs, c_embed=args.c_embed, n_heads=args.n_heads,
    n_layers_encoder=args.n_layers_encoder,n_layers_decoder=args.n_layers_decoder, dropout=args.dropout, c_hidden=args.c_hidden).to(device)


    ###################
    # Load checkpoint
    ###################

    checkpoint_file = "./checkpointfinal_PointerTransformerMachine.pkl"
    checkpoint = torch.load(checkpoint_file, map_location=device, weights_only=False)

    modelPointer.load_state_dict(checkpoint['model_train'])
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--M", type=int, required=True)
    parser.add_argument("--N", type=int, required=True)
    parser.add_argument("--init_cols", type=int, required=True)
    parser.add_argument("--timelimit", type=int, required=True)

    
    cmd_args = parser.parse_args()

    M = cmd_args.M
    N = cmd_args.N

    iteration = cmd_args.init_cols #initial number of columns per machine in the master problem, set to 1 for the experiments in the paper, but can be increased to reduce the number of iterations in the column generation at the expense of increasing the time of each iteration.
    timelimit = cmd_args.timelimit 
    print('M ', M, 'N ', N,  'Number of Initial Cols per Machine', iteration, 'time limit per seed per method ', timelimit)
    NN_Results = {}
    Greedy_Results = {}
    
    n_jobs = -1  

    # Run experiments in parallel using Joblib
    results = Parallel(n_jobs=n_jobs, backend="loky")(
        delayed(run_exp)(seed, M, N, iteration, modelPointer, timelimit) for seed in seeds
    )

    # Collect results
    for seed, res in results:
        print('seed ', seed )
        NN_Results[seed] = res['Pointer+DP']
        Greedy_Results[seed] = res['Greedy+DP']

    methods_results = {
            "CG NN-DP": NN_Results,
            "CG Heuristic-DP": Greedy_Results,
        }


    with open(f"{M}M_{N}N_methods_results.pkl", "wb") as f:
        pickle.dump(methods_results, f)

        # First, find the global max time across all methods/seeds
    max_time = max(
            df['time'].max() 
            for method_dict in methods_results.values() 
            for df in method_dict.values()
        )

        # Common time grid
    time_grid = np.linspace(0, timelimit, 200)

    plt.figure(figsize=(5,4))

    colors = {"CG NN-DP": "blue", "CG Heuristic-DP": "green"}
    for method_name, results_dict in methods_results.items():
            all_vals = []
            for seed, df in results_dict.items():
                base = df['obj_val'].iloc[0]
                rel = df['obj_val'] / base
                interp_vals = np.interp(time_grid, df['time'], rel)
                all_vals.append(interp_vals)
            all_vals = np.array(all_vals)
            
            mean_vals = all_vals.mean(axis=0)
            std_vals  = all_vals.std(axis=0)
            
            # mean curve
            plt.plot(time_grid, mean_vals, color=colors[method_name], linewidth=2, label=method_name)
            
            # ±1 std shading
            plt.fill_between(time_grid, mean_vals - std_vals, mean_vals + std_vals,
                            color=colors[method_name], alpha=0.2)

    plt.xlabel("Time (seconds)")
    plt.ylabel("Relative RMP Objective Value")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{M}M_{N}N.png") 
    # =============================================================================

