"""Model B: zero-AI Aiyagari benchmark, market clearing, distribution mass, delegation monotone in wealth."""
import numpy as np
import pytest

from age.equilibrium.aiyagari import EconomyB, HAParams, rouwenhorst


def test_rouwenhorst_moments():
    g, P = rouwenhorst(7, 0.96, 0.2)
    pi = np.linalg.matrix_power(P, 2000)[0]
    assert np.isclose(pi @ g, 0.0, atol=1e-8)
    assert np.isclose(np.sqrt(pi @ g ** 2), 0.2 / np.sqrt(1 - 0.96 ** 2), rtol=1e-6)
    assert np.allclose(P.sum(axis=1), 1.0)


@pytest.fixture(scope='module')
def zero_ai():
    p = HAParams(n_a=100, delegation_on=False, KQ0=0.0, m_H=0.0, sigma_rH=1e-6, sigma_rA=1e-6, n_gh=3)
    e = EconomyB(p)
    return e, e.solve(tol=1e-4)


def test_zero_ai_clears_and_is_sane(zero_ai):
    e, res = zero_ai
    assert 0.0 < res['r'] < 1 / e.p.beta - 1          # r below the complete-markets rate (precautionary saving)
    assert np.isclose(e.hh.lam.sum(), 1.0, atol=1e-8)
    agg = e.hh.aggregates()
    assert abs(agg['A'] - res['K']) < 1e-2 * res['K']  # capital market clears within tolerance
    assert res['D'] == 0.0 and res['iota'] == 0.0
    assert 0.3 < res['gini'] < 0.9
    assert np.isclose(res['shares']['labor'] + res['shares']['capital'] + res['shares']['compute'], 1.0, atol=1e-6)


def test_goods_market_identity(zero_ai):
    e, res = zero_ai
    # C + delta K + mistakes + fees-as-compute = Y  (fees are transfers to compute owners here: zero AI => 0)
    assert np.isclose(res['C'] + e.p.delta * res['K'] + res['mistakes'], res['Y'], rtol=2e-2)


def test_delegation_increases_with_wealth():
    p = HAParams(n_a=100, delegation_on=True, KQ0=0.0, m_H=0.01, sigma_rH=0.04, sigma_rA=0.01, phi0=0.02, phi1=0.002,
                 n_gh=3)
    e = EconomyB(p)
    e.excess_assets(0.02)
    ws = e.hh.wealth_stats()
    d = ws['deleg_by_decile']
    assert 0 < e.last['agg']['D'] < 1
    assert d[-1] >= d[0]                                   # top decile delegates at least as much as the bottom
    assert np.corrcoef(np.arange(10), d)[0, 1] > 0.5