"""WP18 — equilibrium computation.

* Stationary recursive competitive equilibrium (SRCE, spec B.10): unknowns (log p, log p~_r,
  estate lottery, transfer rebate) solved by damped Broyden with a finite-difference Jacobian
  for the price block and an identity block for the lottery/rebate fixed points. Residuals:
  relative excess demand on the owner-occupied and rental markets, lottery change, rebate
  change. Every result carries `equilibrium = "GE"`; `solve_stationary(..., fixed_prices=True)`
  only iterates the lottery at given prices and is labelled "PE (prices fixed)".
* Supply normalisation (`calibrate_supply`): the supply scales H_bar, Hr_bar are set so that
  the baseline clears at p = p_r = 1 (KMV-style normalisation); counterfactuals keep the scales.
* Perfect-foresight transition after an unanticipated (MIT) shock: sequence-space shooting on
  (log p_t, log p~_r,t)_{t<T} with a finite-difference Jacobian computed in parallel and
  Broyden updates; durable stocks H_t = (1 - delta_H) H_{t-1} + delta_H S(p_t).
* Static-expectations variant (Assumption A6 of the appendix): households use the stationary
  continuation values, prices clear period by period.
"""
from __future__ import annotations

import json
import os
import time
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from multiprocessing import get_context
from pathlib import Path
from typing import Any

import numpy as np

from . import model as M
from .distributions import (Panel, PeriodRecord, estate_lottery, pool_records, simulate_period,
                            simulate_stationary, stack_policies)
from .household_solver import AgePolicy, bellman_residuals, ownership_threshold, solve_lifecycle
from .model import Params, Prices, stationary_prices


# --------------------------------------------------------------------------------------
# Logging
# --------------------------------------------------------------------------------------
class RunLog:
    """Structured JSON-lines log in results/logs/<name>.jsonl plus a console echo."""

    def __init__(self, name: str, echo: bool = True):
        self.path = M.RESULTS_DIR / "logs" / f"{name}.jsonl"
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.echo = echo
        self.t0 = time.time()
        self.fh = open(self.path, "a")

    def __call__(self, event: str, **kw):
        rec = dict(t=round(time.time() - self.t0, 2), event=event, **{k: _jsonable(v) for k, v in kw.items()})
        self.fh.write(json.dumps(rec) + "\n")
        self.fh.flush()
        if self.echo:
            print(f"[{rec['t']:8.1f}s] {event}: " + ", ".join(f"{k}={_fmt(v)}" for k, v in kw.items()), flush=True)

    def close(self):
        self.fh.close()


def _jsonable(v):
    if isinstance(v, np.ndarray):
        return v.tolist()
    if isinstance(v, (np.floating, np.integer)):
        return v.item()
    return v


def _fmt(v):
    if isinstance(v, float):
        return f"{v:.4g}"
    if isinstance(v, np.ndarray):
        return np.array2string(v, precision=4, suppress_small=True)
    return str(v)


# --------------------------------------------------------------------------------------
# Supply helpers
# --------------------------------------------------------------------------------------
def rental_price(P: Params, p: float, ptil_r: float) -> float:
    """p_r = sigma p~_r + (1 - sigma) p (spec B.6)."""
    return P.sigma_seg * ptil_r + (1.0 - P.sigma_seg) * p


def excess_demands(P: Params, H_d: float, S_d: float, p: float, ptil_r: float,
                   H_stock: float | None = None, Hr_stock: float | None = None) -> np.ndarray:
    """Relative excess demands (owner-occupied, rental) given demands and prices.

    With stocks given (transition), the stocks replace the stationary supply curves. For
    sigma < 1 a share (1 - sigma) of rental demand is served by converted owner-occupied units.
    """
    sg = P.sigma_seg
    S_own = M.supply_stock(P, p, P.H_bar, P.eps) if H_stock is None else H_stock
    S_rent = M.supply_stock(P, p, P.Hr_bar, P.eps_r) if Hr_stock is None else Hr_stock   # converted margin, priced at p
    if sg >= 1.0:
        S_seg = M.supply_stock(P, ptil_r, P.Hr_bar, P.eps_r) if Hr_stock is None else Hr_stock
        return np.array([(H_d - S_own) / S_own, (S_d - S_seg) / S_seg])
    if sg <= 0.0:
        tot = S_own + S_rent
        return np.array([(H_d + S_d - tot) / tot, np.log(ptil_r / p)])
    S_seg = M.supply_stock(P, ptil_r, sg * P.Hr_bar, P.eps_r) if Hr_stock is None else sg * Hr_stock
    conv = S_own + (1 - sg) * S_rent
    return np.array([(H_d + (1 - sg) * S_d - conv) / conv, (sg * S_d - S_seg) / S_seg])


# --------------------------------------------------------------------------------------
# Stationary evaluation
# --------------------------------------------------------------------------------------
@dataclass
class SSEval:
    pols: list[AgePolicy]
    rec: PeriodRecord
    panel: Panel
    prices: Prices
    H_d: float
    S_d: float
    lottery_new: tuple[np.ndarray, np.ndarray]
    rebate_new: float
    forced_share: float
    clip_share: float


def evaluate_ss(P: Params, p: float, ptil_r: float, lottery, rebate: float, panel0: Panel | None = None,
                burn: int | None = None) -> SSEval:
    p_r = rental_price(P, p, ptil_r)
    pr = stationary_prices(P, p, p_r, rebate)
    pols = solve_lifecycle(P, pr, lottery)
    recs, panel = simulate_stationary(P, pols, pr, panel0=panel0, burn=burn)
    rec = pool_records(recs)
    lot_new = estate_lottery(P, rec)
    # transfer tax revenue per age-1 household: tau_b * (mean gross gift at j_gift + mean gross estate at J)
    g_mean = float(rec.gift[rec.age == P.j_gift].mean()) if (rec.age == P.j_gift).any() else 0.0
    e_mean = float(rec.estate[rec.age == P.J].mean()) if (rec.age == P.J).any() else 0.0
    rebate_new = P.tau_b * (g_mean + e_mean)
    return SSEval(pols=pols, rec=rec, panel=panel, prices=pr, H_d=rec.housing_demand(P), S_d=rec.rental_demand(P),
                  lottery_new=lot_new, rebate_new=rebate_new,
                  forced_share=float((rec.flag == 1).mean()), clip_share=float((rec.flag == 2).mean()))


@dataclass
class SSResult:
    P: Params
    p: float
    ptil_r: float
    p_r: float
    rho: float
    lottery: tuple[np.ndarray, np.ndarray]
    rebate: float
    ev: SSEval
    residuals: np.ndarray
    converged: bool
    iterations: int
    history: list[dict]
    equilibrium: str
    time_sec: float
    bellman: dict = field(default_factory=dict)
    thresholds: dict = field(default_factory=dict)

    @property
    def H(self) -> float:
        return self.ev.H_d

    @property
    def Hr(self) -> float:
        return self.ev.S_d


@dataclass
class SSCore:
    """Light-weight stationary-equilibrium data needed by the transition solver (picklable)."""
    p: float
    ptil_r: float
    p_r: float
    rebate: float
    lottery: tuple[np.ndarray, np.ndarray]
    V_ss: list[tuple[np.ndarray, np.ndarray]]
    panel: Panel
    H: float
    Hr: float
    equilibrium: str


def ss_core(res: "SSResult") -> SSCore:
    P = res.P
    return SSCore(p=res.p, ptil_r=res.ptil_r, p_r=res.p_r, rebate=res.rebate, lottery=res.lottery,
                  V_ss=[(pol.V_own, pol.V_rent) for pol in res.ev.pols], panel=res.ev.panel,
                  H=M.supply_stock(P, res.p, P.H_bar, P.eps),
                  Hr=M.supply_stock(P, res.ptil_r if P.sigma_seg > 0 else res.p, P.Hr_bar, P.eps_r),
                  equilibrium=res.equilibrium)


def _unpack(P: Params, u: np.ndarray, K: int):
    p, ptil = float(np.exp(u[0])), float(np.exp(u[1]))
    pts = np.maximum(u[2:2 + K], 0.0)
    rebate = float(u[2 + K]) if P.tau_b > 0 else 0.0
    return p, ptil, (pts, np.full(K, 1.0 / K)), rebate


def _residual(P: Params, ev: SSEval, p: float, ptil: float, lottery, rebate: float, fixed_prices: bool) -> np.ndarray:
    ex = excess_demands(P, ev.H_d, ev.S_d, p, ptil)
    if fixed_prices:
        ex = np.zeros(2)
    lot_res = ev.lottery_new[0] - lottery[0]
    parts = [ex, lot_res]
    if P.tau_b > 0:
        parts.append(np.array([ev.rebate_new - rebate]))
    return np.concatenate(parts)


def illinois(fun, x0: float, slope: float, tol: float, max_iter: int = 14, x_step: float = 0.04,
             log=None, label: str = "") -> tuple[float, float, Any, int]:
    """Root of a decreasing scalar function in log-price space.

    Secant step from the initial slope guess, then bracketing and Illinois regula falsi once a
    sign change is found. Returns (x, f(x), payload, n_evals). `fun` returns (f, payload).
    """
    xa, (fa, pa) = x0, fun(x0)
    n = 1
    if abs(fa) < tol:
        return xa, fa, pa, n
    # second point in the direction that should reduce f
    xb = xa - np.clip(fa / slope, -x_step, x_step)
    fb, pb = fun(xb); n += 1
    lo, flo, plo = (xa, fa, pa) if fa > 0 else (xb, fb, pb)
    hi, fhi, phi = (xb, fb, pb) if fa > 0 else (xa, fa, pa)
    # expand until bracketed (f(lo) > 0 > f(hi) with lo < hi for a decreasing function)
    k = 0
    while not (flo > 0 > fhi) and k < 8:
        k += 1
        if abs(fb) < tol:
            return xb, fb, pb, n
        slope_est = (fb - fa) / (xb - xa) if abs(xb - xa) > 1e-10 and (fb - fa) / (xb - xa) < 0 else slope
        xn = xb - np.clip(fb / slope_est, -2 * x_step, 2 * x_step)
        xa, fa, pa = xb, fb, pb
        xb, (fb, pb) = xn, fun(xn); n += 1
        if fb > 0:
            lo, flo, plo = xb, fb, pb
        else:
            hi, fhi, phi = xb, fb, pb
    best = (xb, fb, pb) if abs(fb) < abs(fa) else (xa, fa, pa)
    if not (flo > 0 > fhi):
        return best[0], best[1], best[2], n
    side = 0
    for _ in range(max_iter):
        if abs(best[1]) < tol:
            break
        xm = (lo * fhi - hi * flo) / (fhi - flo)
        xm = min(max(xm, lo + 0.05 * (hi - lo)), hi - 0.05 * (hi - lo))
        fm, pm = fun(xm); n += 1
        if abs(fm) < abs(best[1]):
            best = (xm, fm, pm)
        if fm > 0:
            lo, flo, plo = xm, fm, pm
            if side == 1:
                fhi *= 0.5
            side = 1
        else:
            hi, fhi, phi = xm, fm, pm
            if side == -1:
                flo *= 0.5
            side = -1
        if hi - lo < 1e-6:
            break
    return best[0], best[1], best[2], n


def clear_prices(P: Params, p0: float, ptil0: float, lottery, rebate: float, tol: float, log: RunLog | None,
                 panel0: Panel | None = None, x_step: float = 0.03) -> tuple[float, float, SSEval, np.ndarray, int]:
    """Nested market clearing: rental market (inner, in p~_r) within the owner market (outer, in p)."""
    n_evals = [0]
    ptil_state = [ptil0]

    def rental_cleared(logp):
        p = float(np.exp(logp))
        if P.sigma_seg <= 0.0:
            ev = evaluate_ss(P, p, p, lottery, rebate, panel0=panel0); n_evals[0] += 1
            ex = excess_demands(P, ev.H_d, ev.S_d, p, p)
            return ex[0], (p, ev, ex)

        def inner(logpt):
            pt = float(np.exp(logpt))
            ev = evaluate_ss(P, p, pt, lottery, rebate, panel0=panel0); n_evals[0] += 1
            ex = excess_demands(P, ev.H_d, ev.S_d, p, pt)
            return ex[1], (pt, ev, ex)
        # slope guess: rental excess falls with ptil through supply (eps_r) and demand (~ -10)
        x, fx, (pt, ev, ex), _ = illinois(inner, np.log(ptil_state[0]), -(15.0 + P.eps_r), tol * 0.6, x_step=x_step)
        ptil_state[0] = pt
        return ex[0], (pt, ev, ex)

    x, fx, (pt, ev, ex), _ = illinois(rental_cleared, np.log(p0), -(5.0 + P.eps), tol, x_step=x_step)
    p = float(np.exp(x))
    if log:
        log("clear_prices", p=p, ptil_r=pt, excess=ex, n_evals=n_evals[0])
    return p, pt, ev, ex, n_evals[0]


def clear_prices_broyden(P: Params, p0: float, ptil0: float, lottery, rebate: float, tol: float, log: RunLog | None,
                         J0: np.ndarray | None = None, panel0: Panel | None = None, max_iter: int = 8):
    """2-D damped Broyden on (log p, log p~_r) from a good starting point (later lottery rounds).
    Returns (p, ptil, ev, ex, n_evals, J, success)."""
    n_evals = 0

    def F(u):
        nonlocal n_evals
        p, pt = float(np.exp(u[0])), float(np.exp(u[1]))
        if P.sigma_seg <= 0.0:
            pt = p
        ev = evaluate_ss(P, p, pt, lottery, rebate, panel0=panel0); n_evals += 1
        ex = excess_demands(P, ev.H_d, ev.S_d, p, pt)
        if P.sigma_seg <= 0.0:
            ex = np.array([ex[0], 0.0])
        return ex, ev

    u = np.array([np.log(p0), np.log(ptil0)])
    f, ev = F(u)
    if np.max(np.abs(f)) < tol:
        return float(np.exp(u[0])), float(np.exp(u[1])), ev, f, n_evals, J0, True
    if J0 is None:
        J0 = np.eye(2)
        for k in range(2 if P.sigma_seg > 0 else 1):
            du = np.zeros(2); du[k] = 0.01
            fk, _ = F(u + du)
            J0[:, k] = (fk - f) / 0.01
        if P.sigma_seg <= 0.0:
            J0[1, 1] = 1.0; J0[0, 1] = 0.0; J0[1, 0] = 0.0
    Jm = J0.copy()
    best = (float(np.max(np.abs(f))), u.copy(), f.copy(), ev)
    fails = 0
    for _ in range(max_iter):
        try:
            step = -np.linalg.solve(Jm, f)
        except np.linalg.LinAlgError:
            break
        step = np.clip(step, -0.05, 0.05)
        if P.sigma_seg <= 0.0:
            step[1] = step[0]
        lam = 1.0
        accepted = False
        for _ in range(3):
            u_new = u + lam * step
            f_new, ev_new = F(u_new)
            if np.linalg.norm(f_new) < np.linalg.norm(f):
                accepted = True
                break
            lam *= 0.5
        if not accepted:
            fails += 1
            if fails >= 2:
                break
            continue
        sv = u_new - u; yv = f_new - f
        if float(sv @ sv) > 1e-16:
            Jm = Jm + np.outer(yv - Jm @ sv, sv) / float(sv @ sv)
        u, f, ev = u_new, f_new, ev_new
        if float(np.max(np.abs(f))) < best[0]:
            best = (float(np.max(np.abs(f))), u.copy(), f.copy(), ev)
        if np.max(np.abs(f)) < tol:
            break
    _, u, f, ev = best
    ok = bool(np.max(np.abs(f)) < tol)
    if log:
        log("clear_prices_broyden", p=float(np.exp(u[0])), ptil_r=float(np.exp(u[1])), excess=f, n_evals=n_evals, success=ok)
    return float(np.exp(u[0])), float(np.exp(u[1])), ev, f, n_evals, Jm, ok


def solve_stationary(P: Params, p0: float = 1.0, ptil0: float = 1.0, lottery0=None, rebate0: float = 0.0,
                     fixed_prices: bool = False, log: RunLog | None = None, name: str = "ss",
                     diagnostics: bool = True, tol: float | None = None, max_iter: int | None = None,
                     panel0: Panel | None = None) -> SSResult:
    """Stationary equilibrium: nested price clearing (Illinois regula falsi on each market) inside
    a damped fixed-point iteration on the estate lottery and the transfer rebate.

    fixed_prices=True keeps (p0, ptil0) and only iterates the lottery/rebate fixed point: the
    result is labelled "PE (prices fixed)".
    """
    t0 = time.time()
    cfg_eq = P.cfg["equilibrium"]
    tol = float(cfg_eq["tol_excess"]) if tol is None else tol
    tol_lot = float(cfg_eq["tol_lottery"])
    max_outer = int(cfg_eq.get("max_outer", 10)) if max_iter is None else max_iter
    damp = float(cfg_eq["damping"])
    K = P.K_estate
    if lottery0 is None:
        lottery0 = (np.array([0.0, 0.0, 0.05, 0.4, 1.0])[:K] if K == 5 else np.linspace(0, 1, K), np.full(K, 1.0 / K))
    log = log or RunLog(name)
    lottery, rebate = (np.array(lottery0[0], dtype=float), lottery0[1]), rebate0
    p, ptil = p0, ptil0
    history = []
    converged = False
    total_evals = 0
    it = 0
    ev = None
    Jm = None
    for it in range(1, max_outer + 1):
        if fixed_prices:
            ev = evaluate_ss(P, p, ptil, lottery, rebate, panel0=panel0); total_evals += 1
            ex = np.zeros(2)
        else:
            # nested regula falsi (robust to the Monte-Carlo noise in excess-demand differences);
            # later rounds start from the previous prices with small initial steps
            p, ptil, ev, ex, n_ev = clear_prices(P, p, ptil, lottery, rebate, tol, log, panel0=panel0,
                                                 x_step=0.03 if it == 1 else 0.008)
            total_evals += n_ev
        lot_new = ev.lottery_new[0]
        lot_change = float(np.max(np.abs(lot_new - lottery[0])) / max(1.0, float(np.max(lot_new))))
        reb_change = float(abs(ev.rebate_new - rebate))
        history.append(dict(iter=it, excess_owner=float(ex[0]), excess_rental=float(ex[1]), lottery_change=lot_change,
                            rebate_change=reb_change, p=p, ptil_r=ptil, evals=total_evals))
        log("ss_outer", iter=it, p=p, ptil_r=ptil, excess=ex, lottery=lot_new, lottery_change=lot_change,
            rebate=ev.rebate_new, rebate_change=reb_change, evals=total_evals)
        if lot_change < tol_lot and reb_change < tol_lot and float(np.max(np.abs(ex))) < tol:
            converged = True
            break
        lottery = ((1.0 - damp) * lottery[0] + damp * lot_new, lottery[1])
        rebate = (1.0 - damp) * rebate + damp * ev.rebate_new
    f = np.concatenate([ex, lot_new - lottery[0], [ev.rebate_new - rebate] if P.tau_b > 0 else []])
    lot = (lot_new, lottery[1]) if converged else lottery
    res = SSResult(P=P, p=p, ptil_r=ptil, p_r=rental_price(P, p, ptil), rho=ev.prices.rho, lottery=lot, rebate=rebate,
                   ev=ev, residuals=f, converged=converged, iterations=it, history=history,
                   equilibrium="PE (prices fixed)" if fixed_prices else "GE", time_sec=time.time() - t0)
    if diagnostics:
        res.bellman = bellman_residuals(P, ev.pols, ev.prices, lot, n_points=400)
        res.thresholds = {j: ownership_threshold(P, ev.pols[j - 1], ev.prices, j) for j in (1, 2, 3)}
        log("ss_diagnostics", converged=converged, bellman=res.bellman, forced=ev.forced_share, clipped=ev.clip_share)
    log("ss_done", p=p, p_r=res.p_r, rho=res.rho, converged=converged, outer_iterations=it, evals=total_evals,
        resid=f, sec=res.time_sec)
    return res


def lottery_pts(u, K):
    return u[2:2 + K]


def calibrate_supply(P: Params, log: RunLog | None = None) -> tuple[float, float, SSResult]:
    """Set H_bar, Hr_bar so that the economy clears at p = p~_r = p_ref (lottery converged)."""
    P.H_bar, P.Hr_bar = 1.0, 1.0   # placeholders, prices fixed
    res = solve_stationary(P, p0=P.p_ref, ptil0=P.p_ref, fixed_prices=True, log=log, name="calibrate_supply",
                           diagnostics=False)
    P.H_bar, P.Hr_bar = res.ev.H_d, res.ev.S_d
    return P.H_bar, P.Hr_bar, res


def price_scan(P: Params, p_values: np.ndarray, lottery, log: RunLog | None = None) -> list[dict]:
    """Owner-market excess demand along a grid of p, with p~_r cleared (1-D secant) and the
    lottery held at the equilibrium one: used to count equilibria (multiplicity search)."""
    out = []
    for p in p_values:
        ptil = p
        f_prev = None; x_prev = None
        for it in range(12):
            ev = evaluate_ss(P, float(p), float(ptil), lottery, 0.0)
            ex = excess_demands(P, ev.H_d, ev.S_d, float(p), float(ptil))
            if abs(ex[1]) < 2e-4 or P.sigma_seg <= 0:
                break
            x = np.log(ptil)
            if f_prev is None:
                x_new = x - ex[1] / (-(1.0 + P.eps_r))   # elasticity guess
            else:
                slope = (ex[1] - f_prev) / (x - x_prev) if abs(x - x_prev) > 1e-9 else -(1 + P.eps_r)
                x_new = x - ex[1] / slope if abs(slope) > 1e-6 else x
            x_prev, f_prev = x, ex[1]
            ptil = float(np.exp(np.clip(x_new, x - 0.3, x + 0.3)))
        own_share = float(ev.rec.owner1().mean())
        out.append(dict(p=float(p), ptil_r=float(ptil), excess_owner=float(ex[0]), excess_rental=float(ex[1]),
                        H_d=ev.H_d, S_d=ev.S_d, ownership=own_share))
        if log:
            log("price_scan", p=float(p), ptil_r=float(ptil), excess_owner=float(ex[0]), ownership=own_share)
    return out


# --------------------------------------------------------------------------------------
# Transition after an MIT shock
# --------------------------------------------------------------------------------------
@dataclass
class Shock:
    kind: str = "inflation"        # inflation | real_rate | none
    pi_u: float = 0.10
    persistence: int = 1           # number of periods with elevated inflation / rate
    contract: str = "frm"
    phi_r: float = 0.0             # real-rate response r_t = r + phi_r * pi_u in shock periods
    dr: float = 0.0                # real-rate shock (period units) for kind == real_rate
    eta_rent: float | None = None
    eta_wage: float | None = None


def shock_paths(P: Params, sh: Shock, T: int) -> dict:
    """Exogenous paths: inflation pi_t (between t and t+1 for the deflation of balances, i.e.
    pi_next at t), real gross rate R_t, nominal rate i_t, surprise at t = 0."""
    pibar, R = P.pibar, P.R
    pi_real = np.full(T + 2, pibar)        # realised inflation INTO period t (deflates balances from t-1 to t)
    Rt = np.full(T + 2, R)
    if sh.kind == "inflation":
        for k in range(sh.persistence):
            pi_real[k] = (1.0 + pibar) * (1.0 + sh.pi_u) - 1.0
        if sh.phi_r != 0.0:
            for k in range(sh.persistence):
                Rt[k] = R * (1.0 + sh.phi_r * sh.pi_u) if k == 0 else R * (1.0 + sh.phi_r * sh.pi_u)
    elif sh.kind == "real_rate":
        for k in range(sh.persistence):
            Rt[k] = R * (1.0 + sh.dr)
    pi_u = sh.pi_u if sh.kind == "inflation" else 0.0
    # nominal rate quoted at t for (t, t+1): i_t = R_t (1 + pi_{t+1}) - 1 (+ spread)
    i_new = Rt[:T + 1] * (1.0 + pi_real[1:T + 2]) - 1.0 + P.spread
    # realised real deposit return at t: (1 + i_{t-1}) / (1 + pi_t); at t = 0 the pre-shock i* applies
    R_a = np.empty(T + 1)
    R_a[0] = (1.0 + P.i_nom) / (1.0 + pi_real[0]) * (1.0 + P.dep_premium)
    R_a[1:] = (1.0 + i_new[:T]) / (1.0 + pi_real[1:T + 1]) * (1.0 + P.dep_premium)
    reset, within = M.contract_flags(sh.contract)
    balance_factor = (1.0 + pibar) / (1.0 + pi_real[0]) if not within else 1.0
    return dict(pi_real=pi_real, Rt=Rt, i_new=i_new, R_a=R_a, pi_u=pi_u, reset=reset, within=within,
                balance_factor=balance_factor)


def frm_rate_by_n(P: Params, paths: dict, t: int) -> np.ndarray:
    """Contract rate of an FRM with n remaining payments at date t: the rate at origination
    t0 = t - (N - n), equal to the stationary rate if t0 < 0."""
    out = np.empty(P.N)
    for n in range(1, P.N + 1):
        t0 = t - (P.N - n)
        out[n - 1] = P.i_nom if t0 < 0 else float(paths["i_new"][t0])
    return out


def build_price_path(P: Params, paths: dict, sh: Shock, T: int, logp: np.ndarray, logptil: np.ndarray,
                     ss: SSCore) -> list[Prices]:
    """Prices objects for t = 0..T (t = T uses stationary prices as the terminal condition)."""
    p = np.concatenate([np.exp(logp), [ss.p]])
    ptil = np.concatenate([np.exp(logptil), [ss.ptil_r]])
    p_r = np.array([rental_price(P, p[t], ptil[t]) for t in range(T + 1)])
    eta_rent = P.eta_rent if sh.eta_rent is None else sh.eta_rent
    eta_wage = P.eta_wage if sh.eta_wage is None else sh.eta_wage
    out = []
    for t in range(T + 1):
        p_r_next = p_r[t + 1] if t + 1 <= T else ss.p_r
        rho = M.rent_from_price(P, p_r[t], p_r_next, float(paths["Rt"][t]))
        pi_u = float(paths["pi_u"]) if (t == 0 and sh.kind == "inflation") else 0.0
        out.append(Prices(p=float(p[t]), p_r=float(p_r[t]), rho=rho, R_a=float(paths["R_a"][t]), R_next=float(paths["Rt"][t]),
                          i_new=float(paths["i_new"][t]), i_q=float(paths["i_new"][t]) + P.stress,
                          pi_next=float(paths["pi_real"][t + 1]), pi_u=pi_u,
                          i_c_by_n=frm_rate_by_n(P, paths, t), reset=bool(paths["reset"]),
                          wage_factor=(1.0 + pi_u) ** (-(1.0 - eta_wage)) if pi_u else 1.0,
                          rent_rigid_factor=(1.0 + pi_u) ** (-(1.0 - eta_rent)) if pi_u else 1.0,
                          p_next=float(p[t + 1]) if t + 1 <= T else ss.p, transfer_rebate=ss.rebate))
    return out


def solve_households_path(P: Params, prices: list[Prices], ss: SSCore, T: int,
                          static_expectations: bool = False) -> list[list[AgePolicy]]:
    """Backward solve over dates t = T-1..0; pols_by_t[t][j-1]. Terminal values at T = SS."""
    V_next = list(ss.V_ss)   # V_{j, T} for all j
    pols_by_t: list[list[AgePolicy] | None] = [None] * T
    for t in range(T - 1, -1, -1):
        # continuation for age j at t is V_{j+1, t+1}
        V_next_by_age = [None] * P.J
        for j in range(1, P.J):
            V_next_by_age[j - 1] = V_next[j] if not static_expectations else ss.V_ss[j]
        pols = solve_lifecycle(P, [prices[t]] * P.J, ss.lottery, V_next_by_age=V_next_by_age)
        pols_by_t[t] = pols
        V_next = [(pol.V_own, pol.V_rent) for pol in pols]
    return pols_by_t  # type: ignore[return-value]


def simulate_path(P: Params, pols_by_t, prices: list[Prices], panel0: Panel, T: int, balance_factor: float,
                  seed_offset: int = 100000) -> list[PeriodRecord]:
    panel = panel0.copy()
    recs = []
    for t in range(T):
        S = stack_policies(pols_by_t[t])
        rec = simulate_period(P, S, prices[t], panel, (P.seed, seed_offset + t),
                              balance_factor=balance_factor if t == 0 else 1.0)
        recs.append(rec)
    return recs


def stocks_along_path(P: Params, p: np.ndarray, ptil: np.ndarray, H0: float, Hr0: float) -> tuple[np.ndarray, np.ndarray]:
    T = len(p)
    H = np.empty(T); Hr = np.empty(T)
    h, hr = H0, Hr0
    for t in range(T):
        h = M.stock_next(P, h, p[t], P.H_bar, P.eps)
        pr_seg = ptil[t] if P.sigma_seg > 0 else p[t]
        hr = M.stock_next(P, hr, pr_seg, P.Hr_bar, P.eps_r) if P.sigma_seg >= 1.0 else \
            M.stock_next(P, hr, pr_seg, P.Hr_bar, P.eps_r)
        H[t], Hr[t] = h, hr
    return H, Hr


def transition_residual(P: Params, ss: SSCore, sh: Shock, paths: dict, T: int, u: np.ndarray,
                        static_expectations: bool = False, return_all: bool = False):
    logp, logptil = u[:T], u[T:2 * T]
    prices = build_price_path(P, paths, sh, T, logp, logptil, ss)
    pols_by_t = solve_households_path(P, prices, ss, T, static_expectations=static_expectations)
    recs = simulate_path(P, pols_by_t, prices, ss.panel, T, float(paths["balance_factor"]))
    p = np.exp(logp); ptil = np.exp(logptil)
    H, Hr = stocks_along_path(P, p, ptil, ss.H, ss.Hr)
    res = np.empty(2 * T)
    for t in range(T):
        ex = excess_demands(P, recs[t].housing_demand(P), recs[t].rental_demand(P), float(p[t]), float(ptil[t]),
                            H_stock=float(H[t]), Hr_stock=float(Hr[t]))
        res[t], res[T + t] = ex
    if return_all:
        return res, dict(prices=prices, pols_by_t=pols_by_t, recs=recs, H=H, Hr=Hr)
    return res


# ---- parallel Jacobian workers -------------------------------------------------------
_CTX: dict[str, Any] = {}


def _init_worker(payload, n_threads: int):
    import numba
    numba.set_num_threads(max(1, n_threads))
    _CTX.update(payload)


def _worker_column(args):
    k, h, u = args
    P, ss, sh, paths, T, static = (_CTX[n] for n in ("P", "ss", "sh", "paths", "T", "static"))
    du = np.zeros_like(u); du[k] = h
    return k, transition_residual(P, ss, sh, paths, T, u + du, static_expectations=static)


def _fd_jacobian_parallel(P, ss, sh, paths, T, u, f0, h, static, n_workers: int, log: RunLog | None):
    n = len(u)
    J = np.empty((n, n))
    payload = dict(P=P, ss=ss, sh=sh, paths=paths, T=T, static=static)
    n_workers = max(1, min(n_workers, n))
    threads = max(1, (os.cpu_count() or 4) // n_workers)
    if n_workers == 1:
        for k in range(n):
            _, fk = _worker_column_local(P, ss, sh, paths, T, static, k, h, u)
            J[:, k] = (fk - f0) / h
        return J
    ctx = get_context("spawn")
    with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx, initializer=_init_worker,
                             initargs=(payload, threads)) as ex:
        for k, fk in ex.map(_worker_column, [(k, h, u) for k in range(n)]):
            J[:, k] = (fk - f0) / h
            if log and k % 8 == 0:
                log("jacobian_column", k=k)
    return J


def _worker_column_local(P, ss, sh, paths, T, static, k, h, u):
    du = np.zeros_like(u); du[k] = h
    return k, transition_residual(P, ss, sh, paths, T, u + du, static_expectations=static)


@dataclass
class TransitionResult:
    P: Params
    shock: Shock
    T: int
    p: np.ndarray
    ptil_r: np.ndarray
    p_r: np.ndarray
    rho: np.ndarray
    H: np.ndarray
    Hr: np.ndarray
    residuals: np.ndarray
    converged: bool
    iterations: int
    recs: list[PeriodRecord]
    prices: list[Prices]
    paths: dict
    equilibrium: str
    time_sec: float
    history: list[dict]


def solve_transition(P: Params, ss_in: "SSResult | SSCore", sh: Shock, T: int | None = None, mode: str = "GE",
                     log: RunLog | None = None, name: str = "transition", n_workers: int | None = None,
                     u0: np.ndarray | None = None) -> TransitionResult:
    """mode: "GE" (perfect foresight, prices clear), "PE" (prices fixed at the SRCE — labelled
    partial equilibrium), "static" (A6: stationary continuation values, prices clear).
    u0: optional warm start, stacked (log p_t, log p~_r,t) for t < T (e.g. the path of a
    converged neighbouring experiment); the Jacobian is then computed at u0."""
    t0 = time.time()
    ss = ss_in if isinstance(ss_in, SSCore) else ss_core(ss_in)
    cfg_tr = P.cfg["transition"]
    T = int(cfg_tr["T"]) if T is None else T
    tol = float(cfg_tr["tol"]); max_iter = int(cfg_tr["max_iter"]); h = float(cfg_tr["jac_step"])
    log = log or RunLog(name)
    paths = shock_paths(P, sh, T)
    static = mode == "static"
    u = np.concatenate([np.full(T, np.log(ss.p)), np.full(T, np.log(ss.ptil_r))])
    if u0 is not None and mode != "PE" and len(u0) == 2 * T:
        u = np.asarray(u0, dtype=float).copy()
        log("transition_warm_start", p0=float(np.exp(u[0])))
    history = []
    f, extra = transition_residual(P, ss, sh, paths, T, u, static_expectations=static, return_all=True)
    log("transition_eval", iter=0, max_resid=float(np.max(np.abs(f))), mode=mode)
    converged = False
    it = 0
    if mode == "PE":
        converged = True
    else:
        # --- damped tatonnement pre-phase for large shocks: move each period's log prices in the
        # direction of its excess demand until the path is within the Jacobian's linear range
        kappa = 0.15
        n_tat = 0
        while float(np.max(np.abs(f))) > 0.01 and n_tat < 80:
            u_try = u + kappa * f
            f_try, extra_try = transition_residual(P, ss, sh, paths, T, u_try, static_expectations=static, return_all=True)
            if np.linalg.norm(f_try) < np.linalg.norm(f):
                u, f, extra = u_try, f_try, extra_try
                kappa = min(kappa * 1.2, 0.5)
            else:
                kappa *= 0.5
                if kappa < 1e-3:
                    break
            n_tat += 1
            log("transition_tatonnement", iter=n_tat, max_resid=float(np.max(np.abs(f))), kappa=kappa, p0=float(np.exp(u[0])))
        n_workers = n_workers or max(1, min(16, (os.cpu_count() or 4) - 2))
        Jm = _fd_jacobian_parallel(P, ss, sh, paths, T, u, f, h, static, n_workers, log)
        log("jacobian_done", n=len(u), sec=time.time() - t0)
        fails = 0
        refreshes = 0
        best = (float(np.max(np.abs(f))), u.copy(), f.copy(), extra)
        for it in range(1, max_iter + 1):
            history.append(dict(iter=it - 1, max_resid=float(np.max(np.abs(f)))))
            if np.max(np.abs(f)) < tol:
                converged = True
                break
            try:
                step = -np.linalg.solve(Jm, f)
            except np.linalg.LinAlgError:
                step = -0.5 * f
            step = np.clip(step, -0.2, 0.2)
            lam = 1.0
            accepted = False
            for _ in range(3):
                u_new = u + lam * step
                f_new, extra_new = transition_residual(P, ss, sh, paths, T, u_new, static_expectations=static, return_all=True)
                if np.linalg.norm(f_new) < np.linalg.norm(f):
                    accepted = True
                    break
                lam *= 0.5
            if accepted:
                s_ = u_new - u; yv = f_new - f
                if float(s_ @ s_) > 1e-16:
                    Jm = Jm + np.outer(yv - Jm @ s_, s_) / float(s_ @ s_)
                u, f, extra = u_new, f_new, extra_new
                fails = 0
            else:
                fails += 1
                if fails >= 2:
                    if refreshes >= 1:
                        log("transition_stall", msg="no improvement after a Jacobian refresh: stopping at the best iterate (noise floor)")
                        break
                    u, f, extra = best[1].copy(), best[2].copy(), best[3]
                    Jm = _fd_jacobian_parallel(P, ss, sh, paths, T, u, f, 1.5 * h, static, n_workers, log)
                    log("jacobian_refresh", iter=it)
                    refreshes += 1
                    fails = 0
            if float(np.max(np.abs(f))) < best[0]:
                best = (float(np.max(np.abs(f))), u.copy(), f.copy(), extra)
            log("transition_eval", iter=it, max_resid=float(np.max(np.abs(f))), lam=lam, accepted=accepted, p0=float(np.exp(u[0])))
        u, f, extra = best[1], best[2], best[3]
        if not converged and np.max(np.abs(f)) < tol:
            converged = True
    p = np.exp(u[:T]); ptil = np.exp(u[T:])
    p_r = np.array([pr.p_r for pr in extra["prices"][:T]]); rho = np.array([pr.rho for pr in extra["prices"][:T]])
    label = {"GE": "GE", "PE": "PE (prices fixed)", "static": "static expectations (A6)"}[mode]
    if mode == "GE" and not converged:
        label = "GE NOT CONVERGED (max residual %.2e)" % float(np.max(np.abs(f)))
    res = TransitionResult(P=P, shock=sh, T=T, p=p, ptil_r=ptil, p_r=p_r, rho=rho, H=extra["H"], Hr=extra["Hr"],
                           residuals=f, converged=converged, iterations=it, recs=extra["recs"], prices=extra["prices"],
                           paths=paths, equilibrium=label, time_sec=time.time() - t0, history=history)
    log("transition_done", mode=mode, converged=converged, max_resid=float(np.max(np.abs(f))), iterations=it,
        sec=res.time_sec, p_path=p[:6], rho_path=rho[:6])
    return res