import json
from gurobipy import Model, GRB, quicksum

model = Model("OptimizationProblem")

with open("data.json", "r") as f:
    data = json.load(f)

# Define the parameters
TotalBatteryCapacity = data["TotalBatteryCapacity"]  # shape: [], definition: The total capacity of the battery
InitialSoc = data["InitialSoc"]  # shape: [], definition: Initial state of charge of the battery
InitialPg1 = data["InitialPg1"]  # shape: [], definition: Initial power of generator 1
InitialPg2 = data["InitialPg2"]  # shape: [], definition: Initial power of generator 2
InitialPg3 = data["InitialPg3"]  # shape: [], definition: Initial power of generator 3

Price = data["Price"]  # shape: [24], definition: Electricity price at each time step
Load = data["Load"]  # shape: [24], definition: Residential load at each time step
Pv = data["Pv"]  # shape: [24], definition: Photovoltaic power generation at each time step
Wind = data["Wind"]  # shape: [24], definition: Wind power generation at each time step

# Define the variables

Ag1 = model.addVars(24, vtype=GRB.CONTINUOUS, name="Ag1")
Ag2 = model.addVars(24, vtype=GRB.CONTINUOUS, name="Ag2")
Ag3 = model.addVars(24, vtype=GRB.CONTINUOUS, name="Ag3")
PvPower = model.addVars(24, vtype=GRB.CONTINUOUS, name="PvPower")
Apv = model.addVars(24, vtype=GRB.CONTINUOUS, name="Apv")
Asoc = model.addVars(24, vtype=GRB.CONTINUOUS, name="Asoc")

ExcessEnergy = model.addVars(24, vtype=GRB.CONTINUOUS, name="ExcessEnergy")
ShortageEnergy = model.addVars(24, vtype=GRB.CONTINUOUS, name="ShortageEnergy")

TotalPowerSupply = model.addVars(24, vtype=GRB.CONTINUOUS, name="TotalPowerSupply")
GridTransactionLimit = model.addVar(vtype=GRB.FLOAT, name="GridTransactionLimit")

# Define the constraints

# Loop over each time step to add the soc constraints
for t in range(24):
    # Calculate soc_t based on the previous soc and the adjustment factor
    soc_t = 0.18 * Asoc[t] + InitialSoc

    # Add the lower bound constraint for soc_t
    model.addConstr(soc_t >= 0.2)
    # Add the upper bound constraint for soc_t
    model.addConstr(soc_t <= 0.8)

    # Update soc to the newly calculated soc_t for the next iteration
    soc = soc_t

# Assuming decision variables Pg1[t] have been defined for each time step t
for t in range(24):
    if t == 0:
        # For the first time step, use the initial value of Pg1
        model.addConstr(Pg1[t] <= 150)
        model.addConstr(Pg1[t] >= 10)
    else:
        # For subsequent time steps, calculate Pg1 based on the previous time step and Ag1
        model.addConstr(Pg1[t] <= 150)
        model.addConstr(Pg1[t] >= 10)
        model.addConstr(Pg1[t] == 100 * Ag1[t] + Pg1[t - 1])
for t in range(1, 24):
    model.addConstr(InitialPg2 + 100 * Ag2[t] >= 50)
    model.addConstr(InitialPg2 + 100 * Ag2[t] <= 375)
for t in range(24):
    model.addConstr(Pg3[t] >= 100)  # Lower bound constraint for each time step
    model.addConstr(Pg3[t] <= 500)  # Upper bound constraint for each time step
# Assuming soc is defined as a list of decision variables for each time step
soc = model.addVars(24, lb=0.2, ub=0.8, vtype=GRB.CONTINUOUS, name="soc")

# Loop over each time step starting from 1 up to 24
for t in range(1, 24):
    # Add the constraint for the state of charge calculation
    model.addConstr(soc[t] == 0.18 * Asoc[t] + soc[t - 1])
for t in range(1, 24):
    model.addConstr(pg1[t] == 100 * Ag1[t] + pg1[t - 1])
# Assuming that `pg2` is a list of Gurobi variables representing the power of Generator 2 at each time step
# and that `InitialPg2` is a parameter representing the initial power of Generator 2 at time step 0

# Define the Generator 2 power for the first time step, which is the initial value and does not have a constraint
pg2 = [model.addVar(vtype=GRB.CONTINUOUS, lb=50, ub=375, name="pg2_0")]

# Now add the constraints for time steps 1 to 23
for t in range(1, 24):
    pg2.append(model.addVar(vtype=GRB.CONTINUOUS, lb=50, ub=375, name=f"pg2_{t}"))
    model.addConstr(pg2[t] == 100 * Ag2[t - 1] + pg2[t - 1])
# Assuming Pg3 is defined as a list of decision variables representing power of generator 3 at each time step
for t in range(1, 24):
    model.addConstr(Pg3[t] == 200 * Ag3[t] + Pg3[t - 1])
for t in range(24):
    model.addConstr(PvPower[t] == Pv[t] * (1 + Apv[t]))
for t in range(24):
    model.addConstr(ExcessEnergy[t] >= 0)
for t in range(24):
    # Define the total power supply at time step t
    TotalSupply_t = (InitialPg1 + Ag1[t] * 100 +
                     InitialPg2 + Ag2[t] * 100 +
                     InitialPg3 + Ag3[t] * 200 +
                     Pv[t] * (1 + Apv[t]) +
                     Wind[t] +
                     InitialSoc + Asoc[t] * 0.18)
    # Add the constraint for excess energy at time step t
    model.addConstr(ExcessEnergy[t] == TotalSupply_t - Load[t])
for t in range(24):
    # Calculate updated generator powers and battery state of charge
    soc_t = InitialSoc + 0.18 * Asoc[t] if t == 0 else soc[t - 1] + 0.18 * Asoc[t]
    pg1_t = InitialPg1 + 100 * Ag1[t] if t == 0 else pg1[t - 1] + 100 * Ag1[t]
    pg2_t = InitialPg2 + 100 * Ag2[t] if t == 0 else pg2[t - 1] + 100 * Ag2[t]
    pg3_t = InitialPg3 + 200 * Ag3[t] if t == 0 else pg3[t - 1] + 200 * Ag3[t]

    # Calculate photovoltaic power
    PvPower_t = Pv[t] * (1 + Apv[t])

    # Add the constraint for total power supply at time t
    model.addConstr(TotalPowerSupply[t] == soc_t + pg1_t + pg2_t + pg3_t + PvPower_t + Wind[t])
for t in range(24):
    model.addConstr(ShortageEnergy[t] >= 0)
for t in range(24):
    model.addConstr(-1 <= Apv[t])
    model.addConstr(Apv[t] <= 1)
    model.addConstr(-1 <= Asoc[t])
    model.addConstr(Asoc[t] <= 1)
    model.addConstr(-1 <= Ag1[t])
    model.addConstr(Ag1[t] <= 1)
    model.addConstr(-1 <= Ag2[t])
    model.addConstr(Ag2[t] <= 1)
    model.addConstr(-1 <= Ag3[t])
    model.addConstr(Ag3[t] <= 1)
for t in range(1, 24):
    model.addConstr(100 * Asoc[t] == soc[t] - soc[t - 1])
for t in range(24):
    # Calculate the power from each generator based on previous state and adjustment
    pg1_power = InitialPg1 + 100 * Ag1[t]
    pg2_power = InitialPg2 + 100 * Ag2[t]
    pg3_power = InitialPg3 + 200 * Ag3[t]

    # Calculate the adjusted photovoltaic power
    pv_power = Pv[t] * (1 + Apv[t])

    # Express the total power supply
    TotalPowerSupply[t] = pg1_power + pg2_power + pg3_power + pv_power + Wind[t]

    # Add the constraint that total power supply cannot exceed the load by more than the grid transaction limit
    model.addConstr(TotalPowerSupply[t] <= Load[t] + GridTransactionLimit)
for t in range(24):
    # Assuming that Pg1, Pg2, Pg3 are defined as variables in the model with the correct initial values and adjustments
    Pg1[t] = model.addVar(vtype=GRB.CONTINUOUS, name=f"Pg1_{t}")
    Pg2[t] = model.addVar(vtype=GRB.CONTINUOUS, name=f"Pg2_{t}")
    Pg3[t] = model.addVar(vtype=GRB.CONTINUOUS, name=f"Pg3_{t}")

    # Define the auxiliary constraint for each time step
    model.addConstr(TotalPowerSupply[t] == Pg1[t] + Pg2[t] + Pg3[t] + PvPower[t] + Wind[t])

# Define the objective
model.setObjective(
    quicksum(
        (Asoc[t] + 0.1 * soc[t - 1] +
         0.0034 * pg1[t - 1] ** 2 + 3 * pg1[t - 1] + 30 +
         0.001 * pg2[t - 1] ** 2 + 10 * pg2[t - 1] + 40 +
         0.001 * pg3[t - 1] ** 2 + 15 * pg3[t - 1] + 70 +
         0.01 * (pv[t - 1] + wind[t - 1]) +
         max(0, (ExcessEnergy[t] - 100) * 50) +
         max(0, (ShortageEnergy[t] - 100) * 50) -
         0.5 * Price[t - 1] * ExcessEnergy[t] +
         Price[t - 1] * ShortageEnergy[t] +
         abs(ExcessEnergy[t]) + abs(ShortageEnergy[t])
         ) for t in range(1, 25)
    ),
    GRB.MINIMIZE
)

# Optimize the model
model.optimize()

# Output optimal objective value
print("Optimal Objective Value: ", model.objVal)

if model.status == GRB.OPTIMAL:
    with open("output_solution.txt", "w") as f:
        f.write(str(model.objVal))
    print("Optimal Objective Value: ", model.objVal)
else:
    with open("output_solution.txt", "w") as f:
        f.write(model.status)