"""WP18 — experiment orchestration (CLAUDE.md Phase 3, experiments 1-10).

Usage:  python -m src.simulation [--groups g1 g2 ...] [--workers N] [--quick]

Outputs (all under results/data/, every number used in a table or figure comes from here):
  parameters.csv                       period-unit parameters with sources
  baseline/summary.json                prices, moments, convergence, Bellman residuals, thresholds
  baseline/profiles_by_age.csv         ownership / wealth / debt / consumption by age
  baseline/by_gift_quintile.csv        ownership and net worth by age x inheritance quintile
  baseline/mobility.json               parent-child tenure transition, rank-rank slopes, lambda
  baseline/policy_age*.csv             policy functions on the cash grid (figures)
  baseline/thresholds.csv              ownership-entry thresholds by age and z (+ HTM check)
  baseline/cross_section.npz           pooled cross-section (compressed)
  baseline/equal_income_paths.csv      experiment 2 (deterministic equal-income paths)
  baseline/price_drop_accounting.csv   experiment 5 (PE accounting of unanticipated price drops)
  <group>/<name>.json, <name>_profiles.csv       stationary scenarios (GE or PE, labelled)
  inflation/<name>.json, _path.csv, _impact.csv  transitions (GE / PE / static, labelled)
  welfare/welfare_<name>.csv           consumption-equivalent variations (experiment 9)
  multiplicity/price_scan_<name>.csv   experiment 3/7 multiplicity search
  diagnostics.json / diagnostics.csv   convergence and residuals of every run
Logs: results/logs/<name>.jsonl (structured) and results/logs/simulate.log (summary).
"""
from __future__ import annotations

import argparse
import json
import os
import pickle
import time
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import get_context
from pathlib import Path

import numpy as np
import pandas as pd

from . import calibration as C
from . import distributions as D
from . import equilibrium as E
from . import model as M
from .household_solver import NEG, ownership_threshold

DATA = M.RESULTS_DIR / "data"
LOGS = M.RESULTS_DIR / "logs"


# --------------------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------------------
def _js(o):
    if isinstance(o, dict):
        return {str(k): _js(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [_js(v) for v in o]
    if isinstance(o, np.ndarray):
        return o.tolist()
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.bool_,)):
        return bool(o)
    if isinstance(o, float) and not np.isfinite(o):
        return None
    return o


def save_json(path: Path, obj: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w") as fh:
        json.dump(_js(obj), fh, indent=1)


def apply_overrides(cfg: dict, overrides: dict | None) -> dict:
    cfg = json.loads(json.dumps(cfg))
    for key, val in (overrides or {}).items():
        node = cfg
        parts = key.split(".")
        for part in parts[:-1]:
            node = node[part]
        node[parts[-1]] = val
    return cfg


# --------------------------------------------------------------------------------------
# summaries of a stationary result
# --------------------------------------------------------------------------------------
def profiles_by_age(P: M.Params, rec: D.PeriodRecord) -> pd.DataFrame:
    nw = rec.net_worth1(P)
    own = rec.owner1()
    rows = []
    for j in range(1, P.J + 1):
        m = rec.age == j
        o = m & own; r = m & ~own
        hv = np.where(own, P.h_grid[np.maximum(rec.h1, 0)], 0.0)
        rows.append(dict(age=j, age_years=f"{25 + 5 * (j - 1)}-{29 + 5 * (j - 1)}", ownership=float(own[m].mean()),
                         buy_rate=float((rec.choice[m] == 2).mean()), sell_to_rent_rate=float(((rec.choice[m] == 3) & (rec.h0[m] >= 0)).mean()),
                         refinance_rate=float((rec.choice[m] == 1).mean()),
                         nw_mean=float(nw[m].mean()), nw_median=float(np.median(nw[m])),
                         nw_p25=float(np.quantile(nw[m], 0.25)), nw_p75=float(np.quantile(nw[m], 0.75)),
                         nw_owner_median=float(np.median(nw[o])) if o.any() else np.nan,
                         nw_renter_median=float(np.median(nw[r])) if r.any() else np.nan,
                         nw_owner_mean=float(nw[o].mean()) if o.any() else np.nan,
                         nw_renter_mean=float(nw[r].mean()) if r.any() else np.nan,
                         liquid_mean=float(rec.a1[m].mean()), debt_mean=float(rec.b1[m].mean()),
                         debt_owner_mean=float(rec.b1[o].mean()) if o.any() else np.nan,
                         ltv_owner_mean=float((rec.b1[o] / (rec.p * hv[o])).mean()) if o.any() else np.nan,
                         share_mortgage=float((rec.b1[m] > 1e-8).mean()),
                         house_size_mean=float(hv[o].mean()) if o.any() else np.nan,
                         rent_size_mean=float(rec.s[r].mean()) if r.any() else np.nan,
                         c_mean=float(rec.c[m].mean()), c_owner=float(rec.c[o].mean()) if o.any() else np.nan,
                         c_renter=float(rec.c[r].mean()) if r.any() else np.nan,
                         income_mean=float((P.ybar[j - 1] * P.z_grid[rec.z[m]]).mean()),
                         gini_nw=D.gini(nw[m]), n=int(m.sum())))
    return pd.DataFrame(rows)


def by_gift_quintile(P: M.Params, rec: D.PeriodRecord) -> pd.DataFrame:
    """Ownership and net worth by age and quintile of the gift received at age 1 (young and old
    members both carry it). Quintiles are defined on the age-1 gift distribution; ties at zero
    are broken by rank so that quintiles have equal mass (zero-gift households fall in the
    bottom groups)."""
    g = rec.gift_received
    ok = np.isfinite(g)
    g1 = np.sort(g[(rec.age == 1) & ok])
    # rank in the age-1 gift distribution with uniform tie-breaking (mass at zero is split evenly)
    rng = np.random.default_rng(3)
    lo = np.searchsorted(g1, g[ok], side="left"); hi = np.searchsorted(g1, g[ok], side="right")
    rank = (lo + rng.uniform(size=ok.sum()) * (hi - lo)) / len(g1)
    q = np.minimum((rank * 5).astype(int), 4)   # 0..4
    nw = rec.net_worth1(P)[ok]; own = rec.owner1()[ok]; age = rec.age[ok]
    rows = []
    for j in range(1, P.J + 1):
        for k in range(5):
            m = (age == j) & (q == k)
            if m.sum() == 0:
                continue
            rows.append(dict(age=j, gift_quintile=k + 1, ownership=float(own[m].mean()), nw_mean=float(nw[m].mean()),
                             nw_median=float(np.median(nw[m])), gift_mean=float(g[ok][m].mean()), n=int(m.sum())))
    return pd.DataFrame(rows)


def mobility_stats(P: M.Params, rec: D.PeriodRecord) -> dict:
    """Parent-child tenure transition at age j_gift (both measured at age 7, the gift age),
    persistence lambda = q_OO - q_RO (Part A analogue), rank-rank slope of net worth at age 7."""
    m = (rec.age == P.j_gift) & rec.is_old & (rec.parent_owner >= 0)
    child_own = rec.owner1()[m].astype(float)
    par_own = rec.parent_owner[m]
    q_OO = float(child_own[par_own == 1].mean()) if (par_own == 1).any() else np.nan
    q_RO = float(child_own[par_own == 0].mean()) if (par_own == 0).any() else np.nan
    nw = rec.net_worth1(P)[m]
    rr = D.rank_rank_slope(rec.parent_nw[m], nw)
    corr = float(np.corrcoef(par_own, child_own)[0, 1]) if par_own.std() > 0 and child_own.std() > 0 else np.nan
    lam = q_OO - q_RO
    # upward mobility of renter dynasties: child owns at age 3 given parent renter at the gift
    m3 = (rec.age == 3) & (~rec.is_old) & (rec.parent_owner >= 0)
    own3 = rec.owner1()[m3]; po3 = rec.parent_owner[m3]
    return dict(q_OO=q_OO, q_RO=q_RO, lambda_persistence=lam, half_life_generations=float(np.log(2) / -np.log(lam)) if 0 < lam < 1 else np.inf,
                tenure_correlation=corr, rank_rank_slope_nw=rr,
                q_RO_age3=float(own3[po3 == 0].mean()) if (po3 == 0).any() else np.nan,
                q_OO_age3=float(own3[po3 == 1].mean()) if (po3 == 1).any() else np.nan,
                share_parent_owner=float((par_own == 1).mean()), n=int(m.sum()))


def policy_tables(P: M.Params, res: E.SSResult) -> dict[str, pd.DataFrame]:
    out = {}
    pr = res.ev.prices
    for j in (1, 3, 7):
        pol = res.ev.pols[j - 1]
        rows = []
        for iz in range(P.n_z):
            y = P.ybar[j - 1] * P.z_grid[iz]
            for ix in range(P.n_x):
                x = P.x_grid[ix]
                wb, wr = pol.W_B[iz, ix], pol.W_Rs[iz, ix]
                rows.append(dict(age=j, z_index=iz, z=float(P.z_grid[iz]), cash=x, a_equiv=(x - y) / pr.R_a,
                                 W_buy=wb if wb > NEG / 2 else np.nan, W_rent=wr if wr > NEG / 2 else np.nan,
                                 buy=int((wb > wr) and wb > NEG / 2), h_buy=int(pol.hB[iz, ix]), B_buy=float(pol.BB[iz, ix]),
                                 a_buy=float(pol.aB[iz, ix]), a_rent=float(pol.aRs[iz, ix]), gift_buy=float(pol.gB[iz, ix]),
                                 gift_rent=float(pol.gRs[iz, ix]),
                                 c_rent=float(max(x - pol.aRs[iz, ix] - pol.gRs[iz, ix], 0) * (1 - P.alpha))))
        out[f"policy_age{j}"] = pd.DataFrame(rows)
    return out


def thresholds_table(P: M.Params, res: E.SSResult) -> pd.DataFrame:
    rows = []
    for j in (1, 2, 3):
        th = ownership_threshold(P, res.ev.pols[j - 1], res.ev.prices, j)
        for iz, v in th.items():
            rows.append(dict(age=j, z_index=iz, **v))
    return pd.DataFrame(rows)


def summarize_ss(P: M.Params, res: E.SSResult, name: str, group: str, scenario: dict, t_sec: float) -> dict:
    rec = res.ev.rec
    mom = C.model_moments(P, rec, res.p, res.rho)
    mob = mobility_stats(P, rec)
    nw = rec.net_worth1(P)
    own = rec.owner1()
    # owner-renter gaps at fixed ages (equal income by design: same income process for all)
    gaps = {}
    for j in (2, 4, 7, 10):
        m = rec.age == j
        gaps[f"nw_gap_owner_renter_age{j}"] = float(nw[m & own].mean() - nw[m & ~own].mean()) if (m & ~own).any() and (m & own).any() else np.nan
    z1 = rec.z[rec.age == 1]
    return dict(name=name, group=group, scenario=scenario, equilibrium=res.equilibrium, converged=bool(res.converged),
                iterations=int(res.iterations), time_sec=float(t_sec), p=res.p, p_r=res.p_r, ptil_r=res.ptil_r, rho=res.rho,
                rent_price_annual=res.rho / res.p_r / P.period_years, H=res.ev.H_d, Hr=res.ev.S_d,
                H_bar=P.H_bar, Hr_bar=P.Hr_bar, lottery=res.lottery[0], rebate=res.rebate,
                residuals=res.residuals, excess_owner=float(res.residuals[0]), excess_rental=float(res.residuals[1]),
                bellman=res.bellman, forced_share=res.ev.forced_share, clip_share=res.ev.clip_share,
                moments=mom, mobility=mob, gaps=gaps,
                thresholds={j: {iz: v["a_hat"] for iz, v in th.items()} for j, th in res.thresholds.items()},
                threshold_sign_changes={j: {iz: v["sign_changes"] for iz, v in th.items()} for j, th in res.thresholds.items()},
                threshold_htm={j: {iz: [v["htm_buyer"], v["htm_renter"]] for iz, v in th.items()} for j, th in res.thresholds.items()},
                mean_nw=float(nw.mean()), gini_nw=D.gini(nw), total_debt=float(rec.b1.sum() / (len(rec.age) / P.J)),
                total_deposits=float(rec.a1.sum() / (len(rec.age) / P.J)),
                params=dict(theta=P.theta, psi=P.psi, eps=P.eps, eps_r=P.eps_r, nu=P.nu, tau_b=P.tau_b, r=P.r, pibar=P.pibar,
                            phi=P.phi, sigma_seg=P.sigma_seg, sigma_z=P.sigma_z, beta=P.beta, gamma=P.gamma, alpha=P.alpha,
                            kappa_b=P.kappa_b, kappa_s=P.kappa_s, chi=P.chi, n_a=P.n_a, n_b=P.n_b, n_x=P.n_x, n_dyn=P.n_dyn,
                            seed=P.seed, dep_premium=P.dep_premium))


def save_cross_section(P: M.Params, rec: D.PeriodRecord, path: Path, subsample: int = 1):
    path.parent.mkdir(parents=True, exist_ok=True)
    sl = slice(None, None, subsample)
    np.savez_compressed(path, age=rec.age[sl], z=rec.z[sl], h0=rec.h0[sl], b0=rec.b0[sl], a0=rec.a0[sl], choice=rec.choice[sl],
                        c=rec.c[sl], s=rec.s[sl], h1=rec.h1[sl], b1=rec.b1[sl], a1=rec.a1[sl], gift=rec.gift[sl], estate=rec.estate[sl],
                        gift_received=rec.gift_received[sl], parent_owner=rec.parent_owner[sl], parent_nw=rec.parent_nw[sl],
                        estate_received=rec.estate_received[sl], is_old=rec.is_old[sl], nw0=rec.net_worth0(P)[sl],
                        nw1=rec.net_worth1(P)[sl], pay=rec.pay[sl], Bnew=rec.Bnew[sl], flag=rec.flag[sl], u=rec.u[sl],
                        p=rec.p, rho=rec.rho, h_grid=P.h_grid, z_grid=P.z_grid)


# --------------------------------------------------------------------------------------
# Experiment 2: equal-income deterministic paths
# --------------------------------------------------------------------------------------
def equal_income_paths(P: M.Params, res: E.SSResult) -> pd.DataFrame:
    """Households with the same (median) income path and different initial transfers; policies
    of the baseline equilibrium; estate at j_estate = the lottery's median point (deterministic).
    Also the two households straddling the age-1 ownership threshold (owner vs renter at equal
    income and almost equal endowment), and a renter who invests the down payment at R_a."""
    S = D.stack_policies(res.ev.pols)
    pr = res.ev.prices
    iz = P.n_z // 2
    g1 = res.ev.rec.gift_received[(res.ev.rec.age == 1) & ~res.ev.rec.is_old]
    qs = {"q10": 0.1, "q25": 0.25, "q50": 0.5, "q75": 0.75, "q90": 0.9}
    a1_vals = {k: float(np.quantile(g1, v)) for k, v in qs.items()}
    th = ownership_threshold(P, res.ev.pols[0], pr, 1)[iz]["a_hat"]
    if np.isfinite(th):
        a1_vals["threshold_minus"] = max(th - 0.01, 0.0)
        a1_vals["threshold_plus"] = th + 0.01
    names = list(a1_vals.keys())
    a1 = np.array([a1_vals[k] for k in names])
    z1 = np.full(len(a1), iz)
    med_lot = (np.array([float(np.median(res.lottery[0]))]), np.array([1.0]))
    out = D.simulate_cohort(P, S, pr, z1, a1, med_lot, fixed_z=True)
    rows = []
    for k, nm in enumerate(names):
        for j in range(1, P.J + 1):
            rows.append(dict(path=nm, a1=float(a1[k]), age=j, owner=int(out["own"][j - 1, k]), nw=float(out["nw"][j - 1, k]),
                             liquid=float(out["a"][j - 1, k]), debt=float(out["b"][j - 1, k]), c=float(out["c"][j - 1, k]),
                             s=float(out["s"][j - 1, k]), h=int(out["h"][j - 1, k]), gift_given=float(out["gift"][k]),
                             lifetime_V=float(out["V"][k])))
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------------------
# Experiment 5 (accounting part): unanticipated price drop on the cross-section
# --------------------------------------------------------------------------------------
def price_drop_accounting(P: M.Params, rec: D.PeriodRecord, drops=(0.1, 0.2, 0.3)) -> pd.DataFrame:
    own = rec.owner1()
    hv = np.where(own, P.h_grid[np.maximum(rec.h1, 0)], 0.0) * rec.p
    nw = rec.net_worth1(P)
    rows = []
    for x in drops:
        eq_after = (1 - P.kappa_s) * (1 - x) * hv - rec.b1
        loss = own * (1 - P.kappa_s) * x * hv
        for j in list(range(1, P.J + 1)) + ["all"]:
            m = own if j == "all" else (own & (rec.age == j))
            if m.sum() == 0:
                continue
            ltv = rec.b1[m] / hv[m]
            rows.append(dict(drop=x, age=j, n_owners=int(m.sum()), underwater_share=float((eq_after[m] < 0).mean()),
                             equity_wiped_share=float((eq_after[m] <= 0.0).mean()),
                             mean_loss_over_nw=float((loss[m] / np.maximum(nw[m], 1e-6)).mean()),
                             mean_loss_over_income=float(loss[m].mean() / (P.ybar[min(j, P.J) - 1] if j != "all" else 1.0)),
                             mean_ltv=float(ltv.mean()), share_ltv_above_x_star=float((ltv > (1 - P.kappa_s) * (1 - x)).mean())))
        # new borrowers evaluated at origination (before the first payment): LTV = Bnew / (p h)
        m = rec.Bnew > 0
        if m.any():
            ltv0 = rec.Bnew[m] / hv[m]
            rows.append(dict(drop=x, age="new_loans_at_origination", n_owners=int(m.sum()),
                             underwater_share=float((((1 - P.kappa_s) * (1 - x) * hv[m]) < rec.Bnew[m]).mean()),
                             equity_wiped_share=float((((1 - P.kappa_s) * (1 - x) * hv[m]) <= rec.Bnew[m]).mean()),
                             mean_loss_over_nw=float(((1 - P.kappa_s) * x * hv[m] / np.maximum(nw[m], 1e-6)).mean()),
                             mean_loss_over_income=float(((1 - P.kappa_s) * x * hv[m]).mean()),
                             mean_ltv=float(ltv0.mean()), share_ltv_above_x_star=float((ltv0 > (1 - P.kappa_s) * (1 - x)).mean())))
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------------------
# Experiment 9: consumption-equivalent welfare across stationary equilibria
# --------------------------------------------------------------------------------------
def welfare_vs_baseline(P_base: M.Params, base: E.SSResult, P_alt: M.Params, alt: E.SSResult, name: str) -> pd.DataFrame:
    """CEV at age 1 by (z_1, baseline gift quintile): lambda such that scaling non-housing
    consumption by (1 + lambda) in every period of the BASELINE makes the age-1 household as
    well off as under the alternative, at the same initial state (z_1, a_1).
    V_base(lambda) = (1+lambda)^((1-alpha)(1-gamma)) U_base + Phi_base, where U_base (flow
    utility) and Phi_base (warm glow) are obtained by cohort simulation under baseline policies;
    V_alt is the alternative's age-1 value function at (z_1, a_1).
    Note: the comparison is between stationary equilibria (not along a transition)."""
    S_b = D.stack_policies(base.ev.pols)
    g1 = base.ev.rec.gift_received[(base.ev.rec.age == 1) & ~base.ev.rec.is_old]
    edges = np.quantile(g1, [0.2, 0.4, 0.6, 0.8])
    rows = []
    expo = (1 - P_base.alpha) * (1 - P_base.gamma)
    n_per = 3000
    rng = np.random.default_rng(5)
    th = {iz: v["a_hat"] for iz, v in ownership_threshold(P_base, base.ev.pols[0], base.ev.prices, 1).items()}
    for k in range(6):
        if k < 5:
            lo = 0.0 if k == 0 else edges[k - 1]; hi = edges[k] if k < 4 else np.quantile(g1, 0.995)
            a1 = rng.uniform(lo, max(hi, lo + 1e-6), size=n_per) if hi > lo else np.full(n_per, lo)
        for iz in range(P_base.n_z):
            if k == 5:
                # marginal households: within +/- 10 % of the age-1 ownership threshold for this z
                if not np.isfinite(th[iz]) or th[iz] <= 0:
                    continue
                a1 = rng.uniform(0.9 * th[iz], 1.1 * th[iz], size=n_per)
            z1 = np.full(n_per, iz)
            coh = D.simulate_cohort(P_base, S_b, base.ev.prices, z1, a1, base.lottery, seed=100 + k)
            U_sim, Phi = coh["U"].mean(), coh["Phi"].mean()
            Vb_sim = U_sim + Phi
            Vb = np.mean(np.interp(a1, P_base.a_grid, base.ev.pols[0].V_rent[iz]))
            Va = np.mean(np.interp(a1 + alt.rebate - base.rebate, P_alt.a_grid, alt.ev.pols[0].V_rent[iz]))
            # consistent decomposition: the flow part is V_base - Phi_base (Phi from the simulation),
            # so that the identity holds exactly at the baseline; U_sim is kept as a check
            U = Vb - Phi
            ratio = (Va - Phi) / U
            cev = ratio ** (1.0 / expo) - 1.0 if ratio > 0 else np.nan
            rows.append(dict(scenario=name, gift_quintile=(k + 1) if k < 5 else "marginal", z_index=iz, z=float(P_base.z_grid[iz]), a1_mean=float(a1.mean()),
                             V_base=float(Vb), V_base_simulated=float(Vb_sim), V_alt=float(Va), U_base=float(U), U_base_simulated=float(U_sim), Phi_base=float(Phi),
                             cev=float(cev), weight=float(P_base.pi_z[iz] / 5)))
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------------------
# transition summaries
# --------------------------------------------------------------------------------------
def transition_path_table(P: M.Params, tr: E.TransitionResult, ss: E.SSResult) -> pd.DataFrame:
    rows = []
    for t, rec in enumerate(tr.recs):
        nw = rec.net_worth1(P); own = rec.owner1()
        young = rec.age <= 3
        rows.append(dict(t=t, years=5 * t, p=float(tr.p[t]), p_r=float(tr.p_r[t]), rho=float(tr.rho[t]),
                         p_pct=100 * (tr.p[t] / ss.p - 1), rho_pct=100 * (tr.rho[t] / ss.rho - 1),
                         H=float(tr.H[t]), Hr=float(tr.Hr[t]), ownership=float(own.mean()), ownership_young=float(own[young].mean()),
                         buy_rate_young=float((rec.choice[young] == 2).mean()),
                         excess_owner=float(tr.residuals[t]), excess_rental=float(tr.residuals[tr.T + t]),
                         debt_total=float(rec.b1.sum() / (len(rec.age) / P.J)), deposits_total=float(rec.a1.sum() / (len(rec.age) / P.J)),
                         nw_mean=float(nw.mean()), nw_owner_mean=float(nw[own].mean()), nw_renter_mean=float(nw[~own].mean()),
                         gini_nw=D.gini(nw), c_mean=float(rec.c.mean()),
                         i_new=float(tr.prices[t].i_new), R_a=float(tr.prices[t].R_a), pi_next=float(tr.prices[t].pi_next)))
    return pd.DataFrame(rows)


def impact_table(P: M.Params, tr: E.TransitionResult, ss: E.SSResult, pe: E.TransitionResult | None) -> pd.DataFrame:
    """Impact-period (t = 0) redistribution by age and start-of-period tenure: real revaluation
    of nominal positions (identity, PE), consumption and net-worth changes vs the stationary
    path, with and without price adjustment."""
    rec0 = tr.recs[0]
    pr0 = tr.prices[0]
    pi_u = tr.paths["pi_u"]
    ss_rec = ss.ev.rec
    # stationary comparison period: same households at t=0 under SS policies and prices (the
    # panel state is identical, only the shock differs) -> simulate one SS period from ss.ev.panel
    S_ss = D.stack_policies(ss.ev.pols)
    panel = ss.ev.panel.copy()
    rec_ss = D.simulate_period(P, S_ss, ss.ev.prices, panel, (P.seed, 100000))
    rows = []
    own0 = rec0.h0 >= 0
    b_pre = rec_ss.b0; a_pre = rec_ss.a0
    bf = float(tr.paths["balance_factor"])
    # revaluation gains: debt erosion (1 - bf) * b_pre ; deposit loss a_pre * (R - R_a0)
    debt_gain = (1 - bf) * b_pre
    dep_loss = a_pre * (P.R * (1 + P.dep_premium) - pr0.R_a)
    wage_loss = (P.ybar[rec0.age - 1] * P.z_grid[rec0.z]) * (1 - pr0.wage_factor)
    rent_gain = np.where(rec0.h0 < 0, rec_ss.s * ss.rho * (1 - pr0.rent_rigid_factor), 0.0)
    net_nominal = debt_gain - dep_loss
    for grp_name, grp in [("all", np.ones(len(rec0.age), bool)), ("owners", own0), ("renters", ~own0)]:
        for j in list(range(1, P.J + 1)) + ["all"]:
            m = grp if j == "all" else (grp & (rec0.age == j))
            if m.sum() == 0:
                continue
            rows.append(dict(scenario=tr.shock.contract, group=grp_name, age=j, n=int(m.sum()),
                             debt_erosion_gain=float(debt_gain[m].mean()), deposit_loss=float(dep_loss[m].mean()),
                             wage_loss=float(wage_loss[m].mean()), rent_rigidity_gain=float(rent_gain[m].mean()),
                             net_nominal_gain=float(net_nominal[m].mean()),
                             c_ss=float(rec_ss.c[m].mean()),
                             d_consumption=float(rec0.c[m].mean() - rec_ss.c[m].mean()),
                             d_consumption_pct=float(100 * (rec0.c[m].mean() / rec_ss.c[m].mean() - 1)),
                             d_consumption_pe_pct=float(100 * (pe.recs[0].c[m].mean() / rec_ss.c[m].mean() - 1)) if pe is not None else np.nan,
                             d_nw=float(rec0.net_worth1(P)[m].mean() - rec_ss.net_worth1(P)[m].mean()),
                             d_ownership=float(rec0.owner1()[m].mean() - rec_ss.owner1()[m].mean()),
                             d_consumption_pe=float(pe.recs[0].c[m].mean() - rec_ss.c[m].mean()) if pe is not None else np.nan,
                             d_nw_pe=float(pe.recs[0].net_worth1(P)[m].mean() - rec_ss.net_worth1(P)[m].mean()) if pe is not None else np.nan,
                             d_ownership_pe=float(pe.recs[0].owner1()[m].mean() - rec_ss.owner1()[m].mean()) if pe is not None else np.nan))
    # lender side (per unit mass of one age): mirror image of households' net nominal gain
    mass = len(rec0.age) / P.J
    lender = dict(scenario=tr.shock.contract, group="lender", age="all", n=int(len(rec0.age)),
                  debt_erosion_gain=-float(debt_gain.sum() / mass), deposit_loss=-float(dep_loss.sum() / mass),
                  wage_loss=0.0, rent_rigidity_gain=0.0, net_nominal_gain=-float(net_nominal.sum() / mass),
                  c_ss=np.nan, d_consumption=np.nan, d_consumption_pct=np.nan, d_consumption_pe_pct=np.nan, d_nw=np.nan, d_ownership=np.nan,
                  d_consumption_pe=np.nan, d_nw_pe=np.nan, d_ownership_pe=np.nan)
    rows.append(lender)
    return pd.DataFrame(rows)


def transition_welfare(P: M.Params, tr: E.TransitionResult, ss: E.SSResult, pols0: list) -> pd.DataFrame:
    """Group-level CEV for households alive at the shock: V_{j,0}(state after the shock) vs
    the stationary V_j(state); Phi_base (warm-glow part) by remaining-lifetime simulation under
    the stationary policies, U_base = V_ss - Phi_base (consistent flow part; the simulated U is
    reported as a check). lambda_group = ((mean V_alt - mean Phi_base)/mean U_base)^(1/expo) - 1."""
    rec0 = tr.recs[0]
    S_ss = D.stack_policies(ss.ev.pols)
    U, Phi = D.remaining_lifetime(P, S_ss, ss.ev.prices, rec0, ss.lottery)
    # stationary and post-shock values at the start-of-period state (interpolated)
    bf = float(tr.paths["balance_factor"])
    V_ss = _value_at_states(P, ss.ev.pols, rec0, b_scale=1.0 / bf if bf > 0 else 1.0)   # pre-shock balances
    V_alt = _value_at_states(P, pols0, rec0, b_scale=1.0)
    expo = (1 - P.alpha) * (1 - P.gamma)
    rows = []
    own0 = rec0.h0 >= 0
    hasdebt = rec0.b0 > 1e-8
    for grp_name, grp in [("all", np.ones(len(rec0.age), bool)), ("owners", own0), ("owners_with_mortgage", own0 & hasdebt),
                          ("owners_no_mortgage", own0 & ~hasdebt), ("renters", ~own0)]:
        for j in list(range(1, P.J + 1)) + ["all"]:
            m = grp if j == "all" else (grp & (rec0.age == j))
            ok = m & np.isfinite(V_ss) & np.isfinite(V_alt) & (V_ss > NEG / 2) & (V_alt > NEG / 2)
            if ok.sum() < 20:
                continue
            Ub_sim, Pb = U[ok].mean(), Phi[ok].mean()
            Vs = V_ss[ok].mean()
            Ub = Vs - Pb                       # consistent flow part (identity exact at the baseline)
            ratio = (V_alt[ok].mean() - Pb) / Ub
            cev = ratio ** (1 / expo) - 1 if ratio > 0 else np.nan
            rows.append(dict(scenario=tr.shock.contract, group=grp_name, age=j, n=int(ok.sum()), V_ss=float(Vs),
                             V_shock=float(V_alt[ok].mean()), U_base=float(Ub), U_base_simulated=float(Ub_sim), Phi_base=float(Pb),
                             V_check_sim_vs_fn=float((Ub_sim + Pb) / Vs - 1), cev=float(cev),
                             share_better_off=float((V_alt[ok] > V_ss[ok]).mean())))
    return pd.DataFrame(rows)


def _value_at_states(P: M.Params, pols: list, rec: D.PeriodRecord, b_scale: float = 1.0) -> np.ndarray:
    V = np.full(len(rec.age), np.nan)
    for j in range(1, P.J + 1):
        m = rec.age == j
        pol = pols[j - 1]
        for iz in range(P.n_z):
            mm = m & (rec.z == iz)
            if not mm.any():
                continue
            r = mm & (rec.h0 < 0)
            if r.any():
                V[r] = np.interp(rec.a0[r], P.a_grid, pol.V_rent[iz])
            o = mm & (rec.h0 >= 0)
            if o.any():
                idx = np.where(o)[0]
                for i in idx:
                    ih, n_idx = rec.h0[i], rec.n0[i] - 1
                    bb = rec.b0[i] * b_scale
                    jb, wb = _loc(P.b_grid, bb); ja, wa = _loc(P.a_grid, rec.a0[i])
                    v = pol.V_own[iz, ih, jb:jb + 2, n_idx, ja:ja + 2]
                    if v.min() <= NEG / 2:
                        V[i] = NEG
                    else:
                        V[i] = (1 - wb) * ((1 - wa) * v[0, 0] + wa * v[0, 1]) + wb * ((1 - wa) * v[1, 0] + wa * v[1, 1])
    return V


def _loc(grid, x):
    n = len(grid)
    if x <= grid[0]:
        return 0, 0.0
    if x >= grid[-1]:
        return n - 2, 1.0
    k = int(np.searchsorted(grid, x, side="right")) - 1
    return k, (x - grid[k]) / (grid[k + 1] - grid[k])


def summarize_transition(P: M.Params, tr: E.TransitionResult, ss: E.SSResult, name: str, scen: dict, path_df: pd.DataFrame,
                         impact_df: pd.DataFrame, welfare_df: pd.DataFrame | None, pe: E.TransitionResult | None) -> dict:
    rec0 = tr.recs[0]
    mass = len(rec0.age) / P.J
    bf = float(tr.paths["balance_factor"])
    S_ss = D.stack_policies(ss.ev.pols)
    panel = ss.ev.panel.copy()
    rec_ss = D.simulate_period(P, S_ss, ss.ev.prices, panel, (P.seed, 100000))
    debt_gain = (1 - bf) * rec_ss.b0; dep_loss = rec_ss.a0 * (P.R * (1 + P.dep_premium) - tr.prices[0].R_a)
    # mobility along the path: parent-child tenure at age 7 in the first generation after the shock
    mob = []
    for t, rec in enumerate(tr.recs):
        mob.append(mobility_stats(P, rec))
    return dict(name=name, group="inflation", scenario=scen, equilibrium=tr.equilibrium, converged=bool(tr.converged),
                iterations=int(tr.iterations), time_sec=float(tr.time_sec), max_residual=float(np.max(np.abs(tr.residuals))),
                T=tr.T, pi_u=float(tr.paths["pi_u"]), contract=tr.shock.contract, persistence=tr.shock.persistence,
                balance_factor=bf, p_impact_pct=float(100 * (tr.p[0] / ss.p - 1)), rho_impact_pct=float(100 * (tr.rho[0] / ss.rho - 1)),
                p_peak_pct=float(100 * (np.max(np.abs(tr.p / ss.p - 1)) * np.sign(tr.p[np.argmax(np.abs(tr.p / ss.p - 1))] - ss.p))),
                p_path_pct=(100 * (tr.p / ss.p - 1)).tolist(), rho_path_pct=(100 * (tr.rho / ss.rho - 1)).tolist(),
                ownership_impact=float(rec0.owner1().mean()), ownership_ss=float(rec_ss.owner1().mean()),
                d_ownership_impact_pp=float(100 * (rec0.owner1().mean() - rec_ss.owner1().mean())),
                d_ownership_impact_pe_pp=float(100 * (pe.recs[0].owner1().mean() - rec_ss.owner1().mean())) if pe is not None else None,
                aggregate_debt_erosion=float(debt_gain.sum() / mass), aggregate_deposit_loss=float(dep_loss.sum() / mass),
                lender_revaluation=-float((debt_gain - dep_loss).sum() / mass),
                # lender side computed from its balance sheet (M - D) x erosion: must equal minus the household sum
                zero_sum_check=float(((1 - bf) * rec_ss.b0.sum() - (P.R * (1 + P.dep_premium) - tr.prices[0].R_a) * rec_ss.a0.sum()) / mass
                                     + (-(debt_gain - dep_loss).sum() / mass)),
                p_path_sign_changes=int(np.sum(np.sign(tr.p[1:] / ss.p - 1) != np.sign(tr.p[:-1] / ss.p - 1))),
                residual_at_T=float(max(abs(tr.residuals[tr.T - 1]), abs(tr.residuals[2 * tr.T - 1]))),
                mobility_path=mob, mobility_ss=mobility_stats(P, ss.ev.rec),
                welfare_all=float(welfare_df.loc[(welfare_df.group == "all") & (welfare_df.age == "all"), "cev"].iloc[0]) if welfare_df is not None else None,
                path=path_df.to_dict(orient="list"))


# --------------------------------------------------------------------------------------
# runners
# --------------------------------------------------------------------------------------
def run_baseline(cfg: dict, quick: bool = False) -> tuple[M.Params, E.SSResult, dict]:
    t0 = time.time()
    P = M.build_params(cfg)
    log = E.RunLog("baseline")
    log("start", quick=quick, n_dyn=P.n_dyn)
    H_bar, Hr_bar, cal = E.calibrate_supply(P, log)
    log("calibrated_supply", H_bar=H_bar, Hr_bar=Hr_bar)
    res = E.solve_stationary(P, p0=P.p_ref, ptil0=P.p_ref, lottery0=cal.lottery, log=log, name="baseline")
    out = DATA / "baseline"
    out.mkdir(parents=True, exist_ok=True)
    summ = summarize_ss(P, res, "baseline", "baseline", {}, time.time() - t0)
    summ["supply_scales"] = dict(H_bar=H_bar, Hr_bar=Hr_bar, calibration_iterations=cal.iterations)
    save_json(out / "summary.json", summ)
    save_json(out / "supply_scales.json", dict(H_bar=H_bar, Hr_bar=Hr_bar, lottery=res.lottery[0].tolist(), p=res.p, ptil_r=res.ptil_r))
    rec = res.ev.rec
    profiles_by_age(P, rec).to_csv(out / "profiles_by_age.csv", index=False)
    by_gift_quintile(P, rec).to_csv(out / "by_gift_quintile.csv", index=False)
    save_json(out / "mobility.json", mobility_stats(P, rec))
    for k, df in policy_tables(P, res).items():
        df.to_csv(out / f"{k}.csv", index=False)
    thresholds_table(P, res).to_csv(out / "thresholds.csv", index=False)
    save_cross_section(P, rec, out / "cross_section.npz", subsample=1)
    equal_income_paths(P, res).to_csv(out / "equal_income_paths.csv", index=False)
    price_drop_accounting(P, rec).to_csv(out / "price_drop_accounting.csv", index=False)
    pd.DataFrame(M.parameter_table(P)).to_csv(DATA / "parameters.csv", index=False)
    C.targets_table(P, summ["moments"]).to_csv(DATA / "calibration_targets.csv", index=False)
    # income process
    pd.DataFrame(dict(z=P.z_grid, pi=P.pi_z)).to_csv(out / "income_grid.csv", index=False)
    pd.DataFrame(P.Pi_z).to_csv(out / "income_transition.csv", index=False)
    log("baseline_done", sec=time.time() - t0, p=res.p, converged=res.converged)
    with open(out / "ss_result.pkl", "wb") as fh:
        pickle.dump(dict(P=P, res=res), fh)
    return P, res, summ


def run_ss_scenario(args) -> dict:
    """Worker: one stationary scenario (GE or PE at baseline prices)."""
    cfg_base, scen, group, scales, base_lottery, base_prices, n_threads = args
    import numba
    numba.set_num_threads(max(1, n_threads))
    t0 = time.time()
    name = scen["name"]
    cfg = apply_overrides(cfg_base, scen.get("overrides"))
    P = M.build_params(cfg)
    P.H_bar, P.Hr_bar = scales["H_bar"], scales["Hr_bar"]
    log = E.RunLog(name, echo=False)
    kind = scen["kind"]
    K = P.K_estate
    lot0 = (np.array(base_lottery)[:K], np.full(K, 1.0 / K)) if len(base_lottery) >= K else None
    try:
        if kind == "ss_pe":
            res = E.solve_stationary(P, p0=base_prices["p"], ptil0=base_prices["ptil_r"], lottery0=lot0, fixed_prices=True,
                                     log=log, name=name)
        else:
            res = E.solve_stationary(P, p0=base_prices["p"], ptil0=base_prices["ptil_r"], lottery0=lot0, log=log, name=name)
        summ = summarize_ss(P, res, name, group, scen, time.time() - t0)
        out = DATA / group
        out.mkdir(parents=True, exist_ok=True)
        save_json(out / f"{name}.json", summ)
        profiles_by_age(P, res.ev.rec).to_csv(out / f"{name}_profiles.csv", index=False)
        by_gift_quintile(P, res.ev.rec).to_csv(out / f"{name}_by_gift_quintile.csv", index=False)
        thresholds_table(P, res).to_csv(out / f"{name}_thresholds.csv", index=False)
        if name.startswith("theta") or name.startswith("eps") or name in ("bequests_off", "noseg"):
            save_cross_section(P, res.ev.rec, out / f"{name}_cross_section.npz", subsample=4)
        # welfare vs baseline (experiment 9) for GE scenarios
        with open(out / f"{name}_ss.pkl", "wb") as fh:
            pickle.dump(dict(P=P, core=E.ss_core(res), pols0=res.ev.pols[0], lottery=res.lottery, rebate=res.rebate,
                             p=res.p, ptil_r=res.ptil_r), fh)
        log("scenario_done", sec=time.time() - t0, converged=res.converged, p=res.p)
        return dict(name=name, group=group, ok=True, converged=bool(res.converged), p=res.p, sec=time.time() - t0,
                    equilibrium=res.equilibrium, residuals=res.residuals.tolist(), bellman=res.bellman)
    except Exception as exc:  # noqa: BLE001 — report computational failures, never hide them
        import traceback
        log("scenario_failed", error=repr(exc), trace=traceback.format_exc())
        return dict(name=name, group=group, ok=False, error=repr(exc), sec=time.time() - t0)


def run_welfare(cfg_base: dict, groups: dict, base_P: M.Params, base_res: E.SSResult):
    """Experiment 9: CEV by gift quintile and z_1 for every converged GE scenario with a saved pickle."""
    out = DATA / "welfare"
    out.mkdir(parents=True, exist_ok=True)
    frames = []
    for group, scens in groups.items():
        for scen in scens:
            if scen["kind"] != "ss" or scen["name"] == "baseline":
                continue
            pkl = DATA / group / f"{scen['name']}_ss.pkl"
            if not pkl.exists():
                continue
            with open(pkl, "rb") as fh:
                d = pickle.load(fh)
            P_alt = d["P"]
            # light-weight alternative result object with age-1 policies only

            class _Alt:
                pass
            alt = _Alt(); alt.rebate = d["rebate"]
            alt.ev = _Alt(); alt.ev.pols = [d["pols0"]]
            if P_alt.n_a != base_P.n_a:
                continue  # grid variants: value functions not comparable on the same grid
            df = welfare_vs_baseline(base_P, base_res, P_alt, alt, scen["name"])
            df.to_csv(out / f"welfare_{scen['name']}.csv", index=False)
            frames.append(df)
    if frames:
        pd.concat(frames).to_csv(out / "welfare_all.csv", index=False)


def run_transition_scenario(cfg_base: dict, scen: dict, scales: dict, base: dict, n_workers: int,
                            warm_from: str | None = None) -> dict:
    """warm_from: name of a converged transition whose (p, p~_r) path initialises the shooting."""
    t0 = time.time()
    name = scen["name"]
    cfg = apply_overrides(cfg_base, scen.get("overrides"))
    shk = scen.get("shock", {})
    cfg["mortgage"]["contract"] = shk.get("contract", "frm")
    P = M.build_params(cfg)
    P.H_bar, P.Hr_bar = scales["H_bar"], scales["Hr_bar"]
    log = E.RunLog(name)
    T = int(P.cfg["transition"]["T"])
    # stationary equilibrium of this parameterisation (baseline unless overrides change the SS)
    ss_overrides = {k: v for k, v in (scen.get("overrides") or {}).items() if not k.startswith("rental.eta")}
    if ss_overrides:
        lot0 = (np.array(base["lottery"]), np.full(P.K_estate, 1.0 / P.K_estate))
        ss = E.solve_stationary(P, p0=base["p"], ptil0=base["ptil_r"], lottery0=lot0, log=log, name=name, diagnostics=False)
    else:
        with open(DATA / "baseline" / "ss_result.pkl", "rb") as fh:
            d = pickle.load(fh)
        ss = d["res"]
        ss.P = P  # same parameters except the contract flag / indexation (SS-neutral)
    kind = shk.get("type", "inflation")
    sh = E.Shock(kind=kind, pi_u=float(P.cfg["transition"]["pi_surprise"]), persistence=int(shk.get("persistence", 1)),
                 contract=shk.get("contract", "frm"), phi_r=float(shk.get("phi_r", 0.0)),
                 dr=M._to_period(float(shk.get("dr_annual", 0.0)), P.period_years))
    core = E.ss_core(ss)
    mode = "static" if scen["kind"] == "static" else "GE"
    pe = E.solve_transition(P, core, sh, T=T, mode="PE", log=log, name=name)
    u0 = None
    if warm_from is not None and (DATA / "inflation" / f"{warm_from}_path.csv").exists():
        wp = pd.read_csv(DATA / "inflation" / f"{warm_from}_path.csv")
        if len(wp) == T:
            # scale the neighbour's path to this experiment's own stationary prices
            with open(DATA / "inflation" / f"{warm_from}.json") as fh:
                wj = json.load(fh)
            u0 = np.concatenate([np.log(ss.p * (1 + np.array(wj["p_path_pct"]) / 100)),
                                 np.log(ss.ptil_r * (wp.p_r.values / wp.p_r.values[-1]))])
    tr = E.solve_transition(P, core, sh, T=T, mode=mode, log=log, name=name, n_workers=n_workers, u0=u0)
    out = DATA / "inflation"
    out.mkdir(parents=True, exist_ok=True)
    path_df = transition_path_table(P, tr, ss)
    path_df.to_csv(out / f"{name}_path.csv", index=False)
    pe_df = transition_path_table(P, pe, ss)
    pe_df.to_csv(out / f"{name}_path_pe.csv", index=False)
    imp = impact_table(P, tr, ss, pe)
    imp.to_csv(out / f"{name}_impact.csv", index=False)
    # value functions at t = 0 come from the converged path: re-solve the household block once
    prices = tr.prices
    pols_by_t = E.solve_households_path(P, prices, core, T, static_expectations=(mode == "static"))
    wel = transition_welfare(P, tr, ss, pols_by_t[0])
    wel.to_csv(out / f"{name}_welfare.csv", index=False)
    # ownership-entry threshold at impact vs stationary (Proposition 9.3)
    thr0 = {j: {iz: v["a_hat"] for iz, v in ownership_threshold(P, pols_by_t[0][j - 1], prices[0], j).items()} for j in (1, 2, 3)}
    thr_ss = {j: {iz: v["a_hat"] for iz, v in ownership_threshold(P, ss.ev.pols[j - 1], ss.ev.prices, j).items()} for j in (1, 2, 3)}
    pe_pols = E.solve_households_path(P, pe.prices, core, T)
    wel_pe = transition_welfare(P, pe, ss, pe_pols[0])
    wel_pe.to_csv(out / f"{name}_welfare_pe.csv", index=False)
    summ = summarize_transition(P, tr, ss, name, scen, path_df, imp, wel, pe)
    summ["threshold_impact"] = thr0
    summ["threshold_ss"] = thr_ss
    summ["pe"] = dict(equilibrium=pe.equilibrium, max_residual=float(np.max(np.abs(pe.residuals))),
                      d_ownership_impact_pp=summ["d_ownership_impact_pe_pp"],
                      welfare_all=float(wel_pe.loc[(wel_pe.group == "all") & (wel_pe.age == "all"), "cev"].iloc[0]))
    summ["time_sec"] = time.time() - t0
    save_json(out / f"{name}.json", summ)
    log("transition_scenario_done", sec=time.time() - t0, converged=tr.converged, p_impact_pct=summ["p_impact_pct"])
    return dict(name=name, group="inflation", ok=True, converged=bool(tr.converged), equilibrium=tr.equilibrium,
                max_residual=float(np.max(np.abs(tr.residuals))), sec=time.time() - t0)


def _transition_job(args) -> dict:
    cfg, scen, scales, base, n_workers = args
    import numba
    numba.set_num_threads(max(2, n_workers))
    try:
        return run_transition_scenario(cfg, scen, scales, base, n_workers=n_workers)
    except Exception as exc:  # noqa: BLE001 — computational failures are reported, never hidden
        import traceback
        return dict(name=scen["name"], group="inflation", ok=False, error=repr(exc), trace=traceback.format_exc())


def run_rent_growth_pe(cfg_base: dict, scen: dict, scales: dict, n_workers: int) -> dict:
    """Experiment 8 (rent growth): anticipated real rent growth g for T periods at fixed house
    prices (PE): households solve with the time-varying rent path; we report the tenure and
    wealth response by age. Labelled PE (prices fixed)."""
    t0 = time.time()
    name = scen["name"]
    cfg = apply_overrides(cfg_base, scen.get("overrides"))
    P = M.build_params(cfg)
    P.H_bar, P.Hr_bar = scales["H_bar"], scales["Hr_bar"]
    log = E.RunLog(name)
    with open(DATA / "baseline" / "ss_result.pkl", "rb") as fh:
        d = pickle.load(fh)
    ss = d["res"]; ss.P = P
    core = E.ss_core(ss)
    T = int(P.cfg["transition"]["T"])
    g = M._to_period(float(scen["shock"]["g_rent_annual"]), P.period_years)
    sh = E.Shock(kind="none")
    paths = E.shock_paths(P, sh, T)
    logp = np.full(T, np.log(ss.p)); logpt = np.full(T, np.log(ss.ptil_r))
    prices = E.build_price_path(P, paths, sh, T, logp, logpt, core)
    for t, pr in enumerate(prices):
        pr.rho = ss.rho * (1 + g) ** min(t, T)
    pols_by_t = E.solve_households_path(P, prices, core, T)
    recs = E.simulate_path(P, pols_by_t, prices, core.panel, T, 1.0)
    rows = []
    for t, rec in enumerate(recs):
        own = rec.owner1(); young = rec.age <= 3
        rows.append(dict(t=t, rho=prices[t].rho, ownership=float(own.mean()), ownership_young=float(own[young].mean()),
                         buy_rate_young=float((rec.choice[young] == 2).mean()), nw_mean=float(rec.net_worth1(P).mean()),
                         rental_demand=rec.rental_demand(P), housing_demand=rec.housing_demand(P)))
    out = DATA / "sensitivity"
    out.mkdir(parents=True, exist_ok=True)
    df = pd.DataFrame(rows)
    df.to_csv(out / f"{name}_path.csv", index=False)
    save_json(out / f"{name}.json", dict(name=name, group="sensitivity", scenario=scen, equilibrium="PE (prices fixed)",
                                         g_rent_period=g, path=df.to_dict(orient="list"), time_sec=time.time() - t0))
    return dict(name=name, group="sensitivity", ok=True, converged=True, equilibrium="PE (prices fixed)", sec=time.time() - t0)


def run_multiplicity(cfg_base: dict, scales: dict, base: dict, n_threads: int = 0) -> dict:
    """Experiment 3/7: scan the owner-market excess demand in p under near-fixed supply (and the
    baseline) with the rental market cleared, lottery fixed at the equilibrium one; count sign
    changes = number of candidate equilibria; also report d(ownership)/dp along the curve."""
    out = DATA / "multiplicity"
    out.mkdir(parents=True, exist_ok=True)
    results = {}
    for tag, ov in [("baseline", {}), ("eps_005", {"supply.eps": 0.05, "supply.eps_r": 0.05}),
                    ("eps_005_theta005", {"supply.eps": 0.05, "supply.eps_r": 0.05, "mortgage.theta": 0.05})]:
        cfg = apply_overrides(cfg_base, ov)
        P = M.build_params(cfg)
        P.H_bar, P.Hr_bar = scales["H_bar"], scales["Hr_bar"]
        log = E.RunLog(f"multiplicity_{tag}", echo=False)
        lot = (np.array(base["lottery"]), np.full(P.K_estate, 1.0 / P.K_estate))
        p_values = np.round(np.arange(0.80, 1.251, 0.025), 3)
        scan = E.price_scan(P, p_values, lot, log)
        df = pd.DataFrame(scan)
        df.to_csv(out / f"price_scan_{tag}.csv", index=False)
        sg = np.sign(df.excess_owner.values)
        n_cross = int(np.sum(sg[1:] != sg[:-1]))
        slope = np.gradient(df.ownership.values, df.p.values)
        results[tag] = dict(n_sign_changes=n_cross, p_grid=p_values.tolist(), excess_owner=df.excess_owner.tolist(),
                            ownership=df.ownership.tolist(), d_ownership_dp=slope.tolist(),
                            monotone_decreasing=bool(np.all(np.diff(df.excess_owner.values) < 0)))
        log("scan_done", n_sign_changes=n_cross)
    save_json(out / "multiplicity_summary.json", results)
    return results


# --------------------------------------------------------------------------------------
# main
# --------------------------------------------------------------------------------------
def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--groups", nargs="*", default=None, help="scenario groups to run (default: all)")
    ap.add_argument("--workers", type=int, default=max(1, min(9, (os.cpu_count() or 4) // 2)))
    ap.add_argument("--transition-workers", type=int, default=2, help="transitions solved concurrently")
    ap.add_argument("--quick", action="store_true", help="smaller panel / fewer scenarios (smoke test)")
    ap.add_argument("--skip-baseline", action="store_true")
    ap.add_argument("--rerun-transition", nargs="*", default=None, help="re-run only these transition scenarios")
    ap.add_argument("--warm-from", default=None, help="converged transition used as warm start for --rerun-transition")
    args = ap.parse_args(argv)
    T0 = time.time()
    DATA.mkdir(parents=True, exist_ok=True); LOGS.mkdir(parents=True, exist_ok=True)
    main_log = E.RunLog("simulate")
    cfg = M.load_config()
    if args.quick:
        cfg["simulation"]["n_dynasties"] = 24000
        cfg["transition"]["T"] = 8
        cfg["transition"]["max_iter"] = 3
    scenarios = M.load_scenarios()["groups"]
    groups = args.groups or list(scenarios.keys())
    diagnostics = []
    if args.rerun_transition:
        with open(DATA / "baseline" / "supply_scales.json") as fh:
            scales = json.load(fh)
        base = dict(p=scales["p"], ptil_r=scales["ptil_r"], lottery=scales["lottery"])
        all_scens = {sc["name"]: sc for g in scenarios.values() for sc in g}
        out = []
        for nm in args.rerun_transition:
            r = run_transition_scenario(cfg, all_scens[nm], scales, base, n_workers=min(16, (os.cpu_count() or 4) - 2),
                                        warm_from=args.warm_from)
            out.append(r)
            main_log("transition_rerun", **{k: v for k, v in r.items() if k != "trace"})
        # refresh diagnostics.json entries
        dpath = DATA / "diagnostics.json"
        if dpath.exists():
            with open(dpath) as fh:
                dj = json.load(fh)
            names = {r["name"] for r in out}
            dj["runs"] = [r for r in dj["runs"] if r.get("name") not in names] + out
            save_json(dpath, dj)
            pd.DataFrame([{k: v for k, v in r.items() if k not in ("bellman", "residuals", "trace")} for r in dj["runs"]]).to_csv(
                DATA / "diagnostics.csv", index=False)
        return
    # ---- baseline
    if not args.skip_baseline:
        P, res, summ = run_baseline(cfg, quick=args.quick)
        diagnostics.append(dict(name="baseline", group="baseline", ok=True, converged=summ["converged"], p=summ["p"],
                                sec=summ["time_sec"], equilibrium=summ["equilibrium"], residuals=summ["residuals"], bellman=summ["bellman"]))
        main_log("baseline_done", sec=time.time() - T0, p=summ["p"], ownership=summ["moments"]["ownership_rate_all"])
    with open(DATA / "baseline" / "supply_scales.json") as fh:
        scales = json.load(fh)
    base = dict(p=scales["p"], ptil_r=scales["ptil_r"], lottery=scales["lottery"])
    # ---- stationary scenarios in parallel
    ss_jobs = []
    for g in groups:
        for scen in scenarios.get(g, []):
            if scen["kind"] in ("ss", "ss_pe") and scen["name"] != "baseline":
                if args.quick and g == "robustness" and scen["name"] not in ("grid_coarse", "sample_small"):
                    continue
                ss_jobs.append((cfg, scen, g, scales, base["lottery"], base, max(1, (os.cpu_count() or 4) // args.workers)))
    if ss_jobs:
        main_log("ss_scenarios_start", n=len(ss_jobs), workers=args.workers)
        ctx = get_context("spawn")
        with ProcessPoolExecutor(max_workers=args.workers, mp_context=ctx) as ex:
            for r in ex.map(run_ss_scenario, ss_jobs):
                diagnostics.append(r)
                main_log("ss_scenario", **{k: v for k, v in r.items() if k not in ("bellman", "residuals")})
    # ---- multiplicity scan
    if any(g in groups for g in ("supply", "theta_sweep")):
        main_log("multiplicity_start")
        run_multiplicity(cfg, scales, base)
    # ---- transitions (a few concurrently; parallel Jacobian inside each)
    tr_jobs = []
    for g in groups:
        for scen in scenarios.get(g, []):
            if scen["kind"] in ("mit", "static"):
                if args.quick and scen["name"] not in ("frm_level", "term5_persistent", "frm_level_static"):
                    continue
                tr_jobs.append((cfg, scen, scales, base, max(2, ((os.cpu_count() or 4) - 2) // max(1, args.transition_workers))))
            elif scen["kind"] == "rent_growth_pe":
                r = run_rent_growth_pe(cfg, scen, scales, args.workers)
                diagnostics.append(r)
    if tr_jobs:
        main_log("transitions_start", n=len(tr_jobs), concurrent=args.transition_workers)
        ctx = get_context("spawn")
        with ProcessPoolExecutor(max_workers=args.transition_workers, mp_context=ctx) as ex:
            for r in ex.map(_transition_job, tr_jobs):
                diagnostics.append(r)
                main_log("transition_done", **{k: v for k, v in r.items() if k != "trace"})
    # ---- welfare (experiment 9) across stationary scenarios
    with open(DATA / "baseline" / "ss_result.pkl", "rb") as fh:
        d = pickle.load(fh)
    main_log("welfare_start")
    run_welfare(cfg, {g: scenarios[g] for g in groups if g in scenarios}, d["P"], d["res"])
    # ---- diagnostics file
    save_json(DATA / "diagnostics.json", dict(runs=diagnostics, total_sec=time.time() - T0, quick=args.quick,
                                              cpu=os.cpu_count(), workers=args.workers))
    pd.DataFrame([{k: v for k, v in r.items() if k not in ("bellman", "residuals", "trace")} for r in diagnostics]).to_csv(
        DATA / "diagnostics.csv", index=False)
    main_log("all_done", sec=time.time() - T0)


if __name__ == "__main__":
    main()