"""Canonical assumption registry and figure generator for the regulator paper.

Single source of truth for every number quoted in
research/spirited-energy-regulator-paper.md is research/assumptions.yaml,
loaded into REGISTRY by this module. Each value carries an evidence class:
    O  = official public dataset ingested in this repository
    C  = computed in this repository and verified by automated test
    E  = engineering estimate from project research, not externally verified
    A  = assumption requiring confirmation (counterparty, agency, or filing)
    PD = characterization of public/PUB disclosures; verify against docket

Figures are derived from these values only. Anything not in this registry
must not appear as a number in the paper.
"""

import os
from pathlib import Path

import matplotlib
import yaml

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch

OUT = os.path.join(os.path.dirname(__file__), "rendered", "regulator-paper")
os.makedirs(OUT, exist_ok=True)

# --- Canonical registry -----------------------------------------------------


def _normalize_value(value):
    """Convert YAML list pairs to tuples for parity with the old REGISTRY."""
    if isinstance(value, list) and len(value) == 2 and all(isinstance(x, int) for x in value):
        return tuple(value)
    return value


def load_registry_from_yaml(path) -> dict:
    """Load assumption registry from YAML; returns {key: (value, evidence, note)}."""
    data = yaml.safe_load(Path(path).read_text())
    registry = {}
    for key, entry in data.items():
        if key.startswith("_"):
            continue
        registry[key] = (
            _normalize_value(entry["value"]),
            entry["evidence"],
            entry["note"],
        )
    return registry


REGISTRY = load_registry_from_yaml(Path(__file__).with_name("assumptions.yaml"))


def v(key):
    return REGISTRY[key][0]


def winter_system_load_mw(temp_c: float) -> float:
    """Research heating model anchored so P(winter_design_temp_c) = peak_record_mw.

    P(T) = base + max(0, (thresh - T) * beta)  [E form, PD peak anchor]
    """
    base = float(v("winter_load_base_mw"))
    thresh = float(v("winter_heat_threshold_c"))
    beta = float(v("winter_heat_beta_mw_per_c"))
    return base + max(0.0, (thresh - float(temp_c)) * beta)


def summer_system_load_mw(temp_c: float) -> float:
    """Fig13-style summer load sketch; P(+35C) matches summer_peak_anchor_mw."""
    base = float(v("summer_load_base_mw"))
    thresh = float(v("summer_heat_threshold_c"))
    beta = float(v("summer_heat_beta_mw_per_c"))
    return base + max(0.0, (float(temp_c) - thresh) * beta)


def summer_load_with_pdrc_mw(temp_c: float) -> float:
    """Unmanaged summer load minus ramped PDRC shave (flat BTES applied separately)."""
    onset = float(v("summer_pdrc_onset_c"))
    beta = float(v("summer_pdrc_beta_mw_per_c"))
    return summer_system_load_mw(temp_c) - max(0.0, (float(temp_c) - onset) * beta)


def year_round_typical_peak_unmanaged_mw() -> list[int]:
    """Monthly typical peak-hour load from ECCC normals and the existing P(T) models.

    Heating months (mean T < winter_heat_threshold_c) use mean daily minimum
    with winter_system_load_mw. Cooling months use mean daily maximum with
    summer_system_load_mw. Rounded MW for the year-round figure. Not metered
    Hydro demand and not an export forecast.
    """
    means = list(v("winnipeg_normals_mean_c"))
    mins = list(v("winnipeg_normals_min_c"))
    maxs = list(v("winnipeg_normals_max_c"))
    thresh = float(v("winter_heat_threshold_c"))
    out = []
    for mean_c, min_c, max_c in zip(means, mins, maxs):
        if float(mean_c) < thresh:
            load = winter_system_load_mw(min_c)
        else:
            load = summer_system_load_mw(max_c)
        out.append(int(round(load)))
    return out


def year_round_typical_peak_managed_mw() -> list[int]:
    """Unmanaged typical peak-hour load minus the recovery tools that apply that month.

    Heating: hub_net_electric_mw all-day if the loop is running; batteries only
    when mean daily min is at or below year_round_bess_overlay_min_c (peak-hour
    tool, not monthly energy). Cooling: PDRC ramp at monthly mean max plus
    btes_summer_chiller_shave_mw. Not the +35 C heatwave identity. Does not
    enter 413 / 688 / 788.
    """
    means = list(v("winnipeg_normals_mean_c"))
    mins = list(v("winnipeg_normals_min_c"))
    maxs = list(v("winnipeg_normals_max_c"))
    thresh = float(v("winter_heat_threshold_c"))
    hub = float(v("hub_net_electric_mw"))
    bess = float(v("bess_mw"))
    bess_min = float(v("year_round_bess_overlay_min_c"))
    btes = float(v("btes_summer_chiller_shave_mw"))
    out = []
    for mean_c, min_c, max_c in zip(means, mins, maxs):
        if float(mean_c) < thresh:
            load = winter_system_load_mw(min_c) - hub
            if float(min_c) <= bess_min:
                load -= bess
        else:
            load = summer_load_with_pdrc_mw(max_c) - btes
        out.append(int(round(max(0.0, load))))
    return out


# --- Shared style -----------------------------------------------------------

GREEN = "#14532d"
MIDGREEN = "#2e7d32"
BLUE = "#1565c0"
GRAY = "#6b7280"
ROSE = "#b91c1c"
AMBER = "#b45309"
plt.rcParams.update({
    "font.family": "DejaVu Sans",
    "font.size": 10,
    "axes.edgecolor": "#9ca3af",
    "axes.linewidth": 0.8,
})


def _box(ax, xy, w, h, text, fc, ec, fontsize=9.5, textcolor="#111827"):
    ax.add_patch(FancyBboxPatch(xy, w, h, boxstyle="round,pad=0.012",
                                fc=fc, ec=ec, lw=1.6))
    ax.text(xy[0] + w / 2, xy[1] + h / 2, text, ha="center", va="center",
            fontsize=fontsize, color=textcolor, fontweight="bold")


def _arrow(ax, a, b, color=GRAY, style="-|>"):
    ax.add_patch(FancyArrowPatch(a, b, arrowstyle=style, mutation_scale=16,
                                 lw=2.0, color=color))


# --- Figure 1: graphical abstract ------------------------------------------

def fig1_graphical_abstract():
    fig = plt.figure(figsize=(12, 5.6))
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")

    ax.text(0.5, 0.945, "Use Every Electron Twice: The NEWPCC Clean Hybrid Portfolio",
            ha="center", fontsize=15, fontweight="bold", color=GREEN)

    xs = [0.030, 0.275, 0.520, 0.765]
    w, h = 0.205, 0.21
    fs = 8.6

    # Row 1: the heat path
    _box(ax, (xs[0], 0.60), w, h,
         "Hydro reservoirs\nstore the\nmulti-day energy", "#e3f2fd", BLUE, fontsize=fs)
    _box(ax, (xs[1], 0.60), w, h,
         f"Compute nodes,\n{v('node_band_kw')[0]}–{v('node_band_kw')[1]} kW each\n(0 MW until 3 gates pass)",
         "#ecfdf5", MIDGREEN, fontsize=fs)
    _box(ax, (xs[2], 0.60), w, h,
         f"NEWPCC thermal hub,\n{v('hub_thermal_mwth')} MWth at full build\n(effluent + compute heat)",
         "#ecfdf5", MIDGREEN, fontsize=fs)
    _box(ax, (xs[3], 0.60), w, h,
         "Winnipeg buildings\nheated by water, not\nby the winter peak", "#fff7ed", AMBER, fontsize=fs)

    # Row 2: the electric ledger
    _box(ax, (xs[0], 0.26), w, h,
         f"Batteries: {v('bess_mw')} MW, 4 h —\nshave the two daily\npeak spikes", "#e3f2fd", BLUE, fontsize=fs)
    _box(ax, (xs[1], 0.26), w, h,
         f"Demand response {v('dsm_mw')} MW;\ninterties up to {v('intertie_mw')} MW\n(accreditation pending)",
         "#e3f2fd", BLUE, fontsize=fs)
    _box(ax, (xs[2], 0.26), w, h,
         f"Hub, electric ledger:\n≤ {v('hub_net_electric_mw')} MW net\n(COP 4.13, minus pumps)",
         "#e3f2fd", BLUE, fontsize=fs)
    _box(ax, (xs[3], 0.26), w, h,
         f"Firm capability:\n{v('scenario_conservative_mw')}–{v('scenario_full_mw')} MW by\naccreditation scenario",
         "#dcfce7", GREEN, fontsize=9.2)

    for y in (0.705, 0.365):
        for a, b in zip(xs[:-1], xs[1:]):
            _arrow(ax, (a + w, y), (b, y), BLUE if y < 0.5 else MIDGREEN)
    _arrow(ax, (xs[3] + w / 2, 0.60), (xs[3] + w / 2, 0.47), AMBER)

    ax.text(0.5, 0.135,
            f"Compared with Brandon gas: ${v('gas_capex_b'):.2f}B for {v('gas_nameplate_mw')} MW nameplate, with "
            f"outages or deratings in {v('gas_cold_events_outage')[0]} of the {v('gas_cold_events_outage')[1]} "
            "highest-demand winter events [PD]",
            ha="center", fontsize=9.5, color="#374151")
    ax.text(0.5, 0.075,
            f"≥{v('remediation_floor_pct')}% of compute surplus flows to Lake Winnipeg remediation as a first lien",
            ha="center", fontsize=9.5, color="#374151")
    ax.text(0.5, 0.022,
            "Thermal and electric megawatts kept on separate ledgers; every number derives from the canonical registry",
            ha="center", fontsize=8.3, style="italic", color=GRAY)

    fig.savefig(os.path.join(OUT, "fig1_graphical_abstract.png"), dpi=200)
    plt.close(fig)


# --- Figure 2: capacity scenarios -------------------------------------------

def fig2_capacity_scenarios():
    fig, ax = plt.subplots(figsize=(9, 5.2))
    labels = [
        "Hybrid\nconservative",
        "Hybrid\nreference",
        "Hybrid\nfull accreditation",
        "Gas\nnameplate",
        "Gas\nderate scenario",
    ]
    values = [v("scenario_conservative_mw"), v("scenario_reference_mw"),
              v("scenario_full_mw"), v("gas_nameplate_mw"), v("gas_derate_scenario_mw")]
    colors = ["#86efac", "#34d399", MIDGREEN, "#9ca3af", "#d1d5db"]
    bars = ax.bar(labels, values, color=colors, edgecolor="#374151", width=0.62)
    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, val + 12, f"{val} MW",
                ha="center", fontweight="bold", fontsize=10)
    ax.axhline(v("shortfall_2030_mw"), color=ROSE, linestyle="--", lw=1.8)
    ax.text(4.45, v("shortfall_2030_mw") + 14, "2030 shortfall (600 MW) [PD]",
            ha="right", color=ROSE, fontsize=9)
    ax.set_ylabel("Winter resource stack (MW)")
    ax.set_title("Planning cases vs Brandon gas",
                 fontweight="bold", fontsize=12, color=GREEN)
    ax.text(0.005, -0.16,
            "Electric ledger only; hub 50 MWth enters as ≤38 MW net electric, never as thermal MW. "
            "688 includes 200 MW undemonstrated winter import.\n"
            "Gas derate is a scenario pending the independent ELCC study requested in §11.",
            transform=ax.transAxes, fontsize=8.2, color=GRAY, va="top")
    fig.tight_layout(rect=(0, 0.05, 1, 1))
    fig.savefig(os.path.join(OUT, "fig2_capacity_scenarios.png"), dpi=200)
    plt.close(fig)


# --- Figure 3: delivered cost matrix ----------------------------------------

def fig3_delivered_cost():
    fig, ax = plt.subplots(figsize=(9, 5.2))
    scenarios = ["Hybrid\nconservative", "Hybrid\nreference", "Hybrid\nfull", "Gas\nnameplate", "Gas\nderate\nscenario"]
    cap_mw = [v("scenario_conservative_mw"), v("scenario_reference_mw"), v("scenario_full_mw"),
              v("gas_nameplate_mw"), v("gas_derate_scenario_mw")]
    gross = [v("hybrid_capex_gross_b") * 1e9 / (m * 1000) for m in cap_mw[:3]] + \
            [v("gas_capex_b") * 1e9 / (m * 1000) for m in cap_mw[3:]]
    net = [v("hybrid_capex_net_b") * 1e9 / (m * 1000) for m in cap_mw[:3]] + [None, None]

    x = range(len(scenarios))
    ax.bar([i - 0.19 for i in x], gross, width=0.38, color="#93c5fd",
           edgecolor="#374151", label="Gross capital (no credits)")
    ax.bar([i + 0.19 for i in x[:3]], net[:3], width=0.38, color=MIDGREEN,
           edgecolor="#374151", label="Net of 30% ITC (if confirmed) [A]")
    for i, g in enumerate(gross):
        ax.text(i - 0.19, g + 60, f"${g:,.0f}", ha="center", fontsize=8.8, fontweight="bold")
    for i, n in enumerate(net[:3]):
        ax.text(i + 0.19, n + 60, f"${n:,.0f}", ha="center", fontsize=8.8,
                fontweight="bold", color=GREEN)
    ax.set_xticks(list(x))
    ax.set_xticklabels(scenarios)
    ax.set_ylabel("Capital cost per firm kW delivered ($/kW)")
    ax.set_title("Delivered-cost matrix: even gross and conservative, the hybrid competes;\n"
                 "at reference accreditation it wins outright", fontweight="bold",
                 fontsize=12, color=GREEN)
    ax.legend(loc="upper left", fontsize=9)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "fig3_delivered_cost.png"), dpi=200)
    plt.close(fig)


# --- Figure 4: 20-year cost bands -------------------------------------------

def fig4_lifecycle_bands():
    fig, ax = plt.subplots(figsize=(9, 4.6))
    plans = ["Brandon gas plan", "NEWPCC clean hybrid"]
    lows = [v("gas_20yr_low_b"), v("hybrid_20yr_low_b")]
    highs = [v("gas_20yr_high_b"), v("hybrid_20yr_high_b")]
    colors = ["#9ca3af", MIDGREEN]
    for i, (lo, hi, c) in enumerate(zip(lows, highs, colors)):
        ax.barh(i, hi - lo, left=lo, height=0.42, color=c, edgecolor="#374151")
        ax.text(lo - 0.08, i, f"${lo:.1f}B", ha="right", va="center", fontweight="bold")
        ax.text(hi + 0.08, i, f"${hi:.1f}B", ha="left", va="center", fontweight="bold")
    ax.set_yticks(range(len(plans)))
    ax.set_yticklabels(plans, fontsize=11)
    ax.set_xlim(0, 7.6)
    ax.set_xlabel("Modeled 20-year all-in cost range ($ billions)")
    ax.set_title("Twenty-year cost bands: the ranges do not overlap",
                 fontweight="bold", fontsize=12, color=GREEN)
    ax.text(0.005, -0.30,
            "Band drivers: fuel-price path, carbon-price trajectory, financing terms, ITC eligibility, heat-revenue "
            "realization [E/A].\nReplaces the two retired point ledgers (Appendix A). Savings band: "
            f"${v('gas_20yr_low_b')-v('hybrid_20yr_high_b'):.1f}–"
            f"${v('gas_20yr_high_b')-v('hybrid_20yr_low_b'):.1f}B.",
            transform=ax.transAxes, fontsize=8.2, color=GRAY, va="top")
    fig.tight_layout(rect=(0, 0.08, 1, 1))
    fig.savefig(os.path.join(OUT, "fig4_lifecycle_bands.png"), dpi=200)
    plt.close(fig)


if __name__ == "__main__":
    fig1_graphical_abstract()
    fig2_capacity_scenarios()
    fig3_delivered_cost()
    fig4_lifecycle_bands()
    # Consistency guards: derived scenario totals must match registry claims.
    assert v("scenario_full_mw") == v("bess_mw") + v("dsm_mw") + v("intertie_mw") + v("hub_net_electric_mw")
    assert v("scenario_reference_mw") == v("bess_mw") + v("dsm_mw") + v("intertie_reference_mw") + v("hub_net_electric_mw")
    assert v("scenario_conservative_mw") == v("bess_mw") + v("dsm_mw") // 2 + v("hub_effluent_stage_net_mw")
    assert abs(v("hybrid_capex_net_b") - (v("hybrid_capex_gross_b") - v("hybrid_itc_b"))) < 1e-9
    print(f"wrote 4 figures to {OUT}; registry consistency checks passed")