"""Experiment program for WP19 (Phase 4). Each experiment has a pre-registered hypothesis (stored with the outputs),
an exact configuration, CSV outputs in results/processed/, and a JSON record in results/raw/.

Run:  python -m src.simulation --exp all      (or --exp 1,3,9)
"""
from __future__ import annotations

import argparse
import json
import logging
import platform
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd
import scipy

from .model import load_params, Exponential, Pareto, Weibull, TruncNormal, Params
from .solver import (solve_equilibrium, partial_equilibrium, free_entry_residual, solve_equilibrium_posted,
                     solve_equilibrium_financing, transition, frictionless_price, welfare_flow, planner_optimum)

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "results" / "raw"
PRO = ROOT / "results" / "processed"
LOG = ROOT / "results" / "logs"
for d in (RAW, PRO, LOG):
    d.mkdir(parents=True, exist_ok=True)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s",
                    handlers=[logging.FileHandler(LOG / "simulation.log"), logging.StreamHandler(sys.stdout)])
log = logging.getLogger("wp19")

HYPOTHESES = {
    1: "Baseline: unique stationary equilibrium; TOM of a few months; urgency discount of a few percent; price dispersion among identical houses with a positive between-type share; population hazard decreasing in duration (Lemma 3).",
    2: "Urgency gradient: V, pbar decrease and eta increases in c (Prop. 1); |pbar'| < |V'|; the gradient steepens (gap widens) when the arrival rate falls (Prop. 2(iii)).",
    3: "Arrival-rate collapse: lower A lowers prices and tightness, raises TOM and inventory, widens the urgency discount; GE response exceeds PE response (amplification); Gamma changes sign between thin and thick markets (Prop. 2(iv)).",
    4: "Financing: in the baseline screen, lower phi acts like lower A. In the contingency extension, mortgage deals need better matches; the cash-mortgage gap = delay-bargaining + outside-option + selection; the bargaining term rises with urgency (Prop. 7); tighter approval (lower phi_m, lower nu) lengthens TOM and lowers prices.",
    5: "Bargaining power: higher seller share beta raises prices and lowers tightness and hazards; total welfare is maximised near the Hosios value beta = 1-alpha with homogeneous sellers; with heterogeneous sellers the welfare-maximising beta deviates (numerical).",
    6: "Matching efficiency vs qualification: observationally equivalent in the baseline; in the extension a matching-efficiency improvement and an approval improvement calibrated to the same sales-hazard change have different price and TOM effects because approval shocks also change the pending wedge.",
    7: "Endogenous entry: the free-entry residual is monotone in theta (unique equilibrium); GE amplifies A and phi shocks relative to PE; lower entry/search costs raise theta and sellers' prices.",
    8: "Heterogeneous sellers: between-type dispersion rises with the spread of F; the transaction-price distribution is a mixture whose left tail is populated by urgent sellers; dispersion increases when liquidity falls.",
    9: "Dynamics: after a permanent negative A shock the observed index falls on impact, inventory builds up, volume drops then recovers; the composition term is transitory when delta=0; its sign depends on G (dampening for the truncated normal, amplifying for the exponential).",
    10: "Selection bias: impact decomposition of the index change into within-type and composition terms across G families and urgency spreads; the composition term is small for light tails, larger and amplifying for exponential/Pareto; steady-state composition effects vanish when delta=0.",
    11: "Posted prices: sellers hold out for better matches (lower hazard, higher price conditional on sale), all dispersion is between-type; the urgency gradient of prices is preserved.",
    12: "Sensitivity and convergence: qualitative results (signs) are robust to +-25% parameter changes and to the G family; equilibrium objects converge in the number of seller types and in tolerances; Monte Carlo simulation of the listing process matches analytic hazards and prices.",
}


def env_record(p: Params) -> dict:
    return {"python": platform.python_version(), "numpy": np.__version__, "scipy": scipy.__version__,
            "pandas": pd.__version__, "config_hash": p.config_hash(), "seed": p.seed,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S")}


def save(exp: int, name: str, df: pd.DataFrame, p: Params, meta: dict | None = None):
    path = PRO / f"exp{exp:02d}_{name}.csv"
    df.to_csv(path, index=False)
    rec = {"experiment": exp, "name": name, "hypothesis": HYPOTHESES[exp], "rows": len(df), "columns": list(df.columns),
           "env": env_record(p), **(meta or {})}
    (RAW / f"exp{exp:02d}_{name}.json").write_text(json.dumps(rec, indent=2, default=str))
    log.info("saved %s (%d rows)", path.name, len(df))


def type_table(eq, p: Params) -> pd.DataFrame:
    sb = eq.sellers
    return pd.DataFrame({
        "c": p.c, "f": p.f, "V": sb.V, "vstar": sb.vstar, "acc": sb.acc, "eta": sb.eta, "pbar": sb.pbar,
        "D": 1 - sb.pbar / sb.pbar[0], "E_tau": 1 / (sb.eta + p.delta), "E_tau_months": 12 / (sb.eta + p.delta),
        "sold_prob": sb.eta / (sb.eta + p.delta), "L_holding_cost": p.c / (p.r + p.delta + sb.eta),
        "stock_share": eq.N, "sales_share": eq.omega, "M": sb.M,
        "price_sd_within": p.beta * np.sqrt([p.G.cond_var(v) for v in sb.vstar]),
    })


# ------------------------------------------------------------------------------------------------
def exp01(p: Params):
    eq = solve_equilibrium(p)
    m = eq.moments(p)
    save(1, "baseline_moments", pd.DataFrame([m]), p, {"theta_roots": eq.theta_roots})
    save(1, "baseline_types", type_table(eq, p), p)
    tau = np.linspace(0, 2.0, 201)
    surv = np.exp(-np.outer(tau, eq.sellers.eta + p.delta))
    pop = surv @ p.f
    haz = -np.gradient(np.log(pop), tau)
    df = pd.DataFrame({"tau": tau, "tau_months": 12 * tau, "survival_population": pop, "hazard_population": haz})
    for i in (0, len(p.c) // 2, len(p.c) - 1):
        df[f"survival_c{i}"] = surv[:, i]
    save(1, "baseline_survival", df, p)
    # free-entry residual curve (uniqueness diagnostic)
    th = np.geomspace(p.theta_bounds[0], p.theta_bounds[1], 80)
    save(1, "free_entry_curve", pd.DataFrame({"theta": th, "residual": [free_entry_residual(t, p) for t in th]}), p)
    return eq


def exp02(p: Params):
    eq = solve_equilibrium(p)
    cgrid = np.linspace(0.0, 0.5, 101)
    rows = []
    for mult in (0.5, 0.75, 1.0, 1.5, 2.0):
        pm = p.copy(c=cgrid, f=np.full_like(cgrid, 1 / len(cgrid)))
        pe = partial_equilibrium(eq.theta, pm.copy(A=p.A * mult))     # partial equilibrium: theta fixed at baseline
        sb = pe.sellers
        for i, c in enumerate(cgrid):
            rows.append({"A_mult": mult, "Lam_s": pe.Lam_s, "c": c, "V": sb.V[i], "vstar": sb.vstar[i], "eta": sb.eta[i],
                         "pbar": sb.pbar[i], "D": 1 - sb.pbar[i] / sb.pbar[0], "E_tau_months": 12 / (sb.eta[i] + p.delta),
                         "dV_dc_analytic": -1 / ((p.r + p.delta) + pe.Lam_s * p.beta * sb.acc[i]),
                         "participation": sb.V[i] >= p.w0})
    df = pd.DataFrame(rows)
    df["dV_dc_fd"] = df.groupby("A_mult")["V"].transform(lambda s: np.gradient(s.values, cgrid))
    df["dpbar_dc_fd"] = df.groupby("A_mult")["pbar"].transform(lambda s: np.gradient(s.values, cgrid))
    save(2, "urgency_gradient", df, p, {"theta_fixed": eq.theta})
    # GE version: urgency discount of each baseline type as a function of A
    rows = []
    for mult in np.linspace(0.4, 1.6, 13):
        e = solve_equilibrium(p.copy(A=p.A * mult))
        for i, c in enumerate(p.c):
            rows.append({"A_mult": mult, "theta": e.theta, "c": c, "pbar": e.sellers.pbar[i], "V": e.sellers.V[i],
                         "D": 1 - e.sellers.pbar[i] / e.sellers.pbar[0], "gapV": e.sellers.V[0] - e.sellers.V[i],
                         "eta": e.sellers.eta[i]})
    save(2, "urgency_discount_vs_A_GE", pd.DataFrame(rows), p)


def exp03(p: Params):
    base = solve_equilibrium(p)
    rows = []
    for mult in np.concatenate([np.linspace(0.2, 1.0, 17), np.linspace(1.1, 3.0, 20)]):
        pm = p.copy(A=p.A * mult)
        ge = solve_equilibrium(pm); pe = partial_equilibrium(base.theta, pm)
        for kind, e in (("GE", ge), ("PE", pe)):
            m = e.moments(p)
            m.update({"A_mult": mult, "kind": kind, "Q_impact": float((e.sellers.eta * base.n).sum())})
            rows.append(m)
    save(3, "arrival_shock", pd.DataFrame(rows), p, {"theta_base": base.theta})
    # Gamma sign change for the patient type in partial equilibrium over a wide Lam_s range
    rows = []
    for L in np.geomspace(0.05, 2000, 60):
        pm = p.copy(c=np.array([0.0]), f=np.array([1.0]), delta=0.0)
        from .solver import solve_sellers
        sb = solve_sellers(L, p.k, pm)
        rows.append({"Lam_s": L, "pbar_patient": sb.pbar[0], "p_star": frictionless_price(p),
                     "Gamma": 1 - sb.pbar[0] / frictionless_price(p), "vstar": sb.vstar[0]})
    save(3, "gamma_sign_change", pd.DataFrame(rows), p)


def exp04(p: Params):
    # (a) baseline screen
    rows = []
    for phi in np.linspace(0.4, 1.0, 13):
        e = solve_equilibrium(p.copy(phi=phi)); m = e.moments(p); m["phi"] = phi; rows.append(m)
    save(4, "qualification_screen", pd.DataFrame(rows), p)
    # (b) contingency extension
    def fin_rows(fe, label, val):
        tw_c = fe.Q_c / fe.Q_c.sum(); tw_m = fe.Q_m / fe.Q_m.sum()
        out = []
        for i, c in enumerate(p.c):
            out.append({label: val, "c": c, "f": p.f[i], "theta": fe.theta, "U_c": fe.U_c, "U_m": fe.U_m, "pi_tilde": fe.pi_tilde,
                        "V": fe.V[i], "vstar_c": fe.vstar_c[i], "vstar_m": fe.vstar_m[i], "eta_c": fe.eta_c[i], "eta_m": fe.eta_m[i],
                        "pbar_c": fe.pbar_c[i], "pbar_m": fe.pbar_m[i], "gap": fe.pbar_m[i] - fe.pbar_c[i],
                        "Delta_b": fe.Delta_b[i], "outside_term": fe.outside_term, "sel_term": fe.sel_term[i],
                        "share_cash_sales": fe.Q_c[i] / (fe.Q_c[i] + fe.Q_m[i]),
                        "E_tau_listed_months": 12 / (fe.eta_c[i] + p.phi_m * fe.eta_m[i] + p.delta),
                        "S": fe.S, "Q": fe.Q, "pending_stock": fe.m_pending.sum(), "converged": fe.converged,
                        "E_total_duration_months": 12 * (fe.S + fe.m_pending.sum()) / p.s,
                        "buyer_participation_ok": bool(min(fe.U_c, fe.U_m) >= 0), "seller_participation_ok": bool(np.all(fe.V >= p.w0)),
                        "price_mean_cash": float((tw_c * fe.pbar_c).sum()), "price_mean_mortgage": float((tw_m * fe.pbar_m).sum()),
                        "cash_share_of_sales": fe.Q_c.sum() / fe.Q})
        return out
    rows = []
    fe0 = solve_equilibrium_financing(p); rows += fin_rows(fe0, "phi_m", p.phi_m)
    for phi in (0.5, 0.6, 0.7, 0.75, 0.8, 0.9, 0.95, 1.0):
        rows += fin_rows(solve_equilibrium_financing(p.copy(phi_m=phi)), "phi_m", phi)
    save(4, "financing_phi", pd.DataFrame(rows), p)
    rows = []
    for nu in (3.0, 6.0, 12.0, 24.0, 52.0, 1e4):
        rows += fin_rows(solve_equilibrium_financing(p.copy(nu=nu)), "nu", nu)
    save(4, "financing_nu", pd.DataFrame(rows), p)
    rows = []
    for pi in (0.1, 0.3, 0.5, 0.7, 0.9):
        rows += fin_rows(solve_equilibrium_financing(p.copy(pi=pi)), "pi", pi)
    save(4, "financing_pi", pd.DataFrame(rows), p)
    rows = []
    for beta in (0.2, 0.35, 0.5, 0.65, 0.8):
        rows += fin_rows(solve_equilibrium_financing(p.copy(beta=beta)), "beta", beta)
    save(4, "financing_beta", pd.DataFrame(rows), p)


def exp05(p: Params):
    rows = []
    hom = p.copy(c=np.array([float((p.c * p.f).sum())]), f=np.array([1.0]))
    for beta in np.linspace(0.1, 0.9, 17):
        for label, pp in (("heterogeneous", p), ("homogeneous", hom)):
            e = solve_equilibrium(pp.copy(beta=beta)); m = e.moments(pp)
            m.update({"beta": beta, "sellers": label, "hosios_beta": 1 - p.alpha})
            rows.append(m)
    df = pd.DataFrame(rows)
    save(5, "bargaining_power", df, p)
    # planner benchmark for both economies
    rows = []
    for label, pp in (("heterogeneous", p), ("homogeneous", hom)):
        e = solve_equilibrium(pp); pl = planner_optimum(pp)
        rows.append({"sellers": label, "theta_decentralised": e.theta, "theta_planner": pl["theta"],
                     "welfare_decentralised": welfare_flow(e, pp)["total"], "welfare_planner": pl["welfare"],
                     "vstar_decentralised": json.dumps(np.round(e.sellers.vstar, 5).tolist()),
                     "vstar_planner": json.dumps(np.round(pl["vstar"], 5).tolist()), "planner_success": pl["success"]})
        for beta in np.linspace(0.1, 0.9, 33):
            pass
    save(5, "planner_benchmark", pd.DataFrame(rows), p)
    # welfare surface over (beta, alpha)
    rows = []
    for alpha in (0.3, 0.4, 0.5, 0.6, 0.7):
        for beta in np.linspace(0.1, 0.9, 17):
            for label, pp in (("heterogeneous", p), ("homogeneous", hom)):
                e = solve_equilibrium(pp.copy(alpha=alpha, beta=beta))
                rows.append({"alpha": alpha, "beta": beta, "sellers": label, "welfare_total": welfare_flow(e, pp)["total"],
                             "theta": e.theta, "converged": e.converged})
    save(5, "welfare_surface", pd.DataFrame(rows), p)


def exp06(p: Params):
    """Matching efficiency vs approval, matched on the change in the aggregate sales hazard (extension)."""
    fe0 = solve_equilibrium_financing(p)
    def agg(fe):
        n = fe.n; eta = fe.eta_c + p.phi_m * fe.eta_m
        return float((eta * n).sum() / n.sum())
    h0 = agg(fe0)
    rows = []
    for A_mult in (0.7, 0.85, 1.0, 1.15, 1.3):
        fe = solve_equilibrium_financing(p.copy(A=p.A * A_mult))
        rows.append({"shock": "A", "value": A_mult, "agg_hazard": agg(fe), "theta": fe.theta,
                     "price_mean": float(((fe.Q_c * fe.pbar_c + fe.Q_m * fe.pbar_m).sum()) / fe.Q), "S": fe.S, "Q": fe.Q,
                     "gap_mean": float(((fe.Q_m * (fe.pbar_m - fe.pbar_c)).sum()) / fe.Q_m.sum()),
                     "mean_tom_months": 12 * float((p.f / (fe.eta_c + p.phi_m * fe.eta_m + p.delta)).sum()),
                     "total_duration_months": 12 * (fe.S + fe.m_pending.sum()) / p.s, "U_m": fe.U_m, "U_c": fe.U_c,
                     "buyer_participation_ok": bool(min(fe.U_c, fe.U_m) >= 0),
                     "cash_share": fe.Q_c.sum() / fe.Q, "converged": fe.converged})
    for phi in (0.55, 0.7, 0.75, 0.8, 0.85, 0.95, 1.0):
        fe = solve_equilibrium_financing(p.copy(phi_m=phi))
        rows.append({"shock": "phi_m", "value": phi, "agg_hazard": agg(fe), "theta": fe.theta,
                     "price_mean": float(((fe.Q_c * fe.pbar_c + fe.Q_m * fe.pbar_m).sum()) / fe.Q), "S": fe.S, "Q": fe.Q,
                     "gap_mean": float(((fe.Q_m * (fe.pbar_m - fe.pbar_c)).sum()) / fe.Q_m.sum()),
                     "mean_tom_months": 12 * float((p.f / (fe.eta_c + p.phi_m * fe.eta_m + p.delta)).sum()),
                     "total_duration_months": 12 * (fe.S + fe.m_pending.sum()) / p.s, "U_m": fe.U_m, "U_c": fe.U_c,
                     "buyer_participation_ok": bool(min(fe.U_c, fe.U_m) >= 0),
                     "cash_share": fe.Q_c.sum() / fe.Q, "converged": fe.converged})
    df = pd.DataFrame(rows); df["agg_hazard_base"] = h0
    save(6, "matching_vs_approval", df, p)


def exp07(p: Params):
    base = solve_equilibrium(p)
    rows = []
    for k_mult in (0.5, 1.0, 2.0):
        for kb_mult in (0.5, 1.0, 2.0):
            e = solve_equilibrium(p.copy(k=p.k * k_mult, kappa_b=p.kappa_b * kb_mult)); m = e.moments(p)
            m.update({"k_mult": k_mult, "kappa_b_mult": kb_mult}); rows.append(m)
    save(7, "entry_costs", pd.DataFrame(rows), p)
    # amplification: elasticities of Lam_s, price, TOM to A and phi in GE vs PE
    rows = []
    for var in ("A", "phi"):
        for h in (0.02,):
            up = solve_equilibrium(p.copy(**{var: getattr(p, var) * (1 + h)})); dn = solve_equilibrium(p.copy(**{var: getattr(p, var) * (1 - h)}))
            upe = partial_equilibrium(base.theta, p.copy(**{var: getattr(p, var) * (1 + h)})); dne = partial_equilibrium(base.theta, p.copy(**{var: getattr(p, var) * (1 - h)}))
            for kind, (u, d) in (("GE", (up, dn)), ("PE", (upe, dne))):
                mu, md = u.moments(p), d.moments(p)
                el = lambda key: (np.log(mu[key]) - np.log(md[key])) / (np.log(1 + h) - np.log(1 - h))
                rows.append({"shock": var, "kind": kind, "el_Lam_s": el("Lam_s"), "el_theta": el("theta") if kind == "GE" else 0.0,
                             "el_price": el("price_mean"), "el_tom": el("mean_tom_sold"), "el_Q": el("Q"), "el_S": el("S"),
                             "el_D_urgent": el("D_urgent")})
    save(7, "amplification", pd.DataFrame(rows), p)
    # uniqueness check across parameter draws: count sign changes of R(theta)
    rng = np.random.default_rng(p.seed)
    rows = []
    for i in range(60):
        pp = p.copy(A=p.A * rng.uniform(0.5, 1.5), beta=rng.uniform(0.2, 0.8), k=p.k * rng.uniform(0.5, 2), kappa_b=p.kappa_b * rng.uniform(0.5, 2),
                    delta=rng.uniform(0.1, 1.0), phi=rng.uniform(0.5, 1.0))
        e = solve_equilibrium(pp)
        rows.append({"draw": i, "A": pp.A, "beta": pp.beta, "k": pp.k, "kappa_b": pp.kappa_b, "delta": pp.delta, "phi": pp.phi,
                     "n_roots": len(e.theta_roots), "theta": e.theta, "converged": e.converged})
    save(7, "uniqueness_draws", pd.DataFrame(rows), p)


def price_density(eq, p: Params, grid: np.ndarray) -> pd.DataFrame:
    """Mixture density of transaction prices: p = (1-beta)V(c) - beta U + beta w, w >= v*(c), weights omega(c)."""
    dens = np.zeros_like(grid); parts = {}
    for i in range(len(p.c)):
        V, vs = eq.sellers.V[i], eq.sellers.vstar[i]
        w = (grid - (1 - p.beta) * V + p.beta * eq.U) / p.beta
        di = np.where(w >= vs, np.asarray(p.G.pdf(w)) / max(float(p.G.surv(vs)), 1e-300) / p.beta, 0.0)
        parts[f"type{i}"] = eq.omega[i] * di; dens += eq.omega[i] * di
    df = pd.DataFrame({"price": grid, "density": dens}); 
    for k_, v in parts.items(): df[k_] = v
    return df


def exp08(p: Params):
    rows = []
    for c_max in (0.05, 0.1, 0.2, 0.3, 0.45, 0.6):
        for n_types in (9,):
            pp = p.with_types(0.02, c_max, n_types, 2.0, 4.0)
            for mult in (0.7, 1.0, 1.3):
                e = solve_equilibrium(pp.copy(A=p.A * mult)); m = e.moments(pp)
                m.update({"c_max": c_max, "A_mult": mult, "c_sd": float(np.sqrt((pp.f * (pp.c - (pp.f * pp.c).sum()) ** 2).sum())),
                          "pbar_range": e.sellers.pbar[0] - e.sellers.pbar[-1], "V_range": e.sellers.V[0] - e.sellers.V[-1]})
                rows.append(m)
    save(8, "heterogeneity_spread", pd.DataFrame(rows), p)
    e = solve_equilibrium(p)
    grid = np.linspace(0.8, 1.25, 451)
    save(8, "price_density_baseline", price_density(e, p, grid), p)
    el = solve_equilibrium(p.copy(A=p.A * 0.6))
    save(8, "price_density_lowA", price_density(el, p, grid), p, {"A_mult": 0.6})
    # beta = 0 limit and beta = 1 limit decomposition
    rows = []
    for beta in (1e-6, 0.25, 0.5, 0.75, 1 - 1e-6):
        e = solve_equilibrium(p.copy(beta=beta)); m = e.moments(p); m["beta"] = beta; rows.append(m)
    save(8, "dispersion_vs_beta", pd.DataFrame(rows), p)


def exp09(p: Params):
    T, dt = p.T, p.dt
    def run(label, p0, path, extra=None):
        tr = transition(p0, path, T=T, dt=dt)
        df = pd.DataFrame({"t": tr.t, "theta": tr.theta, "P_obs": tr.P_obs, "P_fix": tr.P_fix, "Q": tr.Q, "S": tr.S,
                           "share_urgent_sales": tr.share_urgent_sales, "mean_tom_sold_months": 12 * tr.mean_tom_sold})
        for i in (0, len(p0.c) - 1):
            df[f"pbar_c{i}"] = tr.pbar[:, i]; df[f"eta_c{i}"] = tr.eta[:, i]; df[f"n_c{i}"] = tr.n[:, i]
        df["scenario"] = label
        save(9, f"transition_{label}", df, p0, {"converged": tr.converged, "iterations": tr.iterations, "max_update": tr.max_update,
                                                 **(extra or {})})
        return tr
    p_lo = p.copy(A=p.A * 0.8); p_hi = p.copy(A=p.A * 1.2); p_phi = p.copy(phi=p.phi * 0.8)
    run("perm_A_minus20", p, lambda t: p_lo)
    run("perm_A_plus20", p, lambda t: p_hi)
    run("temp_A_minus20_1y", p, lambda t: p_lo if t < 1.0 else p)
    run("perm_phi_minus20", p, lambda t: p_phi)
    pe = p.copy(G=Exponential(rate=8.0, lo=0.7)); run("perm_A_minus20_exponential", pe, lambda t: pe.copy(A=p.A * 0.8))
    pd0 = p.copy(delta=0.0); run("perm_A_minus20_delta0", pd0, lambda t: pd0.copy(A=p.A * 0.8))


def exp10(p: Params):
    dists = {"truncnorm": p.G, "exponential": Exponential(rate=8.0, lo=0.7), "weibull": Weibull(shape=2.0, scale=0.25, lo=0.75),
             "pareto": Pareto(a=4.0, lo=0.8)}
    rows = []
    for gname, G in dists.items():
        for c_max in (0.1, 0.3, 0.6):
            pp = p.with_types(0.02, c_max, 9, 2.0, 4.0).copy(G=G)
            base = solve_equilibrium(pp)
            if not base.converged:
                rows.append({"G": gname, "c_max": c_max, "note": "no equilibrium"}); continue
            for mult in (0.7, 0.8, 0.9, 1.1, 1.2, 1.3):
                pm = pp.copy(A=pp.A * mult)
                # impact (stock predetermined, theta re-solved instantly by free entry? -> use PE at base theta and GE both)
                for kind in ("PE", "GE"):
                    e = partial_equilibrium(base.theta, pm) if kind == "PE" else solve_equilibrium(pm)
                    eta_new = e.sellers.eta
                    om_imp = eta_new * base.n / (eta_new * base.n).sum()
                    P0 = float((base.omega * base.sellers.pbar).sum())
                    P_fix = float((base.omega * e.sellers.pbar).sum()); P_imp = float((om_imp * e.sellers.pbar).sum())
                    P_ss = float((e.omega * e.sellers.pbar).sum())
                    eps = (np.log(eta_new) - np.log(base.sellers.eta)) / np.log(mult)
                    rows.append({"G": gname, "c_max": c_max, "A_mult": mult, "kind": kind, "theta_base": base.theta, "theta_new": e.theta,
                                 "P0": P0, "P_fix": P_fix, "P_impact": P_imp, "P_ss": P_ss,
                                 "within_log": np.log(P_fix / P0), "composition_impact_log": np.log(P_imp / P_fix),
                                 "composition_ss_log": np.log(P_ss / P_fix), "total_impact_log": np.log(P_imp / P0),
                                 "amplification_ratio_impact": np.log(P_imp / P0) / np.log(P_fix / P0),
                                 "eps_patient": eps[0], "eps_urgent": eps[-1], "eps_slope_sign": np.sign(eps[-1] - eps[0]),
                                 "urgent_share_base": base.omega[p.c >= np.median(p.c)].sum() if len(p.c) == 9 else np.nan,
                                 "urgent_share_impact": om_imp[-4:].sum(), "urgent_share_base_": base.omega[-4:].sum(),
                                 "is_ifr": G.is_ifr(), "converged": e.converged})
    save(10, "selection_decomposition", pd.DataFrame(rows), p)
    # hazard-elasticity profiles and h_G*M curves
    rows = []
    for gname, G in dists.items():
        v = np.linspace(G.lo + 1e-3, (G.hi if np.isfinite(G.hi) else G.lo + 1.0) - 1e-3, 200)
        hm = np.asarray(G.hazard(v)) * np.asarray(G.M(v))
        for vi, hi_ in zip(v, hm):
            rows.append({"G": gname, "v": vi, "hG_times_M": hi_})
    save(10, "hazard_times_mrl", pd.DataFrame(rows), p)


def exp11(p: Params):
    rows = []
    for mult in (0.7, 1.0, 1.3):
        pm = p.copy(A=p.A * mult)
        en = solve_equilibrium(pm); ep = solve_equilibrium_posted(pm)
        for proto, e in (("nash", en), ("posted", ep)):
            m = e.moments(pm); m.update({"protocol": proto, "A_mult": mult}); rows.append(m)
            if mult == 1.0:
                save(11, f"types_{proto}", type_table(e, pm), pm)
    save(11, "protocol_comparison", pd.DataFrame(rows), p)


def exp12(p: Params):
    base = solve_equilibrium(p).moments(p)
    keys = ["theta", "price_mean", "mean_tom_sold_months", "D_urgent", "D_mean", "price_between_share", "price_sd_total", "Q", "S", "welfare_total"]
    rows = []
    for par in ("r", "delta", "w0", "k", "kappa_b", "phi", "A", "alpha", "beta"):
        for mult in (0.75, 1.25):
            val = getattr(p, par) * mult
            if par in ("phi", "alpha", "beta"):
                val = min(val, 0.98)
            e = solve_equilibrium(p.copy(**{par: val})); m = e.moments(p)
            row = {"parameter": par, "mult": mult, "value": val, "converged": e.converged}
            for k_ in keys:
                row[k_] = m[k_]; row[f"d_{k_}"] = m[k_] - base[k_]; row[f"dlog_{k_}"] = np.log(m[k_] / base[k_]) if m[k_] > 0 and base[k_] > 0 else np.nan
            rows.append(row)
    for sig in (0.05, 0.15, 0.2):
        G = TruncNormal(1.0, sig, 1 - 3 * sig, 1 + 3 * sig); e = solve_equilibrium(p.copy(G=G)); m = e.moments(p)
        row = {"parameter": "sigma_w", "mult": sig / 0.1, "value": sig, "converged": e.converged}
        for k_ in keys: row[k_] = m[k_]; row[f"d_{k_}"] = m[k_] - base[k_]; row[f"dlog_{k_}"] = np.log(m[k_] / base[k_]) if m[k_] > 0 and base[k_] > 0 else np.nan
        rows.append(row)
    for cm in (0.2, 0.45):
        pp = p.with_types(0.02, cm, 9, 2.0, 4.0); e = solve_equilibrium(pp); m = e.moments(pp)
        row = {"parameter": "c_max", "mult": cm / 0.3, "value": cm, "converged": e.converged}
        for k_ in keys: row[k_] = m[k_]; row[f"d_{k_}"] = m[k_] - base[k_]; row[f"dlog_{k_}"] = np.log(m[k_] / base[k_]) if m[k_] > 0 and base[k_] > 0 else np.nan
        rows.append(row)
    df = pd.DataFrame(rows); 
    for k_ in keys: df[f"base_{k_}"] = base[k_]
    save(12, "sensitivity_oat", df, p)
    # alternative G families at baseline
    rows = []
    for gname, G in (("truncnorm", p.G), ("exponential", Exponential(rate=8.0, lo=0.7)), ("weibull", Weibull(shape=2.0, scale=0.25, lo=0.75)),
                     ("weibull_dfr", Weibull(shape=0.8, scale=0.25, lo=0.75)), ("pareto", Pareto(a=4.0, lo=0.8))):
        e = solve_equilibrium(p.copy(G=G)); m = e.moments(p); m.update({"G": gname, "G_mean": G.mean(), "is_ifr": G.is_ifr()}); rows.append(m)
    save(12, "alternative_distributions", pd.DataFrame(rows), p)
    # grid convergence in the number of seller types and in tolerance
    rows = []
    for n_types in (1, 3, 5, 9, 15, 25, 41, 81):
        pp = p.with_types(0.02, 0.3, n_types, 2.0, 4.0); t0 = time.time(); e = solve_equilibrium(pp); m = e.moments(pp)
        m.update({"n_types": n_types, "seconds": time.time() - t0}); rows.append(m)
    save(12, "convergence_types", pd.DataFrame(rows), p)
    rows = []
    for tol in (1e-6, 1e-8, 1e-10, 1e-12):
        for gp in (20, 60, 120):
            e = solve_equilibrium(p.copy(root_xtol=tol, grid_check_points=gp)); m = e.moments(p); m.update({"root_xtol": tol, "grid_points": gp}); rows.append(m)
    save(12, "convergence_tolerance", pd.DataFrame(rows), p)
    # dt convergence of the transition
    rows = []
    p_lo = p.copy(A=p.A * 0.8)
    for dt in (0.1, 0.05, 0.02, 0.01):
        t0 = time.time(); tr = transition(p, lambda t: p_lo, T=4.0, dt=dt)
        j = int(round(1.0 / dt))
        rows.append({"dt": dt, "theta_impact": tr.theta[0], "P_obs_impact": tr.P_obs[0], "P_obs_1y": tr.P_obs[j], "Q_1y": tr.Q[j],
                     "S_1y": tr.S[j], "theta_end": tr.theta[-1], "converged": tr.converged, "iterations": tr.iterations, "seconds": time.time() - t0})
    save(12, "convergence_dt", pd.DataFrame(rows), p)
    # global sensitivity: Latin-hypercube-like random draws
    rng = np.random.default_rng(p.seed + 1)
    rows = []
    for i in range(150):
        pp = p.copy(r=rng.uniform(0.02, 0.08), delta=rng.uniform(0.1, 1.0), k=rng.uniform(0.005, 0.05), kappa_b=rng.uniform(0.02, 0.1),
                    phi=rng.uniform(0.5, 1.0), A=rng.uniform(2.0, 8.0), alpha=rng.uniform(0.3, 0.7), beta=rng.uniform(0.2, 0.8))
        e = solve_equilibrium(pp); m = e.moments(pp)
        m.update({"draw": i, "r": pp.r, "delta": pp.delta, "k": pp.k, "kappa_b": pp.kappa_b, "phi": pp.phi, "A": pp.A, "alpha": pp.alpha, "beta": pp.beta,
                  "sign_D_urgent_pos": m["D_urgent"] > 0, "sign_eta_increasing": bool(np.all(np.diff(e.sellers.eta) > 0)),
                  "sign_pbar_decreasing": bool(np.all(np.diff(e.sellers.pbar) < 0))})
        rows.append(m)
    save(12, "global_sensitivity", pd.DataFrame(rows), p)


EXPERIMENTS = {1: exp01, 2: exp02, 3: exp03, 4: exp04, 5: exp05, 6: exp06, 7: exp07, 8: exp08, 9: exp09, 10: exp10, 11: exp11, 12: exp12}


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default=str(ROOT / "config" / "baseline.yaml"))
    ap.add_argument("--exp", default="all")
    a = ap.parse_args(argv)
    p = load_params(a.config)
    np.random.seed(p.seed)
    todo = sorted(EXPERIMENTS) if a.exp == "all" else [int(x) for x in a.exp.split(",")]
    log.info("config %s hash %s; experiments %s", a.config, p.config_hash(), todo)
    for e in todo:
        t0 = time.time(); log.info("=== Experiment %d: %s", e, HYPOTHESES[e][:80])
        try:
            EXPERIMENTS[e](p)
            log.info("experiment %d done in %.1f s", e, time.time() - t0)
        except Exception as ex:                       # report, never hide
            log.exception("experiment %d FAILED: %s", e, ex)
            (RAW / f"exp{e:02d}_FAILED.json").write_text(json.dumps({"experiment": e, "error": repr(ex)}))


if __name__ == "__main__":
    main()