"""WP18 — experiment 10: robustness summary (discounting, preferences, transaction costs, grid
resolution, solver tolerance, simulation sample size, seed).

Reads results/data/<group>/<name>.json written by simulation.py and writes
results/data/robustness_summary.csv (key moments of every scenario next to the baseline, with
the relative deviation), results/data/grid_sensitivity.csv (grid / tolerance / sample-size
variants only) and results/data/diagnostics_table.csv (convergence of every run).
Usage: python -m src.robustness
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd

from .model import RESULTS_DIR

DATA = RESULTS_DIR / "data"

KEY_MOMENTS = [
    ("p", lambda d: d["p"]),
    ("rho", lambda d: d["rho"]),
    ("ownership", lambda d: d["moments"]["ownership_rate_all"]),
    ("ownership_under35", lambda d: d["moments"]["ownership_rate_under35"]),
    ("threshold_age1_zmid", lambda d: d["thresholds"]["1"][str(len(d["thresholds"]["1"]) // 2)]),
    ("gini_nw", lambda d: d["gini_nw"]),
    ("mean_nw", lambda d: d["mean_nw"]),
    ("mean_gift", lambda d: d["moments"]["mean_gift_received"]),
    ("lambda_age7", lambda d: d["mobility"]["lambda_persistence"]),
    ("q_RO_age3", lambda d: d["mobility"]["q_RO_age3"]),
    ("q_OO_age3", lambda d: d["mobility"]["q_OO_age3"]),
    ("rank_rank_nw", lambda d: d["mobility"]["rank_rank_slope_nw"]),
    ("nw_gap_age4", lambda d: d["gaps"]["nw_gap_owner_renter_age4"]),
    ("share_mortgage", lambda d: d["moments"]["share_with_mortgage"]),
]

GRID_SCENARIOS = ("grid_fine", "grid_coarse", "tol_tight", "sample_small", "sample_large", "seed_alt")


def load_all() -> dict[str, dict]:
    out = {}
    for f in sorted(DATA.glob("*/*.json")):
        if f.parent.name in ("welfare", "multiplicity"):
            continue
        with open(f) as fh:
            d = json.load(fh)
        if isinstance(d, dict) and "moments" in d:
            out[d["name"]] = d
    return out


def _safe(fn, d):
    try:
        v = fn(d)
        return float(v) if v is not None else np.nan
    except (KeyError, TypeError, IndexError):
        return np.nan


def build_summary(runs: dict[str, dict]) -> pd.DataFrame:
    base = runs.get("baseline")
    rows = []
    for name, d in runs.items():
        row = dict(name=name, group=d.get("group"), equilibrium=d.get("equilibrium"), converged=d.get("converged"),
                   time_sec=d.get("time_sec"), excess_owner=d.get("excess_owner"), excess_rental=d.get("excess_rental"),
                   bellman_owner_mean=(d.get("bellman") or {}).get("owner_mean"),
                   bellman_renter_mean=(d.get("bellman") or {}).get("renter_mean"))
        for k, fn in KEY_MOMENTS:
            v = _safe(fn, d)
            row[k] = v
            if base is not None:
                b = _safe(fn, base)
                row[f"{k}_dev"] = (v - b) if np.isfinite(v) and np.isfinite(b) else np.nan
        rows.append(row)
    return pd.DataFrame(rows)


def main():
    runs = load_all()
    if not runs:
        raise SystemExit("no scenario results found under results/data — run `make simulate` first")
    df = build_summary(runs)
    df.to_csv(DATA / "robustness_summary.csv", index=False)
    grid = df[df.name.isin(GRID_SCENARIOS + ("baseline",))].copy()
    grid.to_csv(DATA / "grid_sensitivity.csv", index=False)
    diag = df[["name", "group", "equilibrium", "converged", "time_sec", "excess_owner", "excess_rental",
               "bellman_owner_mean", "bellman_renter_mean"]]
    diag.to_csv(DATA / "diagnostics_table.csv", index=False)
    # max absolute deviation across the numerical variants, as a share of the baseline value
    base = df[df.name == "baseline"].iloc[0]
    dev = {}
    for k, _ in KEY_MOMENTS:
        vals = grid[k].values
        b = base[k]
        dev[k] = float(np.nanmax(np.abs(vals - b)) / max(abs(b), 1e-9)) if np.isfinite(b) else np.nan
    with open(DATA / "grid_sensitivity_max_relative_deviation.json", "w") as fh:
        json.dump(dev, fh, indent=1)
    print(df[["name", "equilibrium", "converged", "p", "ownership", "gini_nw", "lambda_age7"]].to_string())
    print("max relative deviation across numerical variants:", dev)


if __name__ == "__main__":
    main()