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

#### Util functions


class LowCarbonCloudLP:


    def __init__(self, name, timeslots, DCs, workload, irradiance, C, pv_co2, grid_co2, pIdleDC, pNetIntraDC, PUE, pCore, bat_co2, eta_ch, eta_dch, eta_pv):
        """
        Initialize the LP model with the inputs.
        :param str name: Name of the problem
        :param timeslots: List containing the time slots
        :param DCs: List containing the name of the Data Centers (DC)
        :param workload: List containing the workload that will be executed at each time slot (number of CPU cores)
        :param irradiance: Dictionary, where the key is the DC name and the value is a list with solar irradiation values
        :param C: Dictionary, where the key is the DC name and the value is a list with the CPU cores capacity
        :param pv_co2: Dictionary, where the key is the DC name and the value is the emissions from pvs based in the DC location (g co2 / kWh )
        :param grid_co2: Dictionary, where the key is the DC name and the value is the emissions from local electricity grid based in the DC location
        :param pIdleDC: Dictionary, where the key is the DC name and the value is how much watts its servers consume when IDLE   
        :param pNetIntraDC: Dictionary, where the key is the DC name and the value is how much watts its network equipment consume when IDLE
        :param PUE: The data center Power Usage Effectiveness value
        :param pCore: Dynamic costs of computing a server
        :param bat_co2: Emissions of manufacturing the battery (g co2 / kWh of capacity)
        :param eta_ch: The effiency of the charging process into the batteries
        :param eta_dch: The effiency of the discharging process into the batteries
        :param eta_pv:  The effiency of converting solar irradiation into energy in the PVs   
        
        """
        self.name = name
        self.timeslots = timeslots
        self.DCs = DCs
        self.workload = workload
        self.irradiance = irradiance
        self.C = C
        self.pv_co2 =  pv_co2
        self.grid_co2 = grid_co2
        self.pIdleDC = pIdleDC
        self.pNetIntraDC = pNetIntraDC
        self.PUE = PUE
        self.pCore = pCore
        self.bat_co2= bat_co2
        self.eta_ch = eta_ch
        self.eta_dch= eta_dch
        self.eta_pv =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
        self.loe = LpVariable.dicts("loe_", DCs, 0,    cat='Continuous')                     # DC Level of energy
        
        
    ### Auxiliary functions
    def getDCPowerConsumption(self,d,k):
        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):
        return self.grid_co2[d] * self.Pgrid[d][k]

    def FPpv(self,d,k):
        return self.Pre(d,k) * self.pv_co2[d]

    def FPbat(self,d):
        return self.Bat[d] * self.bat_co2

    def P(self,d,k):
        return self.getDCPowerConsumption(d,k) 

    def Pre(self,d,k):
        return (self.A[d] * self.irradiance[d][k] * self.eta_pv) / 1000 # convert to kw

    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        
        
    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):        
        prob = LpProblem(self.name, LpMinimize)

        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:] :
                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 )
                prob.addConstraint( self.Pch[d][k]  * self.eta_ch <= 0.8 * self.Bat[d] - self.B[d][k-1] )   # To ensure that only charge/discharge if
                prob.addConstraint( self.Pdch[d][k] * self.eta_dch <= self.B[d][k-1] -  0.2 * self.Bat[d] ) # battery capacity >= 0 


        for d in self.DCs:
            for k in self.timeslots :
                prob.addConstraint( self.P(d,k) <= self.Pgrid[d][k] + self.Pre(d,k) + self.Pdch[d][k] - self.Pch[d][k]  )
                prob.addConstraint( self.FPgrid(d,k) >= 0 )
                prob.addConstraint( self.Pre(d,k) >= self.Pch[d][k]   )
                prob.addConstraint( self.B[d][k]   >= 0.2 * self.Bat[d] )
                prob.addConstraint( self.B[d][k]   <= 0.8 * self.Bat[d] )
                prob.addConstraint( self.w[d][k]<=self.C[d])

        for k in self.timeslots:
                prob +=  lpSum([  self.w[d][k]   for d in self.DCs]) == self.workload[k]    
        
        return prob




def write_sol_file(filename,dict_variables):
    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};{dict_variables[d]}\n')                                                          
    result_file.close()

def write_summary(total_emissions,runtime,input_file,output_path):
    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():
    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()  
    
    use_gurobi = res.use_gurobi
    only_write_lp_file = res.only_write_lp_file
    input_file = res.input_file# +'.json'
    inputs = Util.load_inputs(input_file)    
    lpModel = LowCarbonCloudLP("Green_DC", inputs['timeslots'] , inputs['DCs'] , inputs['workload'], inputs['irradiance'], inputs['C'], inputs['pv_co2'], inputs['grid_co2'], inputs['pIdleDC'] , inputs['pNetIntraDC'], inputs['PUE'], inputs['pCore'], inputs['bat_co2'], inputs['eta_ch'], inputs['eta_dch'],inputs['eta_pv'])
    
    lpProb = lpModel.build_lp()
    lpProb = lpModel.use_Pdch_obj_function(lpProb)

    ### The problem data is written to an .lp file
    if(only_write_lp_file):
        lpProb.writeLP(res.input_file.replace('.json','') + '.lp')
        return
 
    start = time.time()
    solver =  PULP_CBC_CMD()
 
    if(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 = 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()