from pulp import *
import time
import argparse
import os
from util import Util

class LowCarbonCloudLP:
    
    def __init__(self,name):
        parser = argparse.ArgumentParser(description='Low-carbon cloud sizing tool')
        parser.add_argument('--input_file', metavar = 'input_file', type=str,
                            help = 'filename with the inputs')
        parser.add_argument('--use_gurobi', help = 'use the gurobi solver', action='store_true')
        parser.add_argument('--only_write_lp_file', help = 'Write the .lp file to be used with an external solver', action='store_true')
        parser.add_argument('--write_sol_file', help = 'Write the variables and their values to a .csv file', action='store_true')
        res = parser.parse_args()  
        
        self.args = res
        self.use_gurobi = res.use_gurobi
        self.only_write_lp_file = res.only_write_lp_file
        self.input_file = res.input_file

        if self.input_file is None or self.input_file == '':
          raise RuntimeError('Need to pass a input file!')

        if not os.path.exists(self.input_file):
          raise RuntimeError('Input file must exist!')

        inputs = Util.load_inputs(self.input_file)    

        self.name = name
        # a list containing the timeslots
        self.timeslots = inputs['timeslots']  
         # a list containing the DCs names
        self.DCs = inputs['DCs']             
        # a list containing the workload. Each element of the list represents the total CPU demand in the time slot of the element's index
        self.workload = inputs['workload']    
        # a hash containing the solar irradiation data for each DC. 
        # the key is the DC name and the value is a list.
        # each element of the list represents solar irradiation (W/m²) in the time slot of the element's index
        self.irradiance = inputs['irradiance'] 
        # a hash containing the total number of avaliable CPU cores for each DC. 
        # the key is the DC name and the value is the number of cores.
        self.C = inputs['C']
        # a hash containing the carbon emissions of using the PVS for each DC. 
        # the key is the DC name and the value is the carbon emissions (g CO2 eq/kWh)
        self.pv_co2 =  inputs['pv_co2']
        # a hash containing the carbon emissions of using the local electricity grid for each DC. 
        # the key is the DC name and the value is the carbon emissions (g CO2 eq/kWh)
        self.grid_co2 = inputs['grid_co2']
        # a hash containing the idle power consumption of each DC. 
        # the key is the DC name and the value is the power consumption  (W)
        self.pIdleDC = inputs['pIdleDC']
        # a hash containing the power consumption from the network devices of each DC. 
        # the key is the DC name and the value is the power consumption  (W)        
        self.pNetIntraDC = inputs['pNetIntraDC']
        # a hash containing the Power Usage Effectiveness (PUE) of each DC. 
        # the key is the DC name and the value is the PUE
        self.PUE = inputs['PUE']
        # the dynamic power consumption of using 1 CPU core (W)
        self.pCore = inputs['pCore']
        # carbon footprint of manufacturing batteries (g CO2 eq/ kWh)
        self.bat_co2= inputs['bat_co2']
        # efficiency of the battery charging process
        self.eta_ch = inputs['eta_ch']
        # efficiency of the battery discharging process
        self.eta_dch= inputs['eta_dch']
        # PV panel efficiency of converting solar irradiation into electricity 
        self.eta_pv =inputs['eta_pv']

        ### Variables    
        self.A   = LpVariable.dicts("A",self.DCs, 0,    cat='Continuous')                           # PV panels area (m²)
        self.Bat = LpVariable.dicts("Bat",self.DCs, 0, cat='Continuous')                            # Battery capacity in Wh
        self.B  = LpVariable.dicts('B_', (self.DCs,self.timeslots),lowBound=0, cat='Continuous')     # Level of energy
        self.Pdch = LpVariable.dicts('Pdch_', (self.DCs,self.timeslots),lowBound=0,cat='Continuous') # Power to discharge (W)
        self.Pch = LpVariable.dicts('Pch_', (self.DCs,self.timeslots),lowBound=0, cat='Continuous')  # Power to charge (W)
        self.w = LpVariable.dicts('w_', (self.DCs,self.timeslots),lowBound=0, cat='Continuous')  # workload to be sent to each DC
        self.Pgrid = LpVariable.dicts('Pgrid_', (self.DCs,self.timeslots),lowBound=0, cat='Continuous')     # Green power surplus sold back to the grid
        
        

    ### Auxiliary functions
    def getDCPowerConsumption(self,d,k):
        """Compute the power consumption of DC d at time slot k
        considering the static power (network and idle server costs),
        the dynamic power ( execution of the workload), and the power
        from cooling the DC (PUE)

        Parameters
        ----------
        d : str
            The name of the DC
        k : int
            The time slot number

        Returns
        -------
        float
            The power consumption of DC d at time slot k in kWh
        """

        if(k ==0):
            return 0
        return self.PUE[d] * (self.pNetIntraDC[d]+ self.pIdleDC[d]  + self.w[d][k] * self.pCore ) * 1/1000
    
    def FPgrid(self,d,k):
        """Compute the carbon emissions from using the local 
        electricity grid at the location of DC d at time slot k

        Parameters
        ----------
        d : str
            The name of the DC
        k : int
            The time slot number

        Returns
        -------
        float
            The carbon emissions in g CO2 eq/kWh
        """        
        return self.grid_co2[d] * self.Pgrid[d][k]

    def FPpv(self,d,k):
        """Compute the carbon emissions from using the power 
        produced from the PVs of DC d at time slot k

        Parameters
        ----------
        d : str
            The name of the DC
        k : int
            The time slot number

        Returns
        -------
        float
            The carbon emissions in g CO2 eq/kWh
        """                
        return self.Pre(d,k) * self.pv_co2[d]

    def FPbat(self,d):
        """Compute the carbon emissions of manufacturing 
        batteries at DC d

        Parameters
        ----------
        d : str
            The name of the DC
            
        Returns
        -------
        float
            The carbon emissions in g CO2 eq/kWh
        """                
        return self.Bat[d] * self.bat_co2

    def P(self,d,k):
        """Compute the power consumption of DC d at time slot k

        Parameters
        ----------
        d : str
            The name of the DC
        k : int
            The time slot number

        Returns
        -------
        float
            The power consumption of DC d at time slot k in kWh
        """

        return self.getDCPowerConsumption(d,k) 

    def Pre(self,d,k):
        """Compute the renewable power production of DC d at time slot k

        Parameters
        ----------
        d : str
            The name of the DC
        k : int
            The time slot number

        Returns
        -------
        float
            The renewable power consumption of DC d at time slot k in kW
        """
        return (self.A[d] * self.irradiance[d][k] * self.eta_pv) / 1000 # convert to kw


    ### Objective functions
    def use_original_objective_function(self,prob):    
        prob +=  lpSum(
            [  self.FPgrid(d,k)  + self.FPpv(d,k)  for k in self.timeslots ]  + self.FPbat(d) for d in self.DCs)  
        return prob        
        
    ### This is a extension of the previous objective function and was created to reduce 
    # the number of charge and discharge simultaneously events, and it dont change the optimal value 
    def use_Pdch_obj_function(self,prob):    
        prob +=  lpSum(
            [  self.FPgrid(d,k)  + self.FPpv(d,k)  + self.Pdch[d][k]*self.pv_co2[d]*0.0000000001 for k in self.timeslots ]  + self.FPbat(d) for d in self.DCs)  
        return prob

    
    def build_lp(self):        
        """Build the LP object by adding the restrictions 
        between the variables

        Returns
        -------
        an LPProblem object
            The LP modeled to be solved using a solver
        """        
        prob = LpProblem(self.name, LpMinimize)


        ### Initialization process
        for d in self.DCs:
            prob.addConstraint( self.Pch[d][0]   == 0.0 )
            prob.addConstraint( self.Pdch[d][0]   == 0.0 )

        for d in self.DCs:
            for k in self.timeslots[1:] :
                ### Restriction for the battery level of energy
                prob.addConstraint( self.B[d][k]  == self.B[d][k-1] + self.Pch[d][k]*self.eta_ch  - self.Pdch[d][k]*self.eta_dch )
                ### Restriction for how much power can be charged into the battery
                prob.addConstraint( self.Pch[d][k]  * self.eta_ch <= 0.8 * self.Bat[d] - self.B[d][k-1] )  
                ### Restriction for how much power can be discharged from the battery
                prob.addConstraint( self.Pdch[d][k] * self.eta_dch <= self.B[d][k-1] -  0.2 * self.Bat[d] )


        for d in self.DCs:
            for k in self.timeslots :
                ### Restriction for the power consumption of the DC:
                ### it cannot be higher than the current renewable production, or discharged from the batteries, or
                ### using the grid as backup
                prob.addConstraint( self.P(d,k) <= self.Pgrid[d][k] + self.Pre(d,k) + self.Pdch[d][k] - self.Pch[d][k]  )
                ### Restriction for emissions of using the grid
                prob.addConstraint( self.FPgrid(d,k) >= 0 )
                ### Cannot charge more power than what is being produced by the PV panels
                prob.addConstraint( self.Pre(d,k) >= self.Pch[d][k]   )

                ### Limit the usage of the batteries to prolong their life
                prob.addConstraint( self.B[d][k]   >= 0.2 * self.Bat[d] )
                prob.addConstraint( self.B[d][k]   <= 0.8 * self.Bat[d] )

                ### Cannot allocate more workload than the DCs core capacity
                prob.addConstraint( self.w[d][k]<=self.C[d])

        for k in self.timeslots:
                ### All the workload must be executed
                prob +=  lpSum([  self.w[d][k]   for d in self.DCs]) == self.workload[k]    
        
        return prob




def write_sol_file(filename,dict_variables):
    """ Write the solution of the LP to an external file
    
    Parameters
    ----------
    filename : str
        The path were the file will be written
    dict_variables : hash
        The variables of the LP and their computed values from the solver

    """

    if not os.path.exists('results/'+filename):
        os.mkdir('results/'+filename)
    result_file = open(f'results/{filename}/solution.csv', 'w')    
    for d in dict_variables:                        
        result_file.write(f'{d};{round(dict_variables[d],2)}\n')                                                          
    result_file.close()

def write_summary(total_emissions,runtime,input_file,output_path):
    """ Write the summary file of running the experiment. This
    file will contain the total carbon emissions (objective value),
    time it took to execute the experiments, and which input file 
    was used to run the experiment.
    
    Parameters
    ----------
    total_emissions : float 
        The total carbon emissions (ton CO2 eq) from manufacturing the PVs
        batteries and operating the cloud federation
    runtime : float
        Total execution time of the experiment (in seconds)
    input_file : str
        The file it was used as input for the experiments
    output_path : str
        The path were the file will be written        
    """
    result_file = open(output_path, 'w')    
    result_file.write('input_file;total_emissions;runtime\n')
    result_file.write(f'{input_file};{total_emissions};{runtime}\n')
    result_file.close()


def main():
    lpModel = LowCarbonCloudLP("Low_carbon_cloud")
    lpProb = lpModel.build_lp()
    lpProb = lpModel.use_Pdch_obj_function(lpProb)

    ### The problem data is written to an .lp file
    if(lpModel.only_write_lp_file):
        lpProb.writeLP(lpModel.input_file.replace('.json','') + '.lp')
        return
 
    start = time.time()

    ### Default PuLP solver:
    solver =  PULP_CBC_CMD()

    ### Needs further configuration. Check https://coin-or.github.io/pulp/guides/how_to_configure_solvers.html for details.
    if(lpModel.use_gurobi):
       solver = GUROBI_CMD()
   
    lpProb.solve(solver)
    
    #### The status of the solution is printed to the screen
    print("Status:", LpStatus[lpProb.status])
 
    #### The optimised objective function value is printed to the screen
    emissions = round(value(lpProb.objective),2) # here we have the value in gramms 
    emissions = emissions / 1000000 # here we have the value in tons 
    print("Total emissions of the cloud = ", emissions)
    
    end = time.time()
    runtime = round(end - start,2)
    print(f"Executed in {end - start} s")

    variables = Util.extract_dict_variables(lpProb)
    file_name = lpModel.input_file

    folder_name = file_name.replace('.json','').replace('input/', '')
    write_sol_file(folder_name,variables)
    write_summary(emissions,runtime,file_name, f'results/{folder_name}/summary_results.csv')



if __name__ == '__main__':
    main()