"""Household solver: feasibility, monotonicity, terminal condition, Bellman residuals, thresholds."""
import numpy as np

from src import household_solver as HS
from src import model as M
from src.household_solver import NEG


def test_values_finite_and_monotone_in_assets(P_small, pols):
    for j in range(1, P_small.J + 1):
        V = pols[j - 1].V_rent
        feas = V > NEG / 2
        assert feas[:, 5:].all()
        for iz in range(P_small.n_z):
            v = V[iz][feas[iz]]
            assert np.all(np.diff(v) >= -1e-9), f"V_rent not increasing in a at age {j}, z {iz}"


def test_owner_value_decreasing_in_debt(P_small, pols):
    for j in (2, 5, 8):
        V = pols[j - 1].V_own
        for iz in range(P_small.n_z):
            for ih in range(P_small.n_h):
                v = V[iz, ih, :, P_small.N - 2, -1]      # across balances at the top asset
                v = v[v > NEG / 2]
                assert np.all(np.diff(v) <= 1e-9)


def test_policies_feasible(P_small, pols, prices):
    P = P_small
    for j in range(1, P.J + 1):
        pol = pols[j - 1]
        assert np.all(pol.aR >= 0) and np.all(pol.aB >= 0) and np.all(pol.aS >= 0) and np.all(pol.aF >= 0)
        feasB = pol.W_B > NEG / 2
        assert np.all(pol.hB[feasB] >= 0)
        # loan never above the LTV / DSTI cap
        cap = pol.maxloan[np.arange(P.n_z)[:, None], np.maximum(pol.hB, 0)]
        assert np.all(pol.BB[feasB] <= cap[feasB] + 1e-9)
        assert np.all(pol.BB[feasB] <= (1 - P.theta) * prices.p * P.h_grid[np.maximum(pol.hB, 0)][feasB] + 1e-9)
        # gifts only at the gift age
        if j != P.j_gift:
            assert np.all(pol.gB == 0) and np.all(pol.gR == 0) and np.all(pol.gS == 0)


def test_terminal_condition_is_warm_glow(P_small, prices):
    P = P_small
    EV_own, EV_rent = HS.terminal_continuation(P, prices)
    a = P.a_grid
    assert np.allclose(EV_rent[0], M.warm_glow(P, (1 - P.tau_b) * a))
    # an owner with no debt: estate = a + (1 - kappa_s) p h
    est = a + (1 - P.kappa_s) * prices.p * P.h_grid[1]
    assert np.allclose(EV_own[0, 1, 0, 0, :], M.warm_glow(P, (1 - P.tau_b) * est))
    # negative estate (underwater at death) is infeasible
    big_b = P.b_grid[-1]
    if big_b > (1 - P.kappa_s) * prices.p * P.h_grid[0]:
        assert EV_own[0, 0, -1, 0, 0] <= NEG / 2


def test_bellman_residuals_small(P_small, pols, prices, lottery):
    res = HS.bellman_residuals(P_small, pols, prices, lottery, n_points=200, n_fine=300)
    assert res["owner_mean"] < 5e-3 and res["renter_mean"] < 5e-3
    assert res["owner_max"] < 5e-2 and res["renter_max"] < 2e-2


def test_ownership_threshold_single_crossing(P_small, pols, prices):
    P = P_small
    th = HS.ownership_threshold(P, pols[0], prices, 1)
    for iz, v in th.items():
        assert np.isfinite(v["x_hat"])
        assert v["sign_changes"] == 1, f"tenure choice at age 1 not single-crossing for z index {iz}"
    # thresholds (in cash) are above the pure feasibility cash need
    assert all(v["x_hat"] > v["feasibility_cash"] for v in th.values())


def test_threshold_rises_with_price_and_down_payment(P_small, lottery):
    P = P_small
    base = HS.ownership_threshold(P, HS.solve_lifecycle(P, M.stationary_prices(P, 1.0, 1.0), lottery)[0], M.stationary_prices(P, 1.0, 1.0), 1)
    hi_p = HS.ownership_threshold(P, HS.solve_lifecycle(P, M.stationary_prices(P, 1.08, 1.0), lottery)[0], M.stationary_prices(P, 1.08, 1.0), 1)
    iz = P.n_z // 2
    assert hi_p[iz]["x_hat"] > base[iz]["x_hat"]
    cfg = M.load_config(overrides={"mortgage.theta": 0.35, "simulation.n_dynasties": P.n_dyn,
                                   "grids.n_a": P.n_a, "grids.n_x": P.n_x})
    P2 = M.build_params(cfg)
    pr2 = M.stationary_prices(P2, 1.0, 1.0)
    hi_theta = HS.ownership_threshold(P2, HS.solve_lifecycle(P2, pr2, lottery)[0], pr2, 1)
    assert hi_theta[iz]["x_hat"] > base[iz]["x_hat"]