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

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

# 总数据量和周期数
total_period = len(data)
days = total_period // 24

# 存储所有结果
all_results = []

for day in range(days):
    # 获取每天的24条数据
    start_idx = day * 24
    end_idx = (day + 1) * 24
    daily_price = price[start_idx:end_idx]
    daily_load = load[start_idx:end_idx]
    daily_pv = pv[start_idx:end_idx]
    daily_wind = wind[start_idx:end_idx]

    # 初始化模型
    m = Model("OptimizationProblem_Day{}".format(day))

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

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

    # 设置交换约束
    grid_energy_export = m.addVars(24, vtype=GRB.CONTINUOUS, lb=0, ub=100, name='export')
    grid_energy_import = m.addVars(24, vtype=GRB.CONTINUOUS, lb=0, ub=100, name='import')
    excess_energy = m.addVars(24, vtype=GRB.CONTINUOUS, lb=0, name='excess')
    shortage_energy = m.addVars(24, 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, 24)), name="soc_update")
    m.addConstrs((pg1[t] == pg1[t - 1] + 100 * Ag1[t] for t in range(1, 24)), name="pg1_update")
    m.addConstrs((pg2[t] == pg2[t - 1] + 100 * Ag2[t] for t in range(1, 24)), name="pg2_update")
    m.addConstrs((pg3[t] == pg3[t - 1] + 200 * Ag3[t] for t in range(1, 24)), name="pg3_update")

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

    # 供需平衡约束
    m.addConstrs((pg1[t] + pg2[t] + pg3[t] + pv_power[t] + daily_wind.iloc[t] + grid_energy_import[t] ==
                  daily_load.iloc[t] + 0.18 * Asoc[t] + grid_energy_export[t] + excess_energy[t] - shortage_energy[t]
                  for t in range(24)), name="power_balance")

    # excess_energy 和 shortage_energy 约束
    m.addConstrs(
        (excess_energy[t] == pg1[t] + pg2[t] + pg3[t] + pv_power[t] + daily_wind.iloc[t] + grid_energy_import[t]
         - daily_load.iloc[t] - 0.18 * Asoc[t] - grid_energy_export[t] for t in range(1, 24)),
        name="excess_energy")
    m.addConstrs((shortage_energy[t] == daily_load.iloc[t] + 0.18 * Asoc[t] + grid_energy_export[t]
                  - pg1[t] - pg2[t] - pg3[t] - pv_power[t] - daily_wind.iloc[t] - grid_energy_import[t]
                  for t in range(1, 24)), name="shortage_energy")

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

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

    # 设置目标函数
    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()

    if m.status == GRB.OPTIMAL:
        day_results = []
        for t in range(24):
            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.01 * daily_wind.iloc[t]
            excess_penalty = excess_energy[t].x * 50
            shortage_penalty = shortage_energy[t].x * 50
            sales_revenue = 0.5 * daily_price.iloc[t] * grid_energy_export[t].x
            purchase_cost = daily_price.iloc[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)

            day_results.append([
                Asoc[t].x,
                Ag1[t].x,
                Ag2[t].x,
                Ag3[t].x,
                Apv[t].x,
            ])

        # 保存当天的结果
        all_results.extend(day_results)

# 最后将所有结果保存到一个JSON文件
with open("./results_day.json", "w") as f:
    json.dump(all_results, f, indent=4)
