"""WP18 — Part A benchmark model (three-stage lifecycle, two-period generations).

Implements the closed-form hand-to-mouth (HTM) solution of `research/model_specification.md`
(Part A) and the objects used by the propositions of `paper/appendix_proofs.tex`:

* mortgage schedule (FRM / ARM, level or persistent surprise), real payments D2, D3;
* ownership premium Delta, Lambda = exp(-Delta), entry threshold b_hat (Proposition 1);
* HTM value functions and a general numerical household problem used to *verify*
  that the HTM corner is optimal (Assumption A4) rather than to assume it;
* dynasty transition probabilities, persistence lambda, stationary owner share (Prop. 5);
* stationary general equilibrium with segmented owner / rental markets (Prop. 6);
* leverage statistics (Prop. 4);
* surprise-inflation accounting, FRM vs ARM, level vs persistent (Prop. 8);
* impact general-equilibrium response to a surprise (Prop. 9) and a myopic-expectations
  transition simulation (diagnostic only — not covered by a proposition).

Units: one model period = one stage of life; all flows real (consumption units); the house
price p and rent rho are per unit of housing; nominal principal B_nom = P * (1-theta) * p.
Parameters are illustrative (no calibration claim).
"""
from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Literal

import numpy as np
from scipy import optimize, stats

Contract = Literal["FRM", "ARM"]
Experiment = Literal["level", "persistent"]


# --------------------------------------------------------------------------------------
# Parameters
# --------------------------------------------------------------------------------------
@dataclass(frozen=True)
class Params:
    # preferences
    beta: float = 0.80       # discount factor per stage
    nu: float = 0.30         # warm-glow weight: transfer share nu/(1+nu) of end-of-life wealth
    chi: float = 0.02        # ownership utility premium (utils, stage 1 + 2 combined)
    # incomes (real, per stage); y3 ~ y_lo + (y_hi - y_lo) * Beta(a_y, b_y)
    y1: float = 1.0
    y2: float = 3.0
    y3_lo: float = 2.0
    y3_hi: float = 7.0
    y3_a: float = 2.5
    y3_b: float = 2.5
    # housing finance and costs
    theta: float = 0.30      # minimum down-payment share (max LTV = 0.70)
    kappa_b: float = 0.03    # buying cost (share of price)
    kappa_s: float = 0.05    # selling cost (share of price)
    m: float = 0.05          # maintenance per stage (share of price)
    amort: float = 0.50      # share of principal amortised at stage 2
    # rates and inflation
    r: float = 0.10          # real rate per stage
    pibar: float = 0.10      # steady-state inflation per stage
    # supply curves (per cohort): S(p) = Hbar p^eps ; S_r(rho) = Hbar_r rho^eps_r
    Hbar: float = 0.60
    eps: float = 1.0
    Hbar_r: float = 0.80
    eps_r: float = 1.0
    # indexation of nominal wages / rents to a surprise (1 = full)
    eta_w: float = 1.0
    eta_rho: float = 1.0
    # quadrature nodes for expectations over y3
    n_quad: int = 64

    @property
    def R(self) -> float:
        return 1.0 + self.r

    @property
    def d(self) -> float:
        """Stage-1 cash outlay per unit of price: down payment + buying cost + maintenance."""
        return self.theta + self.kappa_b + self.m


# --------------------------------------------------------------------------------------
# Income distribution G (stage-3 income) — bounded support, smooth density
# --------------------------------------------------------------------------------------
def y3_dist(par: Params):
    return stats.beta(par.y3_a, par.y3_b, loc=par.y3_lo, scale=par.y3_hi - par.y3_lo)


def y3_quad(par: Params):
    """Gauss–Legendre nodes/weights for E[f(y3)] with y3 ~ G on [y_lo, y_hi]."""
    x, w = np.polynomial.legendre.leggauss(par.n_quad)
    lo, hi = par.y3_lo, par.y3_hi
    y = lo + (hi - lo) * (x + 1) / 2
    pdf = y3_dist(par).pdf(y)
    wts = w * (hi - lo) / 2 * pdf
    wts = wts / wts.sum()  # renormalise (quadrature error of the density integral)
    return y, wts


def E_y3(par: Params, f) -> float:
    y, w = y3_quad(par)
    return float(np.sum(w * f(y)))


def G(par: Params, y) -> np.ndarray | float:
    return y3_dist(par).cdf(y)


def g(par: Params, y) -> np.ndarray | float:
    return y3_dist(par).pdf(y)


# --------------------------------------------------------------------------------------
# Mortgage schedule
# --------------------------------------------------------------------------------------
def nominal_rate(par: Params, pi_e: float, r: float | None = None) -> float:
    """Fisher: 1 + i = R (1 + pi_e)."""
    R = par.R if r is None else 1.0 + r
    return R * (1.0 + pi_e) - 1.0


def real_payments(par: Params, p: float, *, pi1: float, pi2: float, pi_e1: float, pi_e2: float,
                  contract: Contract = "FRM", r2: float | None = None,
                  principal_share: float | None = None) -> tuple[float, float]:
    """Real stage-2 and stage-3 payments of a mortgage originated at price p.

    pi1, pi2  : realised inflation over (t, t+1) and (t+1, t+2)
    pi_e1     : inflation expected at origination for (t, t+1)  -> fixes i_1 for all contracts
    pi_e2     : inflation expected at t+1 for (t+1, t+2)        -> ARM reset rate
    r2        : real rate used at the ARM reset (None = par.r)
    """
    B = ((1.0 - par.theta) if principal_share is None else principal_share) * p
    i1 = nominal_rate(par, pi_e1)
    if contract == "FRM":
        i2 = i1
    elif contract == "ARM":
        i2 = nominal_rate(par, pi_e2, r2)
    else:
        raise ValueError(contract)
    D2 = (i1 + par.amort) * B / (1.0 + pi1)
    D3 = (1.0 + i2) * (1.0 - par.amort) * B / ((1.0 + pi1) * (1.0 + pi2))
    return D2, D3


def stationary_schedule(par: Params) -> tuple[float, float, float]:
    """(delta2, delta3, e3): payments per unit of price and net equity share at sale."""
    D2, D3 = real_payments(par, 1.0, pi1=par.pibar, pi2=par.pibar, pi_e1=par.pibar, pi_e2=par.pibar)
    e3 = (1.0 - par.kappa_s) - D3
    return D2, D3, e3


# --------------------------------------------------------------------------------------
# HTM closed forms (Part A, §A.5–A.6)
# --------------------------------------------------------------------------------------
def ownership_premium(par: Params, p: float, rho: float, *, p_sale: float | None = None,
                      rho2: float | None = None, D2: float | None = None, D3: float | None = None) -> float:
    """Delta = chi + beta E[ln c2^O - ln c2^R] + beta^2 (1+nu) E[ln W3^O - ln W3^R] under HTM."""
    p_sale = p if p_sale is None else p_sale
    rho2 = rho if rho2 is None else rho2
    if D2 is None or D3 is None:
        d2, d3, _ = stationary_schedule(par)
        D2 = d2 * p if D2 is None else D2
        D3 = d3 * p if D3 is None else D3
    c2O = par.y2 - D2 - par.m * p
    c2R = par.y2 - rho2
    if c2O <= 0 or c2R <= 0:
        raise ValueError("Assumption A3 violated: non-positive stage-2 consumption")
    eq3 = (1.0 - par.kappa_s) * p_sale - D3
    if par.y3_lo + eq3 <= 0:
        raise ValueError("Assumption A3 violated: owner end-of-life wealth can be negative")
    term3 = E_y3(par, lambda y: np.log(y + eq3) - np.log(y))
    return par.chi + par.beta * (np.log(c2O) - np.log(c2R)) + par.beta**2 * (1 + par.nu) * term3


def threshold(par: Params, p: float, rho: float, **kw) -> float:
    """Ownership-entry threshold b_hat (Proposition 1). Returns -inf-like value if nobody is excluded."""
    Delta = ownership_premium(par, p, rho, **kw)
    d = par.d
    b_min = d * p - par.y1
    if Delta <= 0:
        return np.inf  # nobody owns (A5 fails): owning is dominated at every wealth level
    Lam = np.exp(-Delta)
    if d * p <= rho:
        return b_min
    return (d * p - Lam * rho) / (1.0 - Lam) - par.y1


def value_htm(par: Params, b: float, p: float, rho: float, tenure: Literal["O", "R"], **kw) -> float:
    """Lifetime value of a hand-to-mouth household (n1 = n2 = 0, max mortgage if owner)."""
    K = par.beta**2 * (par.nu * np.log(par.nu) - (1 + par.nu) * np.log(1 + par.nu))
    d2, d3, _ = stationary_schedule(par)
    D2 = kw.get("D2", d2 * p)
    D3 = kw.get("D3", d3 * p)
    p_sale = kw.get("p_sale", p)
    rho2 = kw.get("rho2", rho)
    if tenure == "O":
        c1 = b + par.y1 - par.d * p
        if c1 <= 0:
            return -np.inf
        c2 = par.y2 - D2 - par.m * p
        eq3 = (1 - par.kappa_s) * p_sale - D3
        return (np.log(c1) + par.chi + par.beta * np.log(c2)
                + par.beta**2 * (1 + par.nu) * E_y3(par, lambda y: np.log(y + eq3)) + K)
    c1 = b + par.y1 - rho
    if c1 <= 0:
        return -np.inf
    c2 = par.y2 - rho2
    return (np.log(c1) + par.beta * np.log(c2)
            + par.beta**2 * (1 + par.nu) * E_y3(par, lambda y: np.log(y)) + K)


def htm_slack(par: Params, b: float, p: float, rho: float, tenure: Literal["O", "R"]) -> tuple[float, float]:
    """Euler-inequality slacks (Assumption A4): positive values mean the household would like
    to borrow more at stage 1 and at stage 2 (so n1 = n2 = 0 is the optimum)."""
    d2, d3, e3 = stationary_schedule(par)
    R = par.R
    if tenure == "O":
        c1 = b + par.y1 - par.d * p
        c2 = par.y2 - d2 * p - par.m * p
        EW3inv = E_y3(par, lambda y: 1.0 / (y + e3 * p))
    else:
        c1 = b + par.y1 - rho
        c2 = par.y2 - rho
        EW3inv = E_y3(par, lambda y: 1.0 / y)
    s1 = 1.0 / c1 - par.beta * R * (1.0 / c2)          # stage 1 vs 2 (c2 deterministic)
    s2 = 1.0 / c2 - par.beta * R * (1 + par.nu) * EW3inv  # stage 2 vs 3
    return s1, s2


# --------------------------------------------------------------------------------------
# General (non-HTM) household problem — numerical, used to verify the corner solution
# --------------------------------------------------------------------------------------
def value_general(par: Params, b: float, p: float, rho: float, tenure: Literal["O", "R"],
                  return_policy: bool = False):
    """Owner chooses mortgage share phi in [0, 1-theta], savings n1 >= 0, n2 >= 0.
    Renter chooses n1 >= 0, n2 >= 0. Stage-3 split by Lemma 1. Returns the max value."""
    y, w = y3_quad(par)
    K = par.beta**2 * (par.nu * np.log(par.nu) - (1 + par.nu) * np.log(1 + par.nu))
    R = par.R

    def obj_owner(x):
        phi, n1, n2 = x
        D2, D3 = real_payments(par, p, pi1=par.pibar, pi2=par.pibar, pi_e1=par.pibar, pi_e2=par.pibar,
                               principal_share=phi)
        c1 = b + par.y1 - (par.theta + par.kappa_b + par.m) * p - ((1 - par.theta) - phi) * p - n1
        c2 = par.y2 - D2 - par.m * p + n1 * R - n2
        W3 = y + (1 - par.kappa_s) * p - D3 + n2 * R
        if c1 <= 1e-9 or c2 <= 1e-9 or np.any(W3 <= 1e-9):
            return 1e6
        return -(np.log(c1) + par.chi + par.beta * np.log(c2)
                 + par.beta**2 * (1 + par.nu) * np.sum(w * np.log(W3)) + K)

    def obj_renter(x):
        n1, n2 = x
        c1 = b + par.y1 - rho - n1
        c2 = par.y2 - rho + n1 * R - n2
        W3 = y + n2 * R
        if c1 <= 1e-9 or c2 <= 1e-9 or np.any(W3 <= 1e-9):
            return 1e6
        return -(np.log(c1) + par.beta * np.log(c2)
                 + par.beta**2 * (1 + par.nu) * np.sum(w * np.log(W3)) + K)

    best = None
    if tenure == "O":
        bounds = [(0.0, 1 - par.theta), (0.0, 5.0), (0.0, 5.0)]
        starts = [(1 - par.theta, 0.0, 0.0), (0.5 * (1 - par.theta), 0.1, 0.1), (0.0, 0.0, 0.5)]
        f = obj_owner
    else:
        bounds = [(0.0, 5.0), (0.0, 5.0)]
        starts = [(0.0, 0.0), (0.2, 0.2), (0.0, 1.0)]
        f = obj_renter
    for x0 in starts:
        res = optimize.minimize(f, np.array(x0), bounds=bounds, method="L-BFGS-B")
        if res.fun < 1e5 and (best is None or res.fun < best.fun):
            best = res
    if best is None:
        return (-np.inf, None) if return_policy else -np.inf
    return (-best.fun, best.x) if return_policy else -best.fun


def threshold_general(par: Params, p: float, rho: float) -> float:
    """Threshold from the general (numerically optimised) values; verifies Proposition 1."""
    lo = par.d * p - par.y1 + 1e-6
    hi = lo + 10.0
    f = lambda b: value_general(par, b, p, rho, "O") - value_general(par, b, p, rho, "R")
    if f(hi) < 0:
        return np.inf
    if f(lo) > 0:
        return lo
    return optimize.brentq(f, lo, hi, xtol=1e-8)


# --------------------------------------------------------------------------------------
# Dynasties and persistence (Proposition 5)
# --------------------------------------------------------------------------------------
def transitions(par: Params, p: float, rho: float, b_hat: float | None = None, *,
                eq3: float | None = None) -> dict:
    """q_RO, q_OO, lambda, stationary owner share mu* for given prices.
    eq3 = owner's net housing equity at sale (default e3 p)."""
    if b_hat is None:
        b_hat = threshold(par, p, rho)
    _, _, e3 = stationary_schedule(par)
    eq3 = e3 * p if eq3 is None else eq3
    yR = (1 + par.nu) / par.nu * b_hat
    yO = yR - eq3
    qRO = 1.0 - float(G(par, yR)) if np.isfinite(b_hat) else 0.0
    qOO = 1.0 - float(G(par, yO)) if np.isfinite(b_hat) else 0.0
    lam = qOO - qRO
    mu_star = qRO / (1.0 - lam) if (1.0 - lam) > 0 else 1.0
    half_life = np.log(2) / (-np.log(lam)) if 0 < lam < 1 else (np.inf if lam >= 1 else 0.0)
    return dict(b_hat=b_hat, yR=yR, yO=yO, qRO=qRO, qOO=qOO, lam=lam, mu_star=mu_star,
                half_life=half_life, transfer_gap=par.nu / (1 + par.nu) * eq3)


# --------------------------------------------------------------------------------------
# Supply and stationary equilibrium (Proposition 6)
# --------------------------------------------------------------------------------------
def S(par: Params, p: float) -> float:
    return par.Hbar * p ** par.eps


def S_inv(par: Params, mu: float) -> float:
    return (mu / par.Hbar) ** (1.0 / par.eps)


def Sr(par: Params, rho: float) -> float:
    return par.Hbar_r * rho ** par.eps_r


def Sr_inv(par: Params, share: float) -> float:
    return (share / par.Hbar_r) ** (1.0 / par.eps_r)


def p_upper(par: Params) -> float:
    """Largest price compatible with Assumption A3 (positive stage-2 owner consumption)."""
    d2, _, _ = stationary_schedule(par)
    return 0.999 * par.y2 / (d2 + par.m)


def excess_map(par: Params, mu: float) -> float:
    """f(mu) = mu*(S^-1(mu), S_r^-1(1-mu)) - mu."""
    p = S_inv(par, mu)
    rho = Sr_inv(par, 1.0 - mu)
    return transitions(par, p, rho)["mu_star"] - mu


def stationary_equilibrium(par: Params, mu_lo: float = 1e-4, mu_hi: float = 1 - 1e-4) -> dict:
    """Solve f(mu) = 0 on [mu_lo, mu_hi] (IVT); returns prices, shares, thresholds, diagnostics.
    With a fixed owner stock (eps = 0) the owner share is mu = Hbar and the price solves
    mu*(p, rho) = Hbar (Corollary to Proposition 6)."""
    if par.eps == 0:
        mu = par.Hbar
        rho = Sr_inv(par, 1.0 - mu)
        h = lambda p: transitions(par, p, rho)["mu_star"] - mu
        p = optimize.brentq(h, 1e-3, p_upper(par), xtol=1e-12)
    else:
        mu_hi = min(mu_hi, S(par, p_upper(par)) * (1 - 1e-9))  # keep A3 along the bracket
        f_lo, f_hi = excess_map(par, mu_lo), excess_map(par, mu_hi)
        if f_lo <= 0:
            mu = mu_lo
        elif f_hi >= 0:
            mu = mu_hi
        else:
            mu = optimize.brentq(lambda x: excess_map(par, x), mu_lo, mu_hi, xtol=1e-12)
        p = S_inv(par, mu)
        rho = Sr_inv(par, 1.0 - mu)
    tr = transitions(par, p, rho)
    out = dict(mu=mu, p=p, rho=rho, **tr)
    out["Delta"] = ownership_premium(par, p, rho)
    out["Lambda"] = np.exp(-out["Delta"])
    out["owner_clearing"] = S(par, p) - mu
    out["rental_clearing"] = Sr(par, rho) - (1 - mu)
    out["A5_cash"] = par.d * p - rho          # > 0 : owning costs more cash up front
    return out


def uniqueness_condition(par: Params, p: float, rho: float, h: float = 1e-5) -> float:
    """Condition (C): (1+nu)/nu * d b_hat/dp - e3 >= 0  (b_hat channel dominates ladder channel)."""
    _, _, e3 = stationary_schedule(par)
    db = (threshold(par, p + h, rho) - threshold(par, p - h, rho)) / (2 * h)
    return (1 + par.nu) / par.nu * db - e3


def no_segmentation_equilibrium(par: Params, phi: float = 0.0) -> dict:
    """Proposition 7 benchmark: a single stock with costless conversion and unit demand.
    Total demand per cohort is 1, so p solves S(p) = 1 and does not depend on theta, G, transfers
    or surprises; the rent follows landlord arbitrage rho = p (1 + m + phi - 1/R)."""
    p = S_inv(par, 1.0)
    rho = p * (1.0 + par.m + phi - 1.0 / par.R)
    bh = threshold(par, p, rho)
    tr = transitions(par, p, rho, bh)
    return dict(p=p, rho=rho, b_hat=bh, **{k: v for k, v in tr.items() if k != "b_hat"})


# --------------------------------------------------------------------------------------
# Leverage amplification with costs (Proposition 4)
# --------------------------------------------------------------------------------------
def leverage_stats(par: Params, p: float, rho: float) -> dict:
    d2, d3, e3 = stationary_schedule(par)
    R = par.R
    multiplier = (1 - par.kappa_s) / par.d                     # dW3 per unit of initial cash, per unit price growth
    rhs = R**2 * (1 + par.kappa_b + par.m) + par.m * R - (rho / p) * (R + R**2)
    g_star = rhs / (1 - par.kappa_s) - 1.0                      # break-even real price growth
    x_star = 1.0 - d3 / (1 - par.kappa_s)                       # underwater threshold (price drop)
    return dict(multiplier=multiplier, g_star=g_star, x_star=x_star, e3=e3,
                wealth_gap=e3 * p, transfer_gap=par.nu / (1 + par.nu) * e3 * p)


def owner_vs_renter_investor(par: Params, p: float, rho: float, growth: float) -> float:
    """Terminal wealth of an owner minus that of a renter who invests the cash-flow differences
    at R (same consumption path). Sign flips at g_star (Proposition 4)."""
    d2, d3, _ = stationary_schedule(par)
    R = par.R
    owner = (1 - par.kappa_s) * p * (1 + growth) - d3 * p
    renter = (par.d * p - rho) * R**2 + (d2 * p + par.m * p - rho) * R
    return owner - renter


# --------------------------------------------------------------------------------------
# Surprise inflation accounting (Proposition 8)
# --------------------------------------------------------------------------------------
def surprise_accounting(par: Params, p: float, rho: float, pi_u: float, *, contract: Contract = "FRM",
                        experiment: Experiment = "level", phi_r: float = 0.0,
                        n_savings_age2: float = 0.0, n_savings_age3: float = 0.0) -> dict:
    """Real gains (+) / losses (-) relative to the no-surprise path at fixed real prices,
    for the cohorts aged 2 and 3 when the surprise hits, the lender, landlords and wages.

    n_savings_age2 / age3: nominal deposits held by (non-HTM) households of those ages, to show
    that nominal savers lose. Returns a dict with a zero-sum check on nominal claims."""
    pb = par.pibar
    pi1 = (1 + pb) * (1 + pi_u) - 1
    if experiment == "level":
        pi2 = pb
        pi_e2 = pb
    elif experiment == "persistent":
        pi2 = pi1
        pi_e2 = pi1
    else:
        raise ValueError(experiment)
    r2 = par.r + phi_r * pi_u
    # baseline (no surprise)
    D2_0, D3_0 = real_payments(par, p, pi1=pb, pi2=pb, pi_e1=pb, pi_e2=pb, contract=contract)
    # cohort aged 2 at the shock (bought at t)
    D2_s, D3_s = real_payments(par, p, pi1=pi1, pi2=pi2, pi_e1=pb, pi_e2=pi_e2, contract=contract, r2=r2)
    gain_age2_D2 = D2_0 - D2_s
    gain_age2_D3 = D3_0 - D3_s
    # cohort aged 3 at the shock (bought at t-1): stage-3 nominal payment fixed before the shock
    gain_age3_D3 = D3_0 - D3_0 / (1 + pi_u)
    # nominal savers
    R = par.R
    loss_savers = (n_savings_age2 + n_savings_age3) * R * (1 - 1 / (1 + pi_u))
    # wages (imperfect indexation) — cohort aged 2 and 3 real income at the shock date
    wage_factor = (1 + pi_u) ** (-(1 - par.eta_w))
    wage_loss_age2 = par.y2 * (1 - wage_factor)
    # rents (rigidity): sitting tenants aged 2 pay less; landlords absorb
    rent_factor = (1 + pi_u) ** (-(1 - par.eta_rho))
    renter_gain_age2 = rho * (1 - rent_factor)
    # lender: mirror image of all nominal revaluations on the shock date and (for the age-2 cohort)
    # on the continuation, both contracts
    lender_loss = (gain_age2_D2 + gain_age2_D3 + gain_age3_D3) - loss_savers
    zero_sum = (gain_age2_D2 + gain_age2_D3 + gain_age3_D3 - loss_savers) - lender_loss
    net_owner_age2 = gain_age2_D2 + gain_age2_D3 - wage_loss_age2 - n_savings_age2 * R * (1 - 1 / (1 + pi_u))
    return dict(D2_base=D2_0, D3_base=D3_0, D2_shock=D2_s, D3_shock=D3_s,
                gain_age2_D2=gain_age2_D2, gain_age2_D3=gain_age2_D3, gain_age3_D3=gain_age3_D3,
                loss_savers=loss_savers, wage_loss_age2=wage_loss_age2,
                renter_gain_age2=renter_gain_age2, landlord_loss=renter_gain_age2,
                lender_loss=lender_loss, zero_sum_residual=zero_sum, net_owner_age2=net_owner_age2)


# --------------------------------------------------------------------------------------
# Impact general-equilibrium response (Proposition 9) under static expectations (A6)
# --------------------------------------------------------------------------------------
def impact_response(par: Params, eq: dict, pi_u: float, *, contract: Contract = "FRM") -> dict:
    """Date t+1 equilibrium after a level surprise pi_u, with (i) stage-2 owners not transacting,
    (ii) stage-3 owner-parents' D3 eroded by 1/(1+pi_u), (iii) the stage-1 cohort's threshold
    evaluated at current (p, rho) and stationary future prices (A6).

    Owner market: S(p) = (mu* + mu_{t+1})/2 ; rental: S_r(rho) = 1 - (mu* + mu_{t+1})/2.
    Solved by a root search on mu_{t+1}."""
    mu_s, p_s, rho_s = eq["mu"], eq["p"], eq["rho"]
    d2, d3, _ = stationary_schedule(par)
    D3_parent = d3 * p_s / (1 + pi_u)  # stage-3 parents bought at p_s; payment eroded

    def mu_next_given(mu1, p_fixed=None):
        avg = (mu_s + mu1) / 2
        p = S_inv(par, avg) if p_fixed is None else p_fixed
        rho = Sr_inv(par, 1 - avg)
        # child's threshold: buys at p now, pays D2, D3 on (1-theta) p, sells at p_s (static exp.), rents at rho now, rho_s later
        bh = threshold(par, p, rho, p_sale=p_s, rho2=rho_s, D2=d2 * p, D3=d3 * p)
        eq3_parent = (1 - par.kappa_s) * p - D3_parent    # parent's equity at sale (sells at p)
        tr = transitions(par, p, rho, bh, eq3=eq3_parent)
        mu_new = mu_s * tr["qOO"] + (1 - mu_s) * tr["qRO"]
        return mu_new, p, rho, bh, tr

    if par.eps == 0:
        # fixed owner stock: mu_{t+1} = mu* and the price clears the owner market
        mu1 = mu_s
        hfun = lambda p: mu_next_given(mu1, p_fixed=p)[0] - mu1
        p_sol = optimize.brentq(hfun, 1e-3, p_upper(par), xtol=1e-12)
        mu_new, p, rho, bh, tr = mu_next_given(mu1, p_fixed=p_sol)
    else:
        f = lambda mu1: mu_next_given(mu1)[0] - mu1
        lo, hi = 1e-6, min(1 - 1e-6, 2 * S(par, p_upper(par)) * (1 - 1e-9) - mu_s)
        if f(lo) <= 0:
            mu1 = lo
        elif f(hi) >= 0:
            mu1 = hi
        else:
            mu1 = optimize.brentq(f, lo, hi, xtol=1e-12)
        mu_new, p, rho, bh, tr = mu_next_given(mu1)
    return dict(mu_next=mu1, p=p, rho=rho, b_hat=bh, qRO=tr["qRO"], qOO=tr["qOO"], lam=tr["lam"],
                dp=p - p_s, drho=rho - rho_s, dmu=mu1 - mu_s, db_hat=bh - eq["b_hat"],
                rent_price=rho / p, rent_price_base=rho_s / p_s,
                seller_gain=d3 * p_s * (1 - 1 / (1 + pi_u)) + (1 - par.kappa_s) * (p - p_s),
                new_owner_cash_change=par.nu / (1 + par.nu) * (d3 * p_s * (1 - 1 / (1 + pi_u))
                                                              + (1 - par.kappa_s) * (p - p_s))
                - par.d * (p - p_s))


def simulate_myopic_transition(par: Params, eq: dict, pi_u: float, T: int = 12,
                               experiment: Experiment = "level", contract: Contract = "FRM") -> dict:
    """Diagnostic simulation of the dynasty composition after a surprise at date 1 under the
    Part A structure with per-period supply and static expectations. Tracks, for each cohort,
    the tenure share and the purchase price, so that parents' equity at sale uses the right
    purchase price and payment erosion. Not a proposition: a check that the economy returns to
    the stationary state and that the impact signs match `impact_response`."""
    d2, d3, _ = stationary_schedule(par)
    pb = par.pibar
    # histories are indexed by cohort c = -1, 0, 1, ... at position c + 1
    mu_hist = [eq["mu"], eq["mu"]]      # cohorts born at t=-1, 0
    p_hist = [eq["p"], eq["p"]]
    rho_hist = [eq["rho"], eq["rho"]]
    rows = []
    for t in range(1, T + 1):
        # parents of this cohort: born at t-2 (stage 3 now), position t-1; previous cohort: position t
        p_buy = p_hist[t - 1]
        mu_parents = mu_hist[t - 1]
        # erosion of the parents' D3 and of D2/D3 of the cohort born at t-1 (not needed for mu)
        if t == 1:
            pi1, pi2, pi_e2 = (1 + pb) * (1 + pi_u) - 1, pb, pb
            _, D3_par = real_payments(par, p_buy, pi1=pb, pi2=pi1, pi_e1=pb, pi_e2=pb, contract=contract)
        elif t == 2 and experiment == "persistent":
            pi1 = (1 + pb) * (1 + pi_u) - 1
            _, D3_par = real_payments(par, p_buy, pi1=pi1, pi2=pi1, pi_e1=pb, pi_e2=pi1, contract=contract)
        elif t == 2:
            pi1 = (1 + pb) * (1 + pi_u) - 1
            _, D3_par = real_payments(par, p_buy, pi1=pi1, pi2=pb, pi_e1=pb, pi_e2=pb, contract=contract)
        else:
            D3_par = d3 * p_buy
        mu_prev = mu_hist[t]

        def mu_given(mu1):
            avg = (mu_prev + mu1) / 2
            p = S_inv(par, avg)
            rho = Sr_inv(par, 1 - avg)
            bh = threshold(par, p, rho, p_sale=eq["p"], rho2=eq["rho"], D2=d2 * p, D3=d3 * p)
            eq3_par = (1 - par.kappa_s) * p - D3_par
            tr = transitions(par, p, rho, bh, eq3=eq3_par)
            return mu_parents * tr["qOO"] + (1 - mu_parents) * tr["qRO"], p, rho, bh

        f = lambda x: mu_given(x)[0] - x
        hi = min(1 - 1e-6, 2 * S(par, p_upper(par)) * (1 - 1e-9) - mu_prev)
        mu1 = optimize.brentq(f, 1e-6, hi, xtol=1e-12) if f(1e-6) > 0 > f(hi) else (1e-6 if f(1e-6) <= 0 else hi)
        _, p, rho, bh = mu_given(mu1)
        mu_hist.append(mu1); p_hist.append(p); rho_hist.append(rho)
        rows.append(dict(t=t, mu=mu1, p=p, rho=rho, b_hat=bh))
    return dict(rows=rows, mu=np.array(mu_hist[2:]), p=np.array(p_hist[2:]), rho=np.array(rho_hist[2:]))


# --------------------------------------------------------------------------------------
# Budget identities (used by the tests)
# --------------------------------------------------------------------------------------
def htm_allocation(par: Params, b: float, p: float, rho: float, tenure: Literal["O", "R"], y3: float,
                   *, pi_u1: float = 0.0, pi_u2: float = 0.0, contract: Contract = "FRM",
                   experiment: Experiment = "level", p_sale: float | None = None) -> dict:
    """Full real flows of an HTM household for a given y3 draw; used to check budget identities
    state by state. Returns sources and uses at each stage, the nominal principal path and the
    real debt path (B_nom / P)."""
    pb = par.pibar
    P0 = 1.0
    pi1 = (1 + pb) * (1 + pi_u1) - 1
    if experiment == "persistent" and pi_u1 != 0:
        pi2, pi_e2 = pi1, pi1
    else:
        pi2 = (1 + pb) * (1 + pi_u2) - 1
        pi_e2 = pb
    P1, P2 = P0 * (1 + pi1), P0 * (1 + pi1) * (1 + pi2)
    p_sale = p if p_sale is None else p_sale
    out = dict(P=(P0, P1, P2))
    if tenure == "O":
        B_nom0 = P0 * (1 - par.theta) * p
        i1 = nominal_rate(par, pb)
        i2 = i1 if contract == "FRM" else nominal_rate(par, pi_e2)
        pay2_nom = (i1 + par.amort) * B_nom0
        B_nom1 = B_nom0 * (1 + i1) - pay2_nom          # balance after stage-2 payment = (1-a) B_nom0
        pay3_nom = (1 + i2) * B_nom1
        D2, D3 = pay2_nom / P1, pay3_nom / P2
        c1 = b + par.y1 - par.d * p
        c2 = par.y2 - D2 - par.m * p
        W3 = y3 + (1 - par.kappa_s) * p_sale - D3
        out.update(B_nom=(B_nom0, B_nom1, 0.0), real_debt=(B_nom0 / P0, B_nom1 / P1, 0.0),
                   D2=D2, D3=D3, c1=c1, c2=c2,
                   sources=(b + par.y1 + B_nom0 / P0, par.y2, y3 + (1 - par.kappa_s) * p_sale),
                   uses=(c1 + (1 + par.kappa_b + par.m) * p, c2 + D2 + par.m * p, None))
        # check D2, D3 against real_payments
        D2c, D3c = real_payments(par, p, pi1=pi1, pi2=pi2, pi_e1=pb, pi_e2=pi_e2, contract=contract)
        out["schedule_check"] = (D2 - D2c, D3 - D3c)
    else:
        c1 = b + par.y1 - rho
        c2 = par.y2 - rho
        W3 = y3
        out.update(c1=c1, c2=c2, sources=(b + par.y1, par.y2, y3), uses=(c1 + rho, c2 + rho, None))
    c3 = W3 / (1 + par.nu)
    bq = par.nu * W3 / (1 + par.nu)
    out.update(W3=W3, c3=c3, bequest=bq)
    out["uses"] = (out["uses"][0], out["uses"][1], c3 + bq + (D3 if tenure == "O" else 0.0))
    return out


if __name__ == "__main__":  # quick report
    par = Params()
    eq = stationary_equilibrium(par)
    print("Stationary equilibrium:")
    for k, v in eq.items():
        print(f"  {k:>18s} = {v:.6f}" if isinstance(v, float) else f"  {k:>18s} = {v}")
    print("Uniqueness condition (C) at eq:", uniqueness_condition(par, eq["p"], eq["rho"]))
    print("Leverage:", leverage_stats(par, eq["p"], eq["rho"]))
    for contract in ("FRM", "ARM"):
        for exp_ in ("level", "persistent"):
            acc = surprise_accounting(par, eq["p"], eq["rho"], 0.10, contract=contract, experiment=exp_)
            print(contract, exp_, {k: round(v, 5) for k, v in acc.items() if k.startswith(("gain", "lender"))})
    imp = impact_response(par, eq, 0.10)
    print("Impact:", {k: round(v, 5) for k, v in imp.items()})