"""Tests of the Part A benchmark (`src/toy_model.py`).

Each test is tied to an object of `research/model_specification.md` / `paper/appendix_proofs.tex`:
budget identities and debt evolution (B.11 / A.4), Fisher present-value identity, Proposition 1
(threshold = numerically optimised threshold), Assumption A4 verified on the support of the
stationary transfer distribution, Proposition 2 (signs), Proposition 3 (tilt / neutrality),
Proposition 4 (leverage), Proposition 5 (persistence, chain), Proposition 6 (clearing, condition C),
Corollary (fixed supply), Proposition 7 (no segmentation), Proposition 8 (accounting, FRM vs ARM),
Proposition 9 (impact signs). Parameters are illustrative.
"""
import os
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

import toy_model as tm  # noqa: E402
from toy_model import Params, replace  # noqa: E402

PAR = Params()
EQ = tm.stationary_equilibrium(PAR)
H = 1e-5


def num_deriv(fun, x, h=H):
    return (fun(x + h) - fun(x - h)) / (2 * h)


# ------------------------------------------------------------------ identities
@pytest.mark.parametrize("tenure", ["O", "R"])
@pytest.mark.parametrize("contract", ["FRM", "ARM"])
@pytest.mark.parametrize("pi_u", [0.0, 0.1, -0.05])
@pytest.mark.parametrize("experiment", ["level", "persistent"])
def test_budget_identities_state_by_state(tenure, contract, pi_u, experiment):
    y3_nodes, _ = tm.y3_quad(PAR)
    for y3 in y3_nodes[::8]:
        for b in [EQ["b_hat"] + 0.05, 1.5]:
            a = tm.htm_allocation(PAR, b, EQ["p"], EQ["rho"], tenure, float(y3), pi_u1=pi_u,
                                  contract=contract, experiment=experiment)
            src, use = a["sources"], a["uses"]
            for j in range(3):
                assert src[j] == pytest.approx(use[j], abs=1e-12), (tenure, j)
            assert a["c1"] > 0 and a["c2"] > 0 and a["W3"] > 0  # Assumption A3
            if tenure == "O":
                assert a["schedule_check"] == pytest.approx((0.0, 0.0), abs=1e-12)


def test_debt_evolution_nominal_vs_real():
    a = tm.htm_allocation(PAR, 1.2, EQ["p"], EQ["rho"], "O", 4.0, pi_u1=0.2)
    B0, B1, B2 = a["B_nom"]
    P0, P1, P2 = a["P"]
    # nominal: balance after stage-2 payment = (1 - amort) * B0 (interest paid in full at stage 2)
    assert B1 == pytest.approx((1 - PAR.amort) * B0, abs=1e-12)
    assert B2 == 0.0
    # real debt = B_nom / CPI; the surprise lowers real debt by 1/(1+pi_u) relative to the no-surprise path
    a0 = tm.htm_allocation(PAR, 1.2, EQ["p"], EQ["rho"], "O", 4.0, pi_u1=0.0)
    assert a["real_debt"][1] == pytest.approx(a0["real_debt"][1] / 1.2, rel=1e-12)
    # nominal house-price index (P p) differs from the CPI path
    assert P1 * EQ["p"] != pytest.approx(P1)


def test_fisher_present_value_identity():
    d2, d3, _ = tm.stationary_schedule(PAR)
    assert d2 + d3 / PAR.R == pytest.approx((1 - PAR.theta) * PAR.R, rel=1e-12)
    for pibar in [0.0, 0.1, 0.5]:
        d2, d3, _ = tm.stationary_schedule(replace(PAR, pibar=pibar))
        assert d2 + d3 / PAR.R == pytest.approx((1 - PAR.theta) * PAR.R, rel=1e-12)


def test_stage3_split_lemma1():
    a = tm.htm_allocation(PAR, 1.0, EQ["p"], EQ["rho"], "R", 4.0)
    assert a["c3"] == pytest.approx(a["W3"] / (1 + PAR.nu))
    assert a["bequest"] == pytest.approx(PAR.nu * a["W3"] / (1 + PAR.nu))


# ------------------------------------------------------------------ Proposition 1 & A4
def test_threshold_equals_numerically_optimised_threshold():
    b_num = tm.threshold_general(PAR, EQ["p"], EQ["rho"])
    assert b_num == pytest.approx(EQ["b_hat"], abs=1e-6)
    assert EQ["b_hat"] > PAR.d * EQ["p"] - PAR.y1  # strictly above the feasibility bound
    assert EQ["A5_cash"] > 0 and EQ["Delta"] > 0  # Assumption A5 holds at the equilibrium


@pytest.mark.parametrize("tenure", ["O", "R"])
def test_htm_corner_is_optimal_on_transfer_support(tenure):
    """Assumption A4 verified (not assumed): Euler slacks positive at the threshold and at the
    largest transfer of the stationary distribution, and the general optimiser returns the corner."""
    _, _, e3 = tm.stationary_schedule(PAR)
    b_max = PAR.nu / (1 + PAR.nu) * (PAR.y3_hi + e3 * EQ["p"])
    for b in [EQ["b_hat"], 0.5 * (EQ["b_hat"] + b_max), b_max]:
        s1, s2 = tm.htm_slack(PAR, b, EQ["p"], EQ["rho"], tenure)
        assert s1 > 0 and s2 > 0
        v_gen, pol = tm.value_general(PAR, b, EQ["p"], EQ["rho"], tenure, return_policy=True)
        v_htm = tm.value_htm(PAR, b, EQ["p"], EQ["rho"], tenure)
        assert v_gen == pytest.approx(v_htm, abs=1e-6)
        if tenure == "O":
            assert pol[0] == pytest.approx(1 - PAR.theta, abs=1e-6)  # max LTV
            assert pol[1] == pytest.approx(0.0, abs=1e-6) and pol[2] == pytest.approx(0.0, abs=1e-6)
        else:
            assert np.allclose(pol, 0.0, atol=1e-6)


def test_single_crossing_general_values():
    """V^O - V^R crosses zero once on a grid (checked with the general optimiser)."""
    grid = np.linspace(PAR.d * EQ["p"] - PAR.y1 + 0.02, 3.0, 25)
    diff = np.array([tm.value_general(PAR, b, EQ["p"], EQ["rho"], "O") - tm.value_general(PAR, b, EQ["p"], EQ["rho"], "R")
                     for b in grid])
    signs = np.sign(diff)
    assert np.sum(signs[1:] != signs[:-1]) == 1
    assert diff[0] < 0 < diff[-1]


# ------------------------------------------------------------------ Proposition 2 & 3
def test_threshold_comparative_statics_signs():
    p, rho = EQ["p"], EQ["rho"]
    assert num_deriv(lambda x: tm.threshold(PAR, x, rho), p) > 0
    assert num_deriv(lambda x: tm.threshold(PAR, p, x), rho) < 0
    for name, sign in [("theta", +1), ("kappa_b", +1), ("kappa_s", +1), ("m", +1), ("chi", -1), ("pibar", +1)]:
        v = getattr(PAR, name)
        der = num_deriv(lambda x: tm.threshold(replace(PAR, **{name: x}), p, rho), v)
        assert np.sign(der) == sign, name
    assert num_deriv(lambda x: tm.threshold(replace(PAR, y1=x), p, rho), PAR.y1) == pytest.approx(-1.0, abs=1e-6)


def test_anticipated_inflation_neutral_in_pv_but_tilts_payments():
    d2a, d3a, _ = tm.stationary_schedule(replace(PAR, pibar=0.0))
    d2b, d3b, _ = tm.stationary_schedule(replace(PAR, pibar=0.3))
    assert d2b > d2a and d3b < d3a                     # tilt
    assert d2a + d3a / PAR.R == pytest.approx(d2b + d3b / PAR.R)  # PV neutrality (Fisher)


# ------------------------------------------------------------------ Proposition 4
def test_leverage_multiplier_and_break_even():
    ls = tm.leverage_stats(PAR, EQ["p"], EQ["rho"])
    assert ls["multiplier"] == pytest.approx((1 - PAR.kappa_s) / PAR.d)
    # multiplier falls with the down payment and with costs
    assert tm.leverage_stats(replace(PAR, theta=0.4), EQ["p"], EQ["rho"])["multiplier"] < ls["multiplier"]
    assert tm.leverage_stats(replace(PAR, kappa_b=0.06), EQ["p"], EQ["rho"])["multiplier"] < ls["multiplier"]
    # owner vs renter-investor terminal wealth changes sign exactly at g*
    g = ls["g_star"]
    assert tm.owner_vs_renter_investor(PAR, EQ["p"], EQ["rho"], g) == pytest.approx(0.0, abs=1e-10)
    assert tm.owner_vs_renter_investor(PAR, EQ["p"], EQ["rho"], g + 0.1) > 0
    assert tm.owner_vs_renter_investor(PAR, EQ["p"], EQ["rho"], g - 0.1) < 0
    # break-even growth is independent of theta, increasing in costs, decreasing in rent
    assert tm.leverage_stats(replace(PAR, theta=0.5), EQ["p"], EQ["rho"])["g_star"] == pytest.approx(g)
    assert tm.leverage_stats(replace(PAR, kappa_s=0.10), EQ["p"], EQ["rho"])["g_star"] > g
    assert tm.leverage_stats(PAR, EQ["p"], EQ["rho"] * 1.2)["g_star"] < g
    # underwater threshold: resale at (1 - x*) p leaves zero housing equity
    d2, d3, _ = tm.stationary_schedule(PAR)
    assert (1 - PAR.kappa_s) * (1 - ls["x_star"]) - d3 == pytest.approx(0.0, abs=1e-12)


# ------------------------------------------------------------------ Proposition 5
def test_dynasty_chain_persistence():
    tr = tm.transitions(PAR, EQ["p"], EQ["rho"])
    assert 0 < tr["qRO"] < tr["qOO"] < 1
    assert tr["lam"] == pytest.approx(tr["qOO"] - tr["qRO"])
    assert tr["lam"] == pytest.approx(float(tm.G(PAR, tr["yR"]) - tm.G(PAR, tr["yO"])))
    # stationary distribution of the two-state chain by power iteration
    Pm = np.array([[1 - tr["qRO"], tr["qRO"]], [1 - tr["qOO"], tr["qOO"]]])
    pi = np.array([0.5, 0.5])
    for _ in range(500):
        pi = pi @ Pm
    assert pi[1] == pytest.approx(tr["mu_star"], abs=1e-10)
    assert np.sort(np.abs(np.linalg.eigvals(Pm)))[0] == pytest.approx(tr["lam"], abs=1e-10)
    # persistence vanishes when the house is not net wealth at the end of life (e3 <= 0)
    tr0 = tm.transitions(PAR, EQ["p"], EQ["rho"], eq3=0.0)
    assert tr0["lam"] == pytest.approx(0.0, abs=1e-12)
    # and increases with the equity carried into stage 3 (threshold fixed)
    assert tm.transitions(PAR, EQ["p"], EQ["rho"], eq3=0.9)["lam"] > tr["lam"]
    assert tr["transfer_gap"] == pytest.approx(PAR.nu / (1 + PAR.nu) * tm.stationary_schedule(PAR)[2] * EQ["p"])


# ------------------------------------------------------------------ Proposition 6 + Corollary + 7
def test_stationary_equilibrium_clears_and_is_unique_under_C():
    assert abs(EQ["owner_clearing"]) < 1e-10 and abs(EQ["rental_clearing"]) < 1e-10
    assert EQ["mu"] == pytest.approx(EQ["mu_star"], abs=1e-10)
    assert tm.uniqueness_condition(PAR, EQ["p"], EQ["rho"]) > 0
    # f(mu) is decreasing on a grid -> single root
    grid = np.linspace(0.05, 0.95, 19)
    f = np.array([tm.excess_map(PAR, m) for m in grid])
    assert np.all(np.diff(f) < 0)
    assert np.sum(np.sign(f[1:]) != np.sign(f[:-1])) == 1


def test_fixed_supply_corollary():
    par0 = replace(PAR, eps=0.0, Hbar=0.72)
    e0 = tm.stationary_equilibrium(par0)
    assert e0["mu"] == pytest.approx(0.72)
    # relaxing the down-payment constraint does not change the ownership rate, only the price
    e1 = tm.stationary_equilibrium(replace(par0, theta=0.25))
    assert e1["mu"] == pytest.approx(0.72)
    assert e1["p"] > e0["p"]


def test_no_segmentation_neutrality():
    base = tm.no_segmentation_equilibrium(PAR, phi=0.3)
    for par_alt in [replace(PAR, theta=0.5), replace(PAR, nu=0.6), replace(PAR, y3_hi=9.0)]:
        alt = tm.no_segmentation_equilibrium(par_alt, phi=0.3)
        assert alt["p"] == pytest.approx(base["p"])
        assert alt["rho"] == pytest.approx(base["rho"])
    assert tm.no_segmentation_equilibrium(replace(PAR, theta=0.5), phi=0.3)["b_hat"] != pytest.approx(base["b_hat"])


# ------------------------------------------------------------------ Proposition 8
def test_surprise_accounting_zero_sum_and_contract_ordering():
    p, rho = EQ["p"], EQ["rho"]
    for contract in ["FRM", "ARM"]:
        for exp_ in ["level", "persistent"]:
            acc = tm.surprise_accounting(PAR, p, rho, 0.1, contract=contract, experiment=exp_,
                                         n_savings_age2=0.2, n_savings_age3=0.3)
            assert acc["zero_sum_residual"] == pytest.approx(0.0, abs=1e-12)
            assert acc["gain_age2_D2"] > 0 and acc["gain_age3_D3"] > 0
            assert acc["loss_savers"] > 0
    frm_l = tm.surprise_accounting(PAR, p, rho, 0.1, contract="FRM", experiment="level")
    arm_l = tm.surprise_accounting(PAR, p, rho, 0.1, contract="ARM", experiment="level")
    frm_p = tm.surprise_accounting(PAR, p, rho, 0.1, contract="FRM", experiment="persistent")
    arm_p = tm.surprise_accounting(PAR, p, rho, 0.1, contract="ARM", experiment="persistent")
    assert frm_l["gain_age2_D3"] == pytest.approx(arm_l["gain_age2_D3"])      # level: identical
    assert frm_p["gain_age2_D3"] > arm_p["gain_age2_D3"]                       # persistent: FRM gains more
    assert arm_p["gain_age2_D3"] == pytest.approx(arm_l["gain_age2_D3"])      # ARM reset recovers one factor
    # erosion factors exactly
    d2, d3, _ = tm.stationary_schedule(PAR)
    assert frm_l["D2_shock"] == pytest.approx(d2 * p / 1.1) and frm_l["D3_shock"] == pytest.approx(d3 * p / 1.1)
    assert frm_p["D3_shock"] == pytest.approx(d3 * p / 1.1**2)
    assert frm_l["lender_loss"] == pytest.approx(frm_l["gain_age2_D2"] + frm_l["gain_age2_D3"] + frm_l["gain_age3_D3"])


def test_owner_can_lose_arm_with_real_rate_response_and_unindexed_wages():
    p, rho = EQ["p"], EQ["rho"]
    arm = tm.surprise_accounting(PAR, p, rho, 0.1, contract="ARM", experiment="persistent", phi_r=2.0)
    assert arm["gain_age2_D3"] < 0                       # stage-3 real payment rises
    nowage = tm.surprise_accounting(replace(PAR, eta_w=0.0), p, rho, 0.1, contract="FRM", experiment="level")
    assert nowage["net_owner_age2"] < 0                  # debt erosion < real income loss
    full = tm.surprise_accounting(PAR, p, rho, 0.1, contract="FRM", experiment="level")
    assert full["net_owner_age2"] > 0


# ------------------------------------------------------------------ Proposition 9
def test_impact_signs_segmented_elastic():
    imp = tm.impact_response(PAR, EQ, 0.1)
    assert imp["dp"] > 0 and imp["drho"] < 0 and imp["dmu"] > 0 and imp["db_hat"] > 0
    assert imp["qRO"] < EQ["qRO"] and imp["lam"] > EQ["lam"]
    assert imp["rent_price"] < imp["rent_price_base"]
    assert imp["seller_gain"] > 0
    zero = tm.impact_response(PAR, EQ, 0.0)
    assert abs(zero["dp"]) < 1e-9 and abs(zero["dmu"]) < 1e-9


def test_impact_fixed_supply_is_pure_composition():
    par0 = replace(PAR, eps=0.0, Hbar=0.72)
    e0 = tm.stationary_equilibrium(par0)
    imp = tm.impact_response(par0, e0, 0.1)
    assert imp["dmu"] == pytest.approx(0.0, abs=1e-12)
    assert imp["dp"] > 0 and imp["db_hat"] > 0 and imp["qRO"] < e0["qRO"] and imp["qOO"] > e0["qOO"]


def test_impact_larger_price_response_when_supply_less_elastic():
    outs = []
    for eps in [2.0, 1.0, 0.3]:
        par = replace(PAR, eps=eps, Hbar=EQ["mu"] / EQ["p"] ** eps)  # same stationary point
        e = tm.stationary_equilibrium(par)
        assert e["p"] == pytest.approx(EQ["p"], rel=1e-6)
        outs.append(tm.impact_response(par, e, 0.1)["dp"])
    assert outs[0] < outs[1] < outs[2]


def test_myopic_transition_matches_impact_and_returns_to_steady_state():
    sim = tm.simulate_myopic_transition(PAR, EQ, 0.1, T=12)
    imp = tm.impact_response(PAR, EQ, 0.1)
    assert sim["mu"][0] == pytest.approx(imp["mu_next"], abs=1e-10)
    assert sim["p"][0] == pytest.approx(imp["p"], abs=1e-10)
    assert abs(sim["mu"][-1] - EQ["mu"]) < 0.1 * abs(sim["mu"][0] - EQ["mu"])
    flat = tm.simulate_myopic_transition(PAR, EQ, 0.0, T=6)
    assert np.allclose(flat["mu"], EQ["mu"], atol=1e-10)