import json
import pandas as pd
from gurobipy import Model, GRB, quicksum

# 读取数据
data = pd.read_csv("data.csv")
price = data['price']
load = data['load']
pv = data['pv']
wind = data['wind']

# 初始化模型
m = Model("OptimizationProblem")
period = len(data)

# 决策变量
Ag1 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=-1, ub=1, name="ag1")
Ag2 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=-1, ub=1, name="ag2")
Ag3 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=-1, ub=1, name="ag3")
Apv = m.addVars(period, vtype=GRB.CONTINUOUS, lb=-1, ub=1, name="apv")
Asoc = m.addVars(period, vtype=GRB.CONTINUOUS, lb=-1, ub=1, name="asoc")

# 定义变量相关的约束
soc = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0.2, ub=0.8, name="soc")
pg1 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, ub=150, name="pg1")
pg2 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, ub=375, name="pg2")
pg3 = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, ub=500, name="pg3")
pv_power = m.addVars(period, vtype=GRB.CONTINUOUS, name="pv_power")

# 设置交换约束
grid_energy_export = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, ub=100, name='export')
grid_energy_import = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, ub=100, name='import')
excess_energy = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, name='excess')
shortage_energy = m.addVars(period, vtype=GRB.CONTINUOUS, lb=0, name='shortage')

# 初始值约束
m.addConstr(soc[0] == 0.2, name="soc_init")
m.addConstr(pg1[0] == 0, name="pg1_init")
m.addConstr(pg2[0] == 0, name="pg2_init")
m.addConstr(pg3[0] == 0, name="pg3_init")

# 添加soc和pg的递推约束
m.addConstrs((soc[t] == soc[t - 1] + 0.18 * Asoc[t] for t in range(1, period)), name="soc_update")
m.addConstrs((pg1[t] == pg1[t - 1] + 100 * Ag1[t] for t in range(1, period)), name="pg1_update")
m.addConstrs((pg2[t] == pg2[t - 1] + 100 * Ag2[t] for t in range(1, period)), name="pg2_update")
m.addConstrs((pg3[t] == pg3[t - 1] + 200 * Ag3[t] for t in range(1, period)), name="pg3_update")

# 光伏发电功率
m.addConstrs((pv_power[t] == pv[t] * (1 + Apv[t]) for t in range(period)), name="pv_power_update")

# 定义供需平衡约束
m.addConstrs((pg1[t] + pg2[t] + pg3[t] + pv_power[t] + wind[t] + grid_energy_import[t] ==
              load[t] + 0.18 * Asoc[t] + grid_energy_export[t] + excess_energy[t] - shortage_energy[t] for t in range(period)),
             name="power_balance")
m.addConstrs((excess_energy[t] == pg1[t] + pg2[t] + pg3[t] + pv_power[t] + wind[t] + grid_energy_import[t]
              - load[t] - 0.18 * Asoc[t] - grid_energy_export[t] for t in range(1, period)),
             name="excess_energy")
m.addConstrs((shortage_energy[t] == load[t] + 0.18 * Asoc[t] + grid_energy_export[t]
              - pg1[t] - pg2[t] - pg3[t] - pv_power[t] - wind[t] - grid_energy_import[t] for t in range(1, period)),
             name="shortage_energy")

# 定义成本函数
battery_cost = quicksum((Asoc[t] + 0.1 * soc[t]) for t in range(1, period))
gen1_cost = quicksum((0.0034 * pg1[t] ** 2 + 3 * pg1[t] + 30) for t in range(1, period))
gen2_cost = quicksum((0.001 * pg2[t] ** 2 + 10 * pg2[t] + 40) for t in range(1, period))
gen3_cost = quicksum((0.001 * pg3[t] ** 2 + 15 * pg3[t] + 70) for t in range(1, period))
renewable_cost = quicksum((0.01 * pv_power[t] + 0.01 * wind[t]) for t in range(period))

# 计算惩罚和收益
excess_penalty = quicksum(excess_energy[t] * 50 for t in range(period))
shortage_penalty = quicksum(shortage_energy[t] * 50 for t in range(period))
sales_revenue = quicksum(0.5 * price[t] * grid_energy_export[t] for t in range(period))
purchase_cost = quicksum(price[t] * grid_energy_import[t] for t in range(period))

# 设置目标函数
objective = (battery_cost + gen1_cost + gen2_cost + gen3_cost + renewable_cost
             + excess_penalty + shortage_penalty - sales_revenue + purchase_cost)

m.setObjective(objective, GRB.MINIMIZE)
m.optimize()
# m.computeIIS()
# m.write("file.ilp")

if m.status == GRB.OPTIMAL:
    results = []
    for t in range(period):
        gen1_cost = 0.0034 * pg1[t].x ** 2 + 3 * pg1[t].x + 30
        gen2_cost = 0.001 * pg2[t].x ** 2 + 10 * pg2[t].x + 40
        gen3_cost = 0.001 * pg3[t].x ** 2 + 15 * pg3[t].x + 70
        battery_cost = Asoc[t].x + 0.1 * soc[t].x
        renewable_cost = 0.01 * pv_power[t].x + 0.1 * wind[t]
        excess_penalty = excess_energy[t].x * 50
        shortage_penalty = shortage_energy[t].x * 50
        sales_revenue = 0.5 * price[t] * grid_energy_export[t].x
        purchase_cost = price[t] * grid_energy_import[t].x

        current_obj = (gen1_cost + gen2_cost + gen3_cost + battery_cost + renewable_cost +
                       excess_penalty + shortage_penalty - sales_revenue + purchase_cost)

        results.append({
            "time": t,
            "asoc": Asoc[t].x,
            "ag1": Ag1[t].x,
            "ag2": Ag2[t].x,
            "ag3": Ag3[t].x,
            "apv": Apv[t].x,
            "reward": current_obj,
            "unbanlance": excess_energy[t].x + shortage_energy[t].x
        })

    with open("../../DRL-for-Energy-Systems/data/results_day.json", "w") as f:
        json.dump(results, f, indent=4)