import gurobipy as grb
from gurobipy import quicksum
import time
import numpy as np
from functions import *

from solvers.CLSP import CLSP

class CLSP_S(CLSP):
    """
    Classe per la risoluzione del problema CLSP (Capacitated Lot-Sizing Problem) 
    con domanda stocastica
    """
    def __init__(self,**setting):
        self.name = "CLSP_S"
        self.setting = setting

    def populate(self, instance, tree, I0):
        """
        Questo metodo, a partire dai dati di input della classe, popola il modello definendone 
        variabili, vincoli e funzione obiettivo
        """
        #parametri del modello (costi, capacità massima)
        h=instance.h
        g=instance.g
        f=instance.f
        proc_time=instance.proc_time
        R=instance.R
        setup_times=instance.setup_times
        num_items=instance.num_items
        time_periods=instance.time_periods
        N=instance.N
        
        #parametri dell'albero di scenari
        p=probs(tree)
        T=time_period(tree)
        L=tree.get_leaves()
        a=predecessor(tree)
        omega=ancestor(tree)
        sigma=successors_t(tree)
        d=np.ndarray(shape=(num_items,N))
        for i in range(N):
            if i==0:
                d[:,i]=tree.get_history_node(i)
            else:
                d[:,i]=tree.get_history_node(i)[T[i],:]
        N_L=[]
        for i in range(N):
            if i not in(L):
                N_L.append(i)
                
        model = grb.Model(self.name)
        
        #DEFINIZIONE DELLE VARIABILI DEL MODELLO
        
        y = model.addMVar(
            shape=(num_items,time_periods,N),
            lb=0,
            vtype=grb.GRB.CONTINUOUS,
            name="y"
        )
        
        self.y=y
        
        z = model.addMVar(
            shape=(num_items,N),
            lb=0,
            vtype=grb.GRB.CONTINUOUS,
            name="z"
        )
        
        self.z=z
        
        I = model.addMVar(
            shape=(num_items,N),
            lb=0,
            vtype=grb.GRB.CONTINUOUS,
            name="I"
        )
        
        self.I=I
        
        s = model.addMVar(
            shape=(num_items,N),
            vtype=grb.GRB.BINARY,
            name="s"
        )
        
        self.s=s
        
        #DEFINIZIONE DEI VINCOLI DEL MODELLO
        
        model.addConstrs((I[i,a[n]]+quicksum(y[i,T[n],omega[n,t]] for t in range(T[n]))+y[i,T[n],n]==d[i,n]+I[i,n]-z[i,n]) for i in range(num_items) for n in range(1,N))
        
        model.addConstrs((I0[i]+y[i,T[0],0]==d[i,0]+I[i,0]-z[i,0]) for i in range(num_items))
        
        model.addConstrs((y[i,t,n]<=(max(d[i,j] for j in sigma[n,t]))*s[i,n]) for i in range(num_items) for n in range(N) for t in range(T[n]+1,time_periods))
        
        model.addConstrs((y[i,T[n],n]<=d[i,n]*s[i,n]) for i in range(num_items) for n in range(N))
        
        model.addConstrs((quicksum(quicksum(proc_time[i]*y[i,t,n] for t in range(T[n],time_periods)) for i in range(num_items))+ quicksum(setup_times[i]*s[(i,n)] for i in range(num_items))<=R) for n in range(N))
        
        #DEFINIZIONE DELLA FUNZIONE OBIETTIVO DEL MODELLO
        
        model.setObjective(quicksum(p[n]*quicksum(f[i]*s[i,n]+h[i]*I[i,n]+g[i]*z[i,n] for i in range(num_items)) for n in range(N))+quicksum(p[n]*quicksum(quicksum(h[i]*(t-T[n]) * y[i,t,n] for t in range(T[n]+1,time_periods)) for i in range(num_items)) for n in N_L), grb.GRB.MINIMIZE)
        
        model.update()
        return model

    def get_solution(self, instance, model, time_limit=None, gap=None, verbose=False):
        """
        Questo metodo risolve un modello definito, restituendo il valore della funzione obiettivo,
        la decisione di produzione all'istante iniziale, le corrispondenti variabili di set up,
        il tempo di risoluzione del modello e il valore del gap della soluzione trovata"""
        #gap (per criterio di arresto) del problemo misto intero
        if gap:
            model.setParam('MIPgap', gap)
        #tempo limite per la risoluzione del problemo misto intero
        if time_limit:
            model.setParam(grb.GRB.Param.TimeLimit, time_limit)
        #opzione per avere un output che mostra i dettagli della risoluzione
        #del modello o meno
        if verbose:
            model.setParam('OutputFlag', 1)
        else:
            model.setParam('OutputFlag', 0)
        if verbose:
            print ('Solving a model with: '+str(model.NumConstrs)+' constraints')
            print ('    and: ' +str(model.NumVars)+ ' variables')
        
        #risoluzione del modello e tempo computazionale
        start = time.time()
        model.optimize()
        end = time.time()
        comp_time = end - start
        
        num_items = instance.num_items
        time_periods=instance.time_periods
        N=instance.N
        
        #CALCOLO DELLE DECISIONI INIZIALI (QUANTITA' PRODOTTE PER OGNI ITEM AL TEMPO 0 Y E RELATIVE VARIABILI DI SETUP S)
        
        sol = np.zeros((num_items,time_periods,N))
        
        for i in range(num_items):
            for j in range(time_periods):
                for k in range(N):
                    sol[i,j,k] = self.y[i,j,k].X
                    
        
        Y=[0]*num_items
        S=[0]*num_items
        for i in range(num_items):
            Y[i]=sum(sol[i,:,0])
            if Y[i]>0:
                S[i]=1

        #VALORE DELLA FUNZIONE OBIETTIVO        
        of = model.getObjective().getValue()

        #VALORE DEL GAP ALLA SOLUZIONE TROVATA
        gap=model.MIPGap
          
        model.reset()


        return of, Y, S, comp_time,gap


    def solve(
        self, instance, tree, I0, time_limit=None, gap=None, verbose=False
    ):
        """
        Questo è il metodo dove il problema viene definito e risolto
        :param instance: dizionario contenente tutti i parametri necessari alla definizione del modello
        :param tree: albero di scenari che rappresenta la domanda stocastica
        :I0: valore iniziale del magazzino
        :param time_limit: per interrompere il solver gurobi dopo "time_limit" secondi; se non è presente
        il solver si interrompe quando trova una soluzione 'distante meno di' "gap" dalla 
        soluzione ottima
        :param gap: il solver gurobi si ferma quando trova una soluzione 'distante meno di' "gap"
        dalla soluzione ottima; il valore standard è 0.0001
        :param verbose: parametro (True/False) da passare per avere un output che stampa i dettagli della 
        soluzione del modello o meno
        :return: valore della funzione obiettivo, decisioni al primo stadio (Y e S), tempo di calcolo e
        gap della soluzione trovata da quella ottima
        """
        model = self.populate(instance, tree, I0)
        return self.get_solution(instance, model, time_limit=time_limit, gap=gap, verbose=verbose)