# -*- coding: utf-8 -*-
"""
Created on Sat Oct  8 18:45:28 2022

@author: amhijazi
"""

from gurobipy import *
import matplotlib.pyplot as plt
import random
import numpy as np
import pandas as pd
import random
import heapq
from functools import lru_cache
from operator import itemgetter 

seeds = [550660, 551660, 550661, 551661, 552660, 987543, 787542, 7875429, 900000, 8373522, 899000, 34876044, 986543213]
import copy
#%% 
'''
number of machines (m) with 6 options: 2, 4, 8, 12, 16, 20
number of jobs (n) with 6 options: 20, 30, 40, 60, 80, 100
weights w_j the weight for each job is an integer number uniformly drawn from [1, 100]
processing time p_ij for the unrelated paralell machines is an integer from the uniform distribution [1,30]
'''


#%% Generate data 
def generate_input_dataUniform(seed, M, N):
    np.random.seed(seed)

    random.seed(seed)
    '''
    Generates randome data to be used in the optimization
    ''' 
    #lambda = 
    Weights = pd.Series([np.random.randint(1,100) for i in range(1, N+1)], index = range(1, N+1), name = 'Weight')
    Weights.index.name = 'Job'
    index= pd.MultiIndex.from_product([range(1, M+1), range(1, N+1)], names = ['Machine', 'Job'])
    ProcessingTime = pd.Series([np.random.randint(1,30) for i in range(1, (N*M)+1)], index = index, name ='ProcessingTime')
    InputDf = pd.merge(ProcessingTime.reset_index(), Weights.reset_index(), on ="Job", how = 'left').set_index(['Machine', 'Job'])


    InputDf['P/W'] =  pd.Series(InputDf['ProcessingTime'].div(InputDf['Weight']))   
    B = []
    A = []
    for m in range(1, M+1):
        temp_sort = np.array(InputDf.loc[m].sort_values('P/W').index)
        for j in range(1, N+1):
            B.append([temp_sort[i] for i in range((temp_sort == j).nonzero()[0][0])])
            A.append([temp_sort[i] for i in range(1+(temp_sort == j).nonzero()[0][0], len(temp_sort))])
            
    InputDf['B'] = B
    InputDf['A'] = A
    return InputDf
#%%
def generate_input_dataWeibull(seed, M, N, shape_weights=1.5, scale_weights=50, shape_processing=2, scale_processing=15):
    '''
    Generates random data to be used in the optimization
    ''' 
    np.random.seed(seed)

    random.seed(seed)
    # Generate weights using Weibull distribution
    Weights = pd.Series(np.random.weibull(shape_weights, N) * scale_weights, index=range(1, N+1), name='Weight')
    Weights = Weights.astype(int)
    Weights.index.name = 'Job'

    # Generate processing times using Weibull distribution
    index = pd.MultiIndex.from_product([range(1, M+1), range(1, N+1)], names=['Machine', 'Job'])
    ProcessingTime = pd.Series(np.random.weibull(shape_processing, N * M) * scale_processing, index=index, name='ProcessingTime')
    ProcessingTime = ProcessingTime.astype(int)

    # Merge into a single DataFrame
    InputDf = pd.merge(ProcessingTime.reset_index(), Weights.reset_index(), on="Job", how='left').set_index(['Machine', 'Job'])

    # Calculate P/W ratio
    InputDf['P/W'] = InputDf['ProcessingTime'].div(InputDf['Weight'])

    B = []
    A = []

    for m in range(1, M+1):
        temp_sort = np.array(InputDf.loc[m].sort_values('P/W').index)
        for j in range(1, N+1):
            B.append([temp_sort[i] for i in range((temp_sort == j).nonzero()[0][0])])
            A.append([temp_sort[i] for i in range(1+(temp_sort == j).nonzero()[0][0], len(temp_sort))])

    InputDf['B'] = B
    InputDf['A'] = A

    return InputDf

#%% initial schedules: 
def generateBasicInitialSchedules(seed, M, N, InputDf, colsperM): 
    np.random.seed(seed)

    random.seed(seed)
    Weight = InputDf['Weight']
    ProcessingTime = InputDf['ProcessingTime']
    ScheduleDf = pd.DataFrame(columns = ['Machine'] +  ['ScheduleID'] + ['JobsInSchedule'] + ['ScheduleCost'] +  ['Job'+str(i) for i in range(1,N+1)] + ['Machine'+str(i) for i in range(1,M+1)]) 
    i = -1
    for m in range(1, M+1):
        temp_sort = np.array(InputDf.loc[m].sort_values('P/W').index)
        subsets = []
        iteration = 1
        while iteration <= colsperM+1: 
                col = [temp_sort[i] for i in sorted(random.sample(range(len(temp_sort)), random.randint(1,N)))]
                if col not in subsets:                    
                    i+=1
                    schedule = {}
                    schedule['Machine'] = m
                    schedule['ScheduleID'] = iteration
                    schedule['JobsInSchedule'] = col
                    schedule['ScheduleCost'] = sum(ProcessingTime.loc[m, col].cumsum()*Weight.loc[m,col])
                    ScheduleDf = ScheduleDf._append(schedule, ignore_index = True)
                    jobs = np.zeros(N+1)
                    jobs[col] = 1
                    machines = np.zeros(M+1)
                    machines[m] = 1
                    ScheduleDf.loc[i, 'Job' + str(1) : 'Job' + str(N)] = jobs[1:N+1]
                    ScheduleDf.loc[i, 'Machine' + str(1): 'Machine' + str(M)] = machines[1:M+1]
                    iteration += 1
    ScheduleDf.set_index(['Machine', 'ScheduleID'], inplace = True)
    return ScheduleDf     
#%%        
class MasterProblem:
    def __init__(self, M, N, ScheduleDf, InputDf):
        self.N = N
        self.M = M
        self.Weight = InputDf['Weight']
        self.ProcessingTime = InputDf['ProcessingTime']
        self.ScheduleCost = ScheduleDf['ScheduleCost']
        #self.A = ScheduleDf[['Job' + str(i) for i in range(1, self.N+1)]]
        self.A = ScheduleDf[['Job' + str(i) for i in range(1, self.N+1)]]
        self.SchedulesIndex = ScheduleDf.index
        self.model = Model("MasterProblem")
        self.model.Params.LogToConsole = 0
    def buildModel(self):
        self.generateVariables()
        self.generateConstraints()
        self.generateObjective()
        self.model.update()
        
    def writeLP(self):
        self.model.write("gurobi_new.lp")
        
    def generateVariables(self):
        # self.scheduleUseVar = self.model.addVars(self.SchedulesIndex, vtype = GRB.BINARY, name = "y")
        self.scheduleUseVar = pd.Series(self.model.addVars(self.SchedulesIndex, 
                                                vtype = GRB.BINARY, name = "y").values(),
              index= self.SchedulesIndex, name='y')

    def generateConstraints(self):
    
        self.model.addConstrs((quicksum(self.A['Job' + str(j)]* self.scheduleUseVar)
                                 == 1 for j in range(1, self.N+1)), 'Job')
             
        for (i, lhs, rhs) in(pd.concat([self.scheduleUseVar.groupby('Machine').agg(quicksum).to_frame('machine_constraint'), 
                                      pd.Series([1]*(self.M), index = range(1, self.M+1))], axis =1).itertuples()):
              self.model.addConstr(lhs <= rhs, name = "Machine"+str(i)+"Constraint")    
          
    def generateObjective(self):
        self.model.setObjective(quicksum(self.scheduleUseVar * self.ScheduleCost),
                                sense = GRB.MINIMIZE)
   
    def SolveRelaxedModel(self):
        self.relaxedModel = self.model.relax()
        self.relaxedModel.optimize()
        
        
    def getObj(self):
        obj = self.relaxedModel.getObjective()
        return (obj.getValue())
        
    def getX(self):
         print(self.relaxedModel.getAttr("X", self.relaxedModel.getVars()))
         
    def getDuals(self):
        return self.relaxedModel.getAttr(GRB.Attr.Pi)
    
    def addColumn(self, Objective, newColumn, ctName):
        newColumn = Column(newColumn, self.model.getConstrs())
        self.model.addVar(vtype = GRB.BINARY, obj = Objective, column = newColumn, name = ctName)
        self.model.update()
    
    def solveModel(self, timeLimit, GAP):
        self.model.setParam('TimeLimit', timeLimit)
        self.model.setParam('MIPGap', GAP)
        self.model.optimize()
        print( 'obj ', self.model.objVal)


#%%
class SubproblemLS_SingleOO:
        def __init__(self, m, N, InputDf, duals, cols):
            self.N = N
            self.m =  m
            self.a = InputDf.loc[self.m,:].sort_values(by = 'P/W')
            
            # Convert to NumPy arrays for fast computation
            self.Weight_pd = self.a.Weight
            self.ProcessingTime_pd = self.a.ProcessingTime
            self.Weight = self.a.Weight.values  # NumPy array
            self.ProcessingTime = self.a.ProcessingTime.values  # NumPy array
            
            self.duals = duals
            self.Jobs = self.a.index.tolist()
            self.cols = cols
            
            # Precompute dual values as NumPy array for vectorization
            self.job_duals = np.array([duals[j] for j in self.Jobs])
            self.machine_dual = duals[self.m + self.N]
            
        def ScheduleCost(self, output_schedule):
            # Use pandas for this since output_schedule contains job IDs
            return sum(self.ProcessingTime_pd[output_schedule].cumsum()*self.Weight_pd[output_schedule])
        
        def smith_order(self, S_m):
            return sorted(S_m, key=lambda j: self.ProcessingTime_pd[j] / self.Weight_pd[j])

        def compute_contribution_vectorized(self, active_mask):
            """
            Vectorized computation of all job contributions.
            active_mask: boolean array indicating which jobs are in S
            Returns: array of contributions for each active job
            """
            n = len(active_mask)
            active_indices = np.where(active_mask)[0]
            
            if len(active_indices) == 0:
                return np.array([])
            
            # Get active processing times and weights
            p_active = self.ProcessingTime[active_mask]
            w_active = self.Weight[active_mask]
            
            cumsum_p = np.cumsum(p_active)  # Cumulative sum of processing times
            cumsum_w = np.cumsum(w_active)  # Cumulative sum of weights
            total_w = cumsum_w[-1]
            
            # First term: w[i] * cumsum_p[i]
            first_term = w_active * cumsum_p
            
            # Second term: p[i] * (total_w - cumsum_w[i])
            second_term = p_active * (total_w - cumsum_w)
            
            contributions = first_term + second_term
            
            return contributions

        def reduced_cost_vectorized(self, active_mask):
            """Vectorized reduced cost calculation"""
            p_active = self.ProcessingTime[active_mask]
            w_active = self.Weight[active_mask]
            duals_active = self.job_duals[active_mask]
            
            cumsum_p = np.cumsum(p_active)
            total_wct = np.sum(w_active * cumsum_p)
            sum_duals = np.sum(duals_active)
            
            return total_wct - sum_duals - self.machine_dual

        def local_search(self):
            if self.cols == 1:
                # Optimized path for single column generation
                n = len(self.Jobs)
                active_mask = np.ones(n, dtype=bool)  # All jobs initially active
                
                while True:
                    # Compute contributions for all active jobs (vectorized)
                    contributions = self.compute_contribution_vectorized(active_mask)
                    
                    if len(contributions) == 0:
                        break
                    
                    # Subtract duals from contributions
                    work_values = contributions - self.job_duals[active_mask]
                    
                    # Find job with maximum work value
                    max_work_idx = np.argmax(work_values)
                    max_work = work_values[max_work_idx]
                    
                    if max_work > 0:
                        # Remove this job (convert relative index to absolute)
                        active_indices = np.where(active_mask)[0]
                        job_to_remove = active_indices[max_work_idx]
                        active_mask[job_to_remove] = False
                    else:
                        break
                
                # Build final schedule
                S = [self.Jobs[i] for i in range(n) if active_mask[i]]
                
                # Compute reduced cost and objective
                rc = self.reduced_cost_vectorized(active_mask)
                obj = self.ScheduleCost(S)
                
                return [S], [obj], [rc]
            
            else:
                # Original path for multiple columns (cols > 1)
                S = self.Jobs
                improvement = True
                Schedules = []
                rc_Schedules = []
                obj_Schedules = []
                
                while improvement and len(Schedules) <= 5:
                    WORK = self.update_work_column(S)
                    improvement = False
                    best_cont_j = WORK.index(max(WORK))
                    if WORK[best_cont_j] > 0:
                        new_S = [job for i, job in enumerate(S) if i != best_cont_j]
                        improvement = True
                        rc = self.reduced_cost(new_S)
                        if rc < -1:
                            obj = self.ScheduleCost(new_S)
                            Schedules.append(new_S)
                            rc_Schedules.append(rc)
                            obj_Schedules.append(obj)
                        S = new_S
                    else:
                        break
                
                combined = list(zip(rc_Schedules, obj_Schedules, Schedules))
                top5 = heapq.nsmallest(5, combined, key=lambda x: x[0])
                rc_Schedules, obj_Schedules, Schedules = zip(*top5) if top5 else ([], [], [])
                
                return list(Schedules), list(obj_Schedules), list(rc_Schedules)
        
        # Keep old methods for cols > 1 path
        def compute_contribution(self, j, S):
            idx = S.index(j)
            first_term = self.Weight_pd[j] * sum(self.ProcessingTime_pd[k] for k in S[:idx+1])
            second_term = self.ProcessingTime_pd[j] * sum(self.Weight_pd[k] for k in S[idx+1:])
            return first_term + second_term

        def update_work_column(self, S):
            WORK = []
            for j in S:
                WORK.append(self.compute_contribution(j, S) - self.duals[j])
            return WORK
        
        def reduced_cost(self, S):
            cumulative_time = 0
            total_wct = 0
            sum_duals_jobs = 0
            for j in S:
                cumulative_time += self.ProcessingTime_pd[j]
                total_wct += self.Weight_pd[j] * cumulative_time
                sum_duals_jobs += self.duals[j]
            rc = total_wct - sum_duals_jobs - self.duals[self.m + self.N]
            return rc

#%%
class Subproblem: 
    #!/usr/bin/env pypy
    def __init__(self, m, N, InputDf, duals, numberSchedules):
        self.N = N
        self.m =  m
        self.a = InputDf.loc[self.m,:].sort_values(by = 'P/W')
        self.Weight = self.a.Weight
        self.ProcessingTime = self.a.ProcessingTime
        self.duals = duals
        self.B = self.a.B
        self.new =[]
        self.numberSchedules = numberSchedules
        self.T = sum(self.ProcessingTime)  
        self.index_series = pd.MultiIndex.from_product([list(range(1, self.N+1)),list(range(self.T+1))], names = ['J', 'T'])
        self.temp = dict.fromkeys(self.index_series, np.inf)
        self.schedule = dict.fromkeys(self.index_series, 0)
        
    
    def ScheduleCost(self, output_schedule):
 
        return sum(self.ProcessingTime[output_schedule].cumsum()*self.Weight[output_schedule])
    
    def DP(self):              
        @lru_cache(maxsize = 10000000) 
       
        def F(j,t):
      
            if t < -0.0000001 and j in range(self.N+1):
                return np.inf, 0
            elif j == 0 and t >= 0:
                return 0, 0
            else:
                
                temp = {i:F(i, t - self.ProcessingTime[j])[0] + t*self.Weight[j] - self.duals[j] for i in self.B[j] +[0]}

                return min(temp.values()), min(temp, key=temp.get)
            
        
        
        for j in self.a.index: 
            t = self.ProcessingTime[j]
            while t <= self.T+1:   
                    output = F(j,t)                  
                    self.temp[j,t] = output[0]
                    if self.B[j] != []:
                       self.schedule[j,t] = output[1]
                    t += 1

        self.temp = dict(sorted(self.temp.items(), key = itemgetter(1))[:self.numberSchedules])
    
        ll = []
        for jt, value in self.temp.items():
            if value -  self.duals[self.m+self.N] < -0.001  :   
    
                output_schedule = [jt[0]]
                new_t = jt[1] - self.ProcessingTime[jt[0]]
                new_j = self.schedule[jt[0], jt[1]]
                while new_j > 0 and new_t > 0:
                                output_schedule[0:0] = [new_j]
                                temp_t = new_t
                                new_t = temp_t - self.ProcessingTime[new_j]
                                new_j = self.schedule[new_j, temp_t]
                ll.append({'Machine':self.m, 'JobsInSchedule': output_schedule, 
                                      'ScheduleCost': self.ScheduleCost(output_schedule)})           
            else:
                 break                       
        return ll
#%%%
#%% Original IP model
class OriginalModel:
    def __init__(self, M, N,  InputDf):
        self.N = N
        self.M = M
        self.Weight = InputDf.loc[1,:].Weight
        self.ProcessingTime = InputDf['ProcessingTime']
        #self.A = ScheduleDf[['Job' + str(i) for i in range(1, self.N+1)]]
        self.A = InputDf['A']
        self.B = InputDf['B']
        self.model = Model("OriginalProblem")
        self.model.Params.LogToConsole = 0
    def buildModel(self):
        self.generateVariables()
        self.generateConstraints()
        self.generateObjective()
        self.model.update()
        
    def writeLP(self):
        self.model.write("gurobi_new.lp")
        
    def generateVariables(self):
        self.Xijk = self.model.addVars(range(1,self.N+1), range(1, self.N+1), range(1, self.M+1), 
                                                vtype = GRB.BINARY, name = "Xijk")
        
        self.X0jk = self.model.addVars(range(1, self.N+1), range(1, self.M+1), 
                                                vtype = GRB.BINARY, name = "X0jk")
        self.Xjlastk = self.model.addVars(range(1, self.N+1), range(1, self.M+1), 
                                                vtype = GRB.BINARY, name = "Xjlastk")
        
        self.C = self.model.addVars(range(1, self.N+1), vtype = GRB.CONTINUOUS, name = 'C')
    def generateConstraints(self):
        #for j in range(1, self.N+1):
        self.model.addConstrs((quicksum(self.Xijk[i, j, k] for k in range(1, self.M+1) for i in self.B[k][j]) + quicksum(self.X0jk[j,k] for k in range(1, self.M+1)) 
                                 == 1  for j in range(1, self.N+1)), 'J')
                    
        #for k in range(1, self.M+1):
        self.model.addConstrs((quicksum(self.X0jk[j,k] for j in range(1, self.N+1))
                                  <= 1 for k in range(1, self.M+1)),  'M') 
       
       #for j in range(1, self.N+1):
        #    for k in range(1, self.M+1):
        self.model.addConstrs((quicksum(self.Xijk[i,j,k]  for i in self.B[k][j]) + self.X0jk[j,k] - (quicksum(self.Xijk[j,i,k] for i in self.A[k][j]) + self.Xjlastk[j,k])
                                  == 0 for j in range(1, self.N+1) for k in range(1, self.M+1)),  'FlowCons')
        
       # for j in range(1, self.N+1):
        self.model.addConstrs((self.C[j] - quicksum(self.ProcessingTime[k][j]*self.X0jk[j,k] + quicksum((self.C[i] + 
                                                        self.ProcessingTime[k][j])*self.Xijk[i,j,k] for i in self.B[k][j]) for k in range(1, self.M+1)) == 0) for j in range(1, self.N+1))
        
    def generateObjective(self):     
         self.model.setObjective(quicksum(self.Weight[j]* self.C[j] for j in range(1, self.N+1)),
                                sense = GRB.MINIMIZE)
    
    def solveModel(self, timeLimit, GAP):
        self.model.setParam('TimeLimit', timeLimit)
        self.model.setParam('MIPGap', GAP)
        self.model.optimize()
        if self.model.status == GRB.OPTIMAL:
            return True 
        
    def getObj(self):
           obj = self.model.getObjective()
           return (obj.getValue())

        
  
    def getGap(self):
        return self.model.MIPGap
      
    