# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Figures (PDF, figures/) from results/evaluation/. No titles inside the figures (captions live in LaTeX);
one idea per figure; colour-blind-safe palette validated with the dataviz checker (6 slots + markers/hatches).

    fig_error_distribution.pdf   log error e by arm (box + whiskers at P5/P95)
    fig_pit.pdf                  PIT histograms per arm (small multiples, uniform reference)
    fig_coverage_vs_nominal.pdf  empirical vs nominal interval coverage per arm
    fig_rounds.pdf               rounds curve (only if results/evaluation/rounds.csv exists)
    fig_cost_accuracy.pdf        cost per property vs mean |e| with the Pareto frontier
    fig_strata.pdf               mean |e| by stratum level × arm (grouped bars, suppressed cells omitted)

    python3 scripts/21_figures.py [--stratum region_group] [--arms D_avm,E_single_llm,F_independent,G_debate_r2,J_cost_matched]
"""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

from agentic_valuation.config import FIGURES, RESULTS  # noqa: E402
from agentic_valuation.evaluation.metrics import PRIMARY_METRIC  # noqa: E402
from agentic_valuation.evaluation.stratify import STRATUM_LABELS  # noqa: E402

PALETTE = ["#1f5fa8", "#c8541c", "#0b7a5a", "#8b4e8d", "#b58b00", "#2a9dcf"]   # validated: dataviz validate_palette.js
MARKERS = ["o", "s", "^", "D", "v", "P"]
HATCHES = ["", "//", "..", "xx", "\\\\", "++"]
GREY = "#7a7a7a"
INK = "#222222"
DEFAULT_ARMS = ["D_avm", "E_single_llm", "F_independent", "G_debate_r2", "I_placebo", "J_cost_matched"]

plt.rcParams.update({
    "font.size": 8.5, "axes.labelsize": 8.5, "xtick.labelsize": 8, "ytick.labelsize": 8, "legend.fontsize": 7.5,
    "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": GREY, "axes.labelcolor": INK,
    "xtick.color": GREY, "ytick.color": GREY, "axes.grid": True, "grid.color": "#e3e3e3", "grid.linewidth": 0.5,
    "axes.axisbelow": True, "pdf.fonttype": 42, "ps.fonttype": 42, "figure.dpi": 150, "legend.frameon": False,
    "text.color": INK, "font.family": "sans-serif",
})


def pretty(arm: str) -> str:
    parts = arm.split("_")
    return f"{parts[0]} ({' '.join(parts[1:])})" if len(parts) > 1 else arm


def arm_style(arms: list[str]) -> dict[str, tuple[str, str, str]]:
    """Fixed, order-stable assignment arm → (colour, marker, hatch); beyond 6 arms → grey, folded."""
    style = {}
    for i, a in enumerate(sorted(arms)):
        style[a] = (PALETTE[i], MARKERS[i], HATCHES[i]) if i < len(PALETTE) else (GREY, "x", "")
    return style


def save(fig: plt.Figure, path: Path) -> None:
    fig.savefig(path, bbox_inches="tight", pad_inches=0.02)
    plt.close(fig)
    print("wrote", path)


def fig_error_distribution(records: pd.DataFrame, out: Path) -> None:
    arms = sorted(records["arm"].unique())
    data = [records.loc[records["arm"] == a, "e"].dropna().to_numpy() for a in arms]
    fig, ax = plt.subplots(figsize=(3.4, 0.32 * len(arms) + 0.8))
    bp = ax.boxplot(data, vert=False, whis=(5, 95), showfliers=False, widths=0.55, patch_artist=True,
                    medianprops={"color": INK, "linewidth": 1.2}, boxprops={"facecolor": "#d6e2f1", "edgecolor": PALETTE[0], "linewidth": 0.9},
                    whiskerprops={"color": PALETTE[0], "linewidth": 0.9}, capprops={"color": PALETTE[0], "linewidth": 0.9})
    ax.axvline(0, color=GREY, linewidth=0.8, linestyle="--")
    ax.set_yticks(range(1, len(arms) + 1), [pretty(a) for a in arms])
    ax.set_xlabel(r"log error $e=\log\hat P-\log P$ (box IQR, whiskers P5–P95)")
    ax.invert_yaxis()
    save(fig, out / "fig_error_distribution.pdf")


def fig_pit(pit: pd.DataFrame, out: Path) -> None:
    arms = sorted(pit["arm"].unique())
    if not arms:
        return
    ncol = min(3, len(arms))
    nrow = int(np.ceil(len(arms) / ncol))
    fig, axes = plt.subplots(nrow, ncol, figsize=(2.1 * ncol, 1.6 * nrow + 0.3), sharex=True, sharey=True, squeeze=False)
    for ax, a in zip(axes.flat, arms):
        g = pit[pit["arm"] == a]
        ax.bar((g["bin_lo"] + g["bin_hi"]) / 2, g["share"], width=(g["bin_hi"] - g["bin_lo"]) * 0.92, color=PALETTE[0], edgecolor="white", linewidth=0.5)
        ax.axhline(g["expected"].iloc[0], color=PALETTE[1], linewidth=1.0, linestyle="--")
        ax.text(0.0, 1.02, pretty(a), transform=ax.transAxes, ha="left", va="bottom", fontsize=7.5)
        ax.set_xlim(0, 1)
        ax.set_xticks([0, 0.5, 1])
    for ax in axes.flat[len(arms):]:
        ax.axis("off")
    for j in range(ncol):
        col_axes = [axes[i, j] for i in range(nrow) if i * ncol + j < len(arms)]
        if col_axes:
            col_axes[-1].set_xlabel("PIT")
            col_axes[-1].tick_params(labelbottom=True)
    for ax in axes[:, 0]:
        ax.set_ylabel("share")
    save(fig, out / "fig_pit.pdf")


def fig_coverage(coverage: pd.DataFrame, out: Path) -> None:
    c = coverage[(coverage["group"] == "all") & (coverage["kind"] == "interval")]
    arms = sorted(c["arm"].unique())
    if not arms:
        return
    style = arm_style(arms)
    fig, ax = plt.subplots(figsize=(3.3, 3.0))
    ax.plot([0.4, 1.0], [0.4, 1.0], color=GREY, linewidth=0.8, linestyle="--", zorder=1)
    for a in arms:
        g = c[c["arm"] == a].sort_values("nominal")
        col, mk, _ = style[a]
        ax.plot(g["nominal"], g["empirical"], color=col, marker=mk, markersize=4.5, linewidth=1.4, label=pretty(a), markeredgecolor="white", markeredgewidth=0.6)
    ax.set_xlabel("nominal coverage")
    ax.set_ylabel("empirical coverage")
    ax.set_xticks([0.5, 0.8, 0.9, 0.95], ["0.50", "0.80", "0.90", "0.95"], rotation=45, ha="right")
    ax.set_xlim(0.45, 1.0)
    ax.set_ylim(min(0.3, float(c["empirical"].min()) - 0.05), 1.0)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.22), ncol=2, handlelength=1.6, columnspacing=1.0)
    save(fig, out / "fig_coverage_vs_nominal.pdf")


def fig_rounds(rounds: pd.DataFrame, out: Path) -> None:
    """Expects columns: round, value (+ optional ci_lo, ci_hi, arm/series, metric)."""
    if "arm" not in rounds.columns:
        rounds = rounds.assign(arm="G_debate")
    if "metric" in rounds.columns:
        rounds = rounds[rounds["metric"] == PRIMARY_METRIC] if (rounds["metric"] == PRIMARY_METRIC).any() else rounds
    arms = sorted(rounds["arm"].unique())
    style = arm_style(arms)
    fig, ax = plt.subplots(figsize=(3.3, 2.5))
    for a in arms:
        g = rounds[rounds["arm"] == a].sort_values("round")
        col, mk, _ = style[a]
        if {"ci_lo", "ci_hi"} <= set(g.columns):
            ax.fill_between(g["round"], g["ci_lo"], g["ci_hi"], color=col, alpha=0.15, linewidth=0)
        ax.plot(g["round"], g["value"], color=col, marker=mk, markersize=4.5, linewidth=1.4, label=pretty(a), markeredgecolor="white", markeredgewidth=0.6)
    ax.set_xlabel("debate rounds")
    ax.set_ylabel(r"mean $|e|$")
    ax.set_xticks(sorted(rounds["round"].unique()))
    if len(arms) > 1:
        ax.legend()
    save(fig, out / "fig_rounds.pdf")


def fig_cost_accuracy(summary: pd.DataFrame, out: Path) -> None:
    s = summary.dropna(subset=[PRIMARY_METRIC]).copy()
    s["cost"] = s["cost_usd_mean"].fillna(0.0)
    fig, ax = plt.subplots(figsize=(3.4, 2.8))
    floor = max(s.loc[s["cost"] > 0, "cost"].min() / 10 if (s["cost"] > 0).any() else 1e-3, 1e-4)
    x = np.where(s["cost"] > 0, s["cost"], floor)
    fam = s["arm"].str[0].le("D")
    ax.scatter(x[fam], s.loc[fam, PRIMARY_METRIC], color=PALETTE[0], marker="s", s=26, label="statistical baselines (A–D, cost ≈ 0)", zorder=3, edgecolor="white", linewidth=0.6)
    ax.scatter(x[~fam], s.loc[~fam, PRIMARY_METRIC], color=PALETTE[1], marker="o", s=26, label="LLM arms (E–J)", zorder=3, edgecolor="white", linewidth=0.6)
    # Pareto frontier: lowest error for each cost level (non-dominated points), sorted by cost
    pts = sorted(zip(x, s[PRIMARY_METRIC], s["arm"]), key=lambda t: (t[0], t[1]))
    frontier, best = [], np.inf
    for cx, cy, a in pts:
        if cy < best:
            frontier.append((cx, cy))
            best = cy
    if len(frontier) > 1:
        ax.plot([f[0] for f in frontier], [f[1] for f in frontier], color=GREY, linewidth=0.9, linestyle="-", zorder=2)
    for cx, cy, a in zip(x, s[PRIMARY_METRIC], s["arm"]):
        ax.annotate(a.split("_")[0], (cx, cy), textcoords="offset points", xytext=(4, 3), fontsize=7, color=INK)
    ax.set_xscale("log")
    ax.set_xlabel("cost per property (USD, log; baselines at floor)")
    ax.set_ylabel(r"mean $|e|$")
    ax.legend(loc="upper right")
    save(fig, out / "fig_cost_accuracy.pdf")


def fig_strata(strata: pd.DataFrame, out: Path, stratum: str, arms: list[str]) -> None:
    s = strata[(strata["metric"] == PRIMARY_METRIC) & (strata["stratum"] == stratum) & (strata["arm"].isin(arms))]
    if s.empty:
        print(f"no strata rows for {stratum}; skipping fig_strata")
        return
    levels = list(dict.fromkeys(s["level"]))
    arms = [a for a in sorted(arms) if a in set(s["arm"])]
    style = arm_style(arms)
    w = 0.8 / max(len(arms), 1)
    fig, ax = plt.subplots(figsize=(max(3.3, 0.9 * len(levels) + 1.2), 2.6))
    for j, a in enumerate(arms):
        g = s[s["arm"] == a].set_index("level").reindex(levels)
        vals = g["value"].to_numpy(dtype=float)
        col, _, hatch = style[a]
        xs = np.arange(len(levels)) + (j - (len(arms) - 1) / 2) * w
        ax.bar(xs, np.nan_to_num(vals, nan=0.0), width=w * 0.92, color=col, hatch=hatch, edgecolor="white", linewidth=0.5, label=pretty(a))
        for xi, v, n in zip(xs, vals, g["n"]):
            if not np.isfinite(v):
                ax.text(xi, 0.002, "n<30", rotation=90, ha="center", va="bottom", fontsize=6, color=GREY)
    ax.set_xticks(range(len(levels)), [f"{lv}\n(n={int(s[s['level'] == lv]['n'].max())})" for lv in levels])
    ax.set_ylabel(r"mean $|e|$")
    ax.set_xlabel(STRATUM_LABELS.get(stratum, stratum))
    ax.legend(ncol=min(3, len(arms)), loc="upper center", bbox_to_anchor=(0.5, -0.32))
    save(fig, out / "fig_strata.pdf")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--eval-dir", default=str(RESULTS / "evaluation"))
    ap.add_argument("--out", default=str(FIGURES))
    ap.add_argument("--stratum", default="region_group")
    ap.add_argument("--arms", default=",".join(DEFAULT_ARMS), help="arms shown in the stratified figure (≤ 6)")
    args = ap.parse_args(argv)
    ev, out = Path(args.eval_dir), Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    if not (ev / "records.parquet").exists():
        print(f"{ev}/records.parquet not found — run scripts/20_evaluate.py first")
        return 1
    records = pd.read_parquet(ev / "records.parquet")
    summary = pd.read_csv(ev / "summary_main.csv")
    fig_error_distribution(records, out)
    if (ev / "pit.csv").exists() and (ev / "pit.csv").stat().st_size > 1:
        pit = pd.read_csv(ev / "pit.csv")
        if not pit.empty:
            fig_pit(pit, out)
    if (ev / "coverage.csv").exists():
        cov = pd.read_csv(ev / "coverage.csv")
        if not cov.empty:
            fig_coverage(cov, out)
    if (ev / "rounds.csv").exists():
        fig_rounds(pd.read_csv(ev / "rounds.csv"), out)
    else:
        print("results/evaluation/rounds.csv absent — fig_rounds.pdf not produced (placeholder)")
    fig_cost_accuracy(summary, out)
    if (ev / "strata.csv").exists():
        strata = pd.read_csv(ev / "strata.csv")
        if not strata.empty:
            fig_strata(strata, out, args.stratum, [a for a in args.arms.split(",") if a][:6])
    return 0


if __name__ == "__main__":
    raise SystemExit(main())