"""Primitives: income process, mortgage arithmetic (Fisher, annuity-due), supply, utility."""
import numpy as np

from src import model as M


def test_rouwenhorst_moments():
    grid, Pi = M.rouwenhorst(5, 0.7, 0.3)
    pi = M.stationary_dist(Pi)
    assert np.allclose(Pi.sum(axis=1), 1.0)
    assert abs(pi @ grid) < 1e-12
    var = pi @ grid**2
    assert abs(var - 0.3**2 / (1 - 0.7**2)) < 1e-8
    # first-order autocorrelation equals rho
    acov = sum(pi[i] * grid[i] * sum(Pi[i, k] * grid[k] for k in range(5)) for i in range(5))
    assert abs(acov / var - 0.7) < 1e-8


def test_period_conversions(P_small):
    P = P_small
    cfg = P.cfg
    assert abs(P.beta - cfg["preferences"]["beta_annual"] ** 5) < 1e-12
    assert abs(P.R - (1 + cfg["rates"]["r_annual"]) ** 5) < 1e-12
    assert abs(P.rho_z - cfg["income"]["rho_annual"] ** 5) < 1e-12
    # Fisher: 1 + i = R (1 + pibar)
    assert abs((1 + P.i_nom) - P.R * (1 + P.pibar)) < 1e-12
    assert abs(P.pi_z @ P.z_grid - 1.0) < 1e-10          # E[z] = 1
    assert np.all(P.ybar[P.j_retire - 1:] == P.rr * P.ybar[P.j_retire - 2])


def test_annuity_due_amortises_in_N_payments(P_small):
    P = P_small
    for i in (0.0, 0.1, P.i_nom):
        b = 1.0
        pays = []
        for n in range(P.N, 0, -1):
            pays.append(M.payment(b, n, i))
            b = M.next_balance(b, n, i, 0.0)          # nominal path (no deflation)
        assert abs(b) < 1e-12
        assert np.allclose(pays, pays[0])             # constant nominal payment
        # present value of the annuity-due at rate i equals the principal
        pv = sum(p / (1 + i) ** k for k, p in enumerate(pays))
        assert abs(pv - 1.0) < 1e-12


def test_real_balance_erodes_with_inflation(P_small):
    P = P_small
    b0 = 1.0
    b_nom = M.next_balance(b0, P.N, P.i_nom, 0.0)
    b_real = M.next_balance(b0, P.N, P.i_nom, P.pibar)
    assert abs(b_real * (1 + P.pibar) - b_nom) < 1e-12
    # a surprise pi_u on top of pibar erodes the real balance by 1/(1+pi_u)
    b_surprise = M.next_balance(b0, P.N, P.i_nom, (1 + P.pibar) * 1.1 - 1)
    assert abs(b_surprise * 1.1 - b_real) < 1e-12


def test_borrowing_is_pv_neutral_against_deposits(P_small):
    """Borrow B, pay the first payment now, deposit the rest: next period the real position equals
    the one of not borrowing (annuity-due timing makes loans and deposits PV-consistent)."""
    P = P_small
    B = 0.8
    pay = M.payment(B, P.N, P.i_nom)
    deposit_next = (B - pay) * P.R
    debt_next = M.next_balance(B, P.N, P.i_nom, P.pibar)
    assert abs(deposit_next - debt_next) < 1e-12


def test_ltv_dsti_caps(P_small):
    P = P_small
    y = 1.0
    cap = M.max_loan(P, 1.0, P.h_grid[0], y, P.i_q)
    assert cap <= (1 - P.theta) * P.h_grid[0] + 1e-12
    assert M.payment(cap, P.N, P.i_q) <= P.psi * y + 1e-10
    # the stress test tightens the cap
    assert M.max_loan(P, 1.0, P.h_grid[-1], y, P.i_q) <= M.max_loan(P, 1.0, P.h_grid[-1], y, P.i_nom) + 1e-12


def test_supply_and_rent(P_small):
    P = P_small
    H = M.supply_stock(P, 1.0, 10.0, 1.0)
    assert abs(H - 10.0) < 1e-12
    assert abs(M.stock_next(P, H, 1.0, 10.0, 1.0) - H) < 1e-12          # stationary stock
    assert M.supply_stock(P, 1.1, 10.0, 2.0) > M.supply_stock(P, 1.1, 10.0, 0.5)
    rho = M.rent_from_price(P, 1.0, 1.0, P.R)
    assert abs(rho - (1 + P.m + P.phi - 1 / P.R)) < 1e-12


def test_utility_homogeneity(P_small):
    P = P_small
    c, s = 1.3, 0.8
    k = 1.7
    expo = (1 - P.alpha) * (1 - P.gamma)
    assert abs(M.utility(P, k * c, s) - k**expo * M.utility(P, c, s)) < 1e-12
    assert M.warm_glow(P, 1.0) > M.warm_glow(P, 0.5)