# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Illustrative simulation of the buyer–seller–lender model (welfare.tex).
Outputs: results/theory/welfare_sim.csv, figures/welfare_failure_vs_s.pdf, figures/welfare_systemic_vs_rho.pdf.
Only LTV = 80 % and the haircut rule are transparent calibrations; all other parameters are ILLUSTRATIVE."""
from __future__ import annotations

import csv
import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from agentic_valuation.theory import welfare as wf  # noqa: E402

SEED = 20261027
OUT = ROOT / "results" / "theory"
FIG = ROOT / "figures"
OUT.mkdir(parents=True, exist_ok=True)
FIG.mkdir(parents=True, exist_ok=True)
C1, C2, C3, C4 = "#2a78d6", "#eb6834", "#1baf7a", "#eda100"
INK, MUTED, GRID = "#1a1a19", "#6b6a66", "#e6e5e1"


def style(ax):
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(GRID)
    ax.tick_params(colors=MUTED, labelsize=8)
    ax.yaxis.grid(True, color=GRID, lw=0.6)
    ax.set_axisbelow(True)


def main() -> None:
    t0 = time.time()
    rng = np.random.default_rng(SEED)
    rows: list[dict] = []
    n = 200_000
    shocks = wf.draw_market_shocks(n, rng)
    s_grid = np.linspace(0.0, 0.25, 26)
    base = wf.WelfareParams()
    variants = {"h1=0,b=0": wf.with_params(base, h1=0.0), "h1=0.5,b=0": wf.with_params(base, h1=0.5)}
    biases = {"b=-0.03": -0.03, "b=0": 0.0, "b=+0.03": 0.03}
    market = {}
    for vname, p in variants.items():
        for bname, b in biases.items():
            for s in s_grid:
                m = wf.simulate_market(b, float(s), p, n, rng, shocks=shocks)
                el = wf.expected_loss(b, float(s), p)
                rows.append({"experiment": "transaction", "variant": vname, "bias": b, "s": round(float(s), 4),
                             "rho": "", **m, "expected_loss": el})
                market[(vname, bname, round(float(s), 4))] = (m, el)

    # Systemic block: N lenders, common-model correlation rho
    N, n_sys = 50, 20_000
    sys_shocks = {"C": rng.standard_normal(n_sys), "E": rng.standard_normal((n_sys, N))}
    rho_grid = np.linspace(0.0, 1.0, 11)
    systemic = {}
    for s in (0.05, 0.10, 0.15):
        for rho in rho_grid:
            r = wf.systemic_loss(float(rho), 0.0, s, base, N, n_sys, rng, shocks=sys_shocks)
            rows.append({"experiment": "systemic", "variant": f"s={s}", "bias": 0.0, "s": s, "rho": round(float(rho), 3),
                         "sys_mean": r["mean"], "sys_sd": r["sd"], "sys_es95": r["es95"], "cov_pair": r["cov_pair"]})
            systemic[(s, round(float(rho), 3))] = r
    # "Adopting the best common AVM": accuracy improves (s falls) but rho rises to 1
    for (s, rho, label) in ((0.15, 0.2, "many_avms"), (0.10, 1.0, "single_best_avm")):
        r = wf.systemic_loss(rho, 0.0, s, base, N, n_sys, rng, shocks=sys_shocks)
        rows.append({"experiment": "concentration", "variant": label, "bias": 0.0, "s": s, "rho": rho,
                     "sys_mean": r["mean"], "sys_sd": r["sd"], "sys_es95": r["es95"], "cov_pair": r["cov_pair"]})

    fields = ["experiment", "variant", "bias", "s", "rho", "beneficial_rate", "failure_rate", "failure_rate_exact",
              "completion_rate", "mean_price", "seller_surplus", "buyer_surplus", "total_surplus", "lost_surplus",
              "expected_loss", "sys_mean", "sys_sd", "sys_es95", "cov_pair"]
    with open(OUT / "welfare_sim.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields)
        w.writeheader()
        for row in rows:
            w.writerow({k: row.get(k, "") for k in fields})

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update({"font.family": "serif", "font.size": 9, "axes.labelcolor": INK, "text.color": INK})

    # Figure 1: failure rate of beneficial transactions vs s (two haircut rules x three biases)
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.2), sharey=True)
    bcols = {"b=-0.03": C2, "b=0": C1, "b=+0.03": C3}
    for ax, (vname, _) in zip(axes, variants.items()):
        for bname, b in biases.items():
            ys = [market[(vname, bname, round(float(s), 4))][0]["failure_rate_exact"] for s in s_grid]
            ax.plot(s_grid, ys, color=bcols[bname], lw=2, label=bname.replace("b=", "bias $b=$"))
        ax.set_title("no uncertainty haircut ($h_1=0$)" if vname.startswith("h1=0,") else
                     "haircut rising with $s$ ($h_1=0.5$)", fontsize=9)
        ax.set_xlabel("valuation error standard deviation $s$")
        style(ax)
    axes[0].set_ylabel("failure rate, beneficial matches")
    axes[0].legend(frameon=False, fontsize=8, loc="upper left")
    fig.tight_layout()
    fig.savefig(FIG / "welfare_failure_vs_s.pdf")
    plt.close(fig)

    # Figure 2: systemic loss vs rho
    fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.2))
    scols = {0.05: "#9ec5f4", 0.10: "#256abf", 0.15: "#104281"}
    for s in (0.05, 0.10, 0.15):
        axes[0].plot(rho_grid, [systemic[(s, round(float(r), 3))]["es95"] for r in rho_grid], color=scols[s], lw=2,
                     label=f"ES$_{{95}}$, $s={s:.2f}$")
        axes[0].plot(rho_grid, [systemic[(s, round(float(r), 3))]["mean"] for r in rho_grid], color=scols[s], lw=1,
                     ls="--")
        axes[1].plot(rho_grid, [systemic[(s, round(float(r), 3))]["sd"] for r in rho_grid], color=scols[s], lw=2,
                     label=f"$s={s:.2f}$")
    axes[0].set_title("system loss: mean (dashed) and 95 % expected shortfall", fontsize=9)
    axes[1].set_title("standard deviation of the system loss", fontsize=9)
    for ax in axes:
        ax.set_xlabel(r"common-model correlation $\rho$ across lenders' AVMs")
        ax.set_ylim(0, None)
        style(ax)
    axes[0].set_ylabel(f"loss per $N={N}$ loans (share of $V$)")
    axes[0].legend(frameon=False, fontsize=7, loc="upper left")
    axes[1].legend(frameon=False, fontsize=8, loc="upper left")
    fig.tight_layout()
    fig.savefig(FIG / "welfare_systemic_vs_rho.pdf")
    plt.close(fig)
    print(f"wrote {OUT / 'welfare_sim.csv'} ({len(rows)} rows) and 2 figures in {time.time() - t0:.1f} s")


if __name__ == "__main__":
    main()