# -*- coding: utf-8 -*-
"""
fig_dgcup2022a.py —— 电工杯 2022 A 题范文配图（24 张 SVG）
================================================
复用 gen_dgcup2022a 的确定性真源，生成三篇范文各 8 张图，共 24 张，
写入 assets/problems/dgcup2022a/figs/，命名 dgcup2022a-{paper}-fig{n}.svg。

用法：cd tools/ && python3 fig_dgcup2022a.py
"""
import os
import gen_dgcup2022a as G
from _svg import (_fig, _save, _bar, _grouped_bar, _line, _scatter,
                  _pie, _heatmap, _stack, _network, _flow, _txt,
                  C_ACC, C_RED, C_GREEN, C_AMB, C_PUR, C_CYAN, C_TEAL, C_PINK, C_GRID, C_MUT, PALETTE)

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.normpath(os.path.join(HERE, "..", "assets", "problems", "dgcup2022a", "figs"))
os.makedirs(OUT, exist_ok=True)

D = G.gen_dgcup2022a()
load_arr = D["load_arr"]
load_norm = D["load_norm"]
wind_norm = D["wind_norm"]
DT = G.DT
hours = [k * DT for k in range(96)]
CARBON = G.CARBON_PRICES


def step_series(units_idx, la, wa):
    res = {i: [] for i in units_idx}
    wu, cu, lo = [], [], []
    for load, w in zip(la, wa):
        Ps, w_used, curt, loss = G.dispatch_no_storage(units_idx, load, w)
        for i in units_idx:
            res[i].append(Ps[i])
        wu.append(w_used)
        cu.append(curt)
        lo.append(loss)
    return res, wu, cu, lo


# ===================== 范文一：基础建模与无储能分析 =====================
def p1_f1():
    s = [("负荷(标幺)", C_ACC, [(hours[k], load_norm[k]) for k in range(96)]),
         ("风电(标幺)", C_GREEN, [(hours[k], wind_norm[k]) for k in range(96)])]
    return _line("图1 典型日负荷与风电归一化出力曲线", s, xlabel="时刻 h", ylabel="标幺值")


def p1_f2():
    q = D["q1_plan"]
    s = [("机组1", C_ACC, [(hours[k], q[1][k]) for k in range(96)]),
         ("机组2", C_RED, [(hours[k], q[2][k]) for k in range(96)]),
         ("机组3", C_GREEN, [(hours[k], q[3][k]) for k in range(96)])]
    return _line("图2 问题一无风电时三机组日发电计划 MW", s, xlabel="时刻 h", ylabel="出力 MW")


def p1_f3():
    q = D["q1"][60]
    return _pie("图3 问题一发电成本构成(碳价60元/t)",
                ["火电煤耗", "运行维护", "碳捕集"],
                [q["coal"], q["om"], q["carbon"]])


def p1_f4():
    res, wu, cu, lo = step_series([1, 2], load_arr, [v * 300 for v in wind_norm])
    layers = [("机组1", C_ACC, res[1]), ("机组2", C_RED, res[2]),
              ("风电消纳", C_GREEN, wu), ("弃风", C_AMB, cu)]
    return _stack("图4 问题二功率平衡(风电300MW替代机组3)", hours, layers, ymax=1000)


def p1_f5():
    _, _, cu, _ = step_series([1, 2], load_arr, [v * 300 for v in wind_norm])
    s = [("弃风功率 MW", C_AMB, [(hours[k], cu[k]) for k in range(96)])]
    return _line("图5 问题二逐时弃风功率 MW", s, xlabel="时刻 h", ylabel="弃风 MW")


def p1_f6():
    res, wu, cu, lo = step_series([1, 3], load_arr, [v * 600 for v in wind_norm])
    layers = [("机组1", C_ACC, res[1]), ("机组3", C_GREEN, res[3]),
              ("风电消纳", C_TEAL, wu), ("弃风", C_AMB, cu), ("失负荷", C_RED, lo)]
    return _stack("图6 问题三功率平衡(风电600MW替代机组2)", hours, layers, ymax=1000)


def p1_f7():
    q2, q3 = D["q2"], D["q3"]
    groups = ["问题二(300MW)", "问题三(600MW)"]
    cats = ["弃风电量", "失负荷电量"]
    vals = [[q2["curt_kwh"], q2["loss_kwh"]], [q3["curt_kwh"], q3["loss_kwh"]]]
    return _grouped_bar("图7 问题二/三 弃风与失负荷电量对比 kWh", groups, cats, vals, fmt="%.0f")


def p1_f8():
    vals = [D["q1"][0]["unit_supply"], D["q2"]["unit_supply"], D["q3"]["unit_supply"]]
    return _bar("图8 问题一/二/三 单位供电成本(碳价0) 元/kWh",
                ["问题一(无风)", "问题二(300MW)", "问题三(600MW)"], vals,
                colors=[C_ACC, C_GREEN, C_TEAL], fmt="%.4f")


# ===================== 范文二：碳捕集、储能配置与单位成本 =====================
def p2_f1():
    q4 = D["q4"]
    groups = ["问题二场景", "问题三场景"]
    cats = ["0", "60", "80", "100"]
    vals = [[q4["Q2"][c] for c in CARBON], [q4["Q3"][c] for c in CARBON]]
    return _grouped_bar("图1 四档碳价下单位供电成本 元/kWh", groups, cats, vals, fmt="%.4f")


def p2_f2():
    w2 = [v * 300 for v in wind_norm]
    vals = []
    for cp in CARBON:
        tot, _, _, _, _ = G.scenario_cost([1, 2], load_arr, w2, cp)
        vals.append(tot["carbon"])
    return _bar("图2 问题二碳捕集成本随碳价变化 元", [str(c) for c in CARBON], vals,
                colors=[C_PUR] * 4, fmt="%.0f")


def p2_f3():
    _, _, _, lo = step_series([1], load_arr, [v * 900 for v in wind_norm])
    s = [("失负荷功率 MW", C_RED, [(hours[k], lo[k]) for k in range(96)])]
    return _line("图3 问题五无储能时逐时失负荷 MW", s, xlabel="时刻 h", ylabel="失负荷 MW")


def p2_f4():
    q5 = D["q5"]
    s = [("可充电功率 MW", C_GREEN, [(hours[k], q5["S"][k]) for k in range(96)]),
         ("需放电功率 MW", C_RED, [(hours[k], q5["D"][k]) for k in range(96)])]
    return _line("图4 问题五储能充放电功率计划 MW", s, xlabel="时刻 h", ylabel="功率 MW")


def p2_f5():
    q5 = D["q5"]
    return _bar("图5 问题五最小储能配置",
                ["储能容量(MWh)", "储能功率(MW)"],
                [q5["sto_capacity"], q5["sto_power"]],
                colors=[C_ACC, C_PUR], fmt="%.1f")


def p2_f6():
    vals = [D["q5"]["unit_supply_no_sto"], D["q5"]["unit_supply"]]
    return _bar("图6 问题五含储能前后单位供电成本 元/kWh",
                ["无储能(碳价60)", "含储能(碳价60)"], vals, colors=[C_RED, C_GREEN], fmt="%.4f")


def p2_f7():
    q4 = D["q4"]
    s = [("问题二场景", C_GREEN, [(c, q4["Q2"][c]) for c in CARBON]),
         ("问题三场景", C_RED, [(c, q4["Q3"][c]) for c in CARBON])]
    return _line("图7 单位供电成本随碳价变化 元/kWh", s, xlabel="碳价 元/t", ylabel="元/kWh")


def p2_f8():
    w5 = [v * 900 for v in wind_norm]
    tot, _, _, _, _ = G.scenario_cost([1], load_arr, w5, 60)
    stor = D["q5"]["sto_daily"]
    labels = ["火电煤耗", "运行维护", "碳捕集", "风电运维", "弃风损失", "储能成本"]
    vals = [tot["coal"], tot["om"], tot["carbon"], tot["wind"], tot["curt"], stor]
    return _bar("图8 问题五含储能成本构成 元/天", labels, vals,
                colors=[C_ACC, C_CYAN, C_PUR, C_GREEN, C_AMB, C_PINK], fmt="%.0f")


# ===================== 范文三：灵敏度、Q6/Q7 与方案设计 =====================
def p3_f1():
    ser = D["q6"]["series"]
    s = [("单位供电成本", C_ACC, [(r["inst"], r["unit_supply"]) for r in ser])]
    return _line("图1 单位供电成本随风电装机变化 元/kWh", s, xlabel="风电装机 MW", ylabel="元/kWh")


def p3_f2():
    ser = D["q6"]["series"]
    s = [("弃风电量 kWh", C_AMB, [(r["pen"], r["curt_kwh"]) for r in ser]),
         ("失负荷电量 kWh", C_RED, [(r["pen"], r["loss_kwh"]) for r in ser])]
    return _line("图2 弃风与失负荷随渗透率变化", s, xlabel="风电渗透率", ylabel="电量 kWh")


def p3_f3():
    q5 = D["q5"]
    mults = [0.5, 1.0, 2.0, 3.0]
    tp = sum(q5["S"]) * DT * 1000 + sum(q5["D"]) * DT * 1000
    base = q5["total_with_sto"] - q5["sto_daily"]
    us = []
    for m in mults:
        inv = (G.STO_PC * m) * q5["sto_power"] * 1000 + (G.STO_EC * m) * q5["sto_capacity"] * 1000
        dinv = inv / (G.STO_LIFE * 365)
        dom = tp * G.STO_OM
        us.append((base + dinv + dom) / D["load_kwh_day"])
    s = [("单位供电成本", C_ACC, [(m, us[i]) for i, m in enumerate(mults)])]
    return _line("图3 单位供电成本对储能成本的敏感性", s, xlabel="储能成本倍数", ylabel="元/kWh")


def p3_f4():
    rows = [0,  # placeholder replaced below
            ]
    insts = [0, 300, 600, 900, 1200]
    M = []
    for inst in insts:
        if inst == 0:
            uidx = [1, 2, 3]
        elif inst <= 300:
            uidx = [1, 2]
        elif inst <= 600:
            uidx = [1, 3]
        else:
            uidx = [1]
        wa = [v * inst for v in wind_norm]
        row = []
        for cp in CARBON:
            tot, _, _, _, _ = G.scenario_cost(uidx, load_arr, wa, cp)
            row.append(tot["unit_supply"])
        M.append(row)
    return _heatmap("图4 风电装机×碳价 单位供电成本热力图 元/kWh",
                    [str(i) for i in insts], [str(c) for c in CARBON], M, fmt="%.3f")


def p3_f5():
    load15 = D["q7"]["load_arr"]
    wind15 = D["q7"]["wind_arr"]
    idx = list(range(0, 1440, 4))
    s = [("负荷 MW", C_ACC, [(k * DT, load15[k]) for k in idx]),
         ("风电 MW", C_GREEN, [(k * DT, wind15[k]) for k in idx])]
    return _line("图5 问题七十五日负荷与风电 MW", s, xlabel="时刻 h", ylabel="功率 MW")


def p3_f6():
    load15 = D["q7"]["load_arr"]
    wind15 = D["q7"]["wind_arr"]
    end = 72
    net = [load15[k] - wind15[k] for k in range(end)]
    s = [("负荷", C_ACC, [(k * DT, load15[k]) for k in range(end)]),
         ("风电", C_GREEN, [(k * DT, wind15[k]) for k in range(end)]),
         ("净负荷(负荷-风电)", C_RED, [(k * DT, net[k]) for k in range(end)])]
    return _line("图6 问题七前3日净负荷(替代机组2、3) MW", s, xlabel="时刻 h", ylabel="功率 MW")


def p3_f7():
    load15 = D["q7"]["load_arr"]
    wind15 = D["q7"]["wind_arr"]
    _, _, _, lo = step_series([1], load15, wind15)
    idx = list(range(0, 1440, 4))
    s = [("失负荷 MW", C_RED, [(k * DT, lo[k]) for k in idx])]
    return _line("图7 问题七逐时失负荷(十五日) MW", s, xlabel="时刻 h", ylabel="失负荷 MW")


def p3_f8():
    load15 = D["q7"]["load_arr"]
    wind15 = D["q7"]["wind_arr"]

    def us_scenario(uidx, with_sto):
        tot, _, _, _, _ = G.scenario_cost(uidx, load15, wind15, 60)
        cost = tot["total"]
        if with_sto:
            st = G.storage_size(uidx, load15, wind15)
            tp = sum(st["S"]) * DT * 1000 + sum(st["D"]) * DT * 1000
            sc = G.storage_daily_cost(st["power_MW"], st["capacity_MWh"], tp)
            cost = cost - tot["loss"] + sc["daily_total"]   # 储能消除失负荷
        lk = sum(load15) * DT * 1000
        return cost / lk

    vals = [us_scenario([1], False),
            us_scenario([1], True),
            us_scenario([1, 2], True),
            us_scenario([1, 2, 3], True)]
    return _bar("图8 问题七不同方案单位供电成本 元/kWh",
                ["无措施", "仅储能", "储能+恢复机组2", "储能+恢复机组2,3"],
                vals, colors=[C_RED, C_GREEN, C_CYAN, C_ACC], fmt="%.4f")


# ===================== 任务表 =====================
JOBS = [p1_f1, p1_f2, p1_f3, p1_f4, p1_f5, p1_f6, p1_f7, p1_f8,
        p2_f1, p2_f2, p2_f3, p2_f4, p2_f5, p2_f6, p2_f7, p2_f8,
        p3_f1, p3_f2, p3_f3, p3_f4, p3_f5, p3_f6, p3_f7, p3_f8]


def save(name, svg):
    _save(os.path.join(OUT, name), svg)


def main():
    n_bad = 0
    for i, job in enumerate(JOBS):
        paper = (i // 8) + 1
        fig = (i % 8) + 1
        name = "dgcup2022a-%d-fig%d.svg" % (paper, fig)
        try:
            svg = job()
            if svg and "<svg" in svg:
                save(name, svg)
            else:
                n_bad += 1
                print("BAD:", name)
        except Exception as e:
            n_bad += 1
            print("ERR:", name, repr(e))
    print("生成 SVG 共 %d 张，失败 %d 张，目录: %s" % (len(JOBS), n_bad, OUT))


if __name__ == "__main__":
    main()
