"""Economy -> game mapping: steady-state identities, zero-AI limit, household linearisation, compute market."""
import numpy as np

from age.equilibrium import lq_map as m
from age.equilibrium.lq_game import LQParams, wedge


def test_steady_state_resource_constraint_and_zero_ai():
    p = m.EconParams()
    ss = m.steady_state(p)
    assert ss['iota'] == 0.0 and ss['pQ'] == 0.0 and ss['share_L'] == 1 - p.alpha
    # resource constraint C + I + IQ + energy = Y
    assert np.isclose(ss['C'] + ss['I'] + ss['IQ'], ss['Y'])
    # factor payments exhaust output
    assert np.isclose(ss['RK'] * ss['K'] + ss['w'] * ss['L'] + ss['pQ'] * ss['QP'], ss['Y'])
    assert np.isclose(ss['R'], 1 / p.beta)


def test_brock_mirman_closed_form():
    """delta=1, gamma=1 (log) zero-AI: K = alpha*beta*Y in steady state."""
    p = m.EconParams(delta=1.0, gamma=1.0)
    ss = m.steady_state(p)
    assert np.isclose(ss['K'], p.alpha * p.beta * ss['Y'], rtol=1e-10)


def test_linearised_household_mpc_and_news():
    p = m.EconParams()
    ss = m.steady_state(p)
    lin = m.linearized_household(p, ss, T=300)
    # with beta R = 1 the MPC out of a windfall is the annuity value r/R-ish: between 0.02 and 0.06
    assert 0.02 < lin['MPC'] < 0.06
    assert lin['Uxx_level'] < 0
    # savings mapping: negligible stake (documented negative result)
    assert abs(lin['kappa0']) < 0.01
    # GE path returns to a stationary level
    assert abs(lin['a_ge'][-1] - lin['a_ge'][-2]) < 1e-8


def test_portfolio_mapping_and_benchmark_efficiency():
    p = m.EconParams()
    ss = m.steady_state(p)
    pf = m.portfolio_lq_params(p, ss)
    assert pf['alpha'] < 0 and pf['kappa0'] > 0
    for D in (0.0, 0.3, 0.9):
        assert abs(wedge(pf['lqp'], D)) < 1e-6
    # separate investors and workers -> positive wedge (over-delegation), Proposition 4(a)
    pf2 = m.portfolio_lq_params(p, ss, investors_earn_wages=False)
    assert wedge(pf2['lqp'], 0.5) > 0


def test_compute_market_short_run_prop5():
    p = m.EconParams(qbar=0.05)      # inference-heavy delegation
    ss = m.steady_state(p, KQ=0.3)
    pf = m.portfolio_lq_params(p, ss)
    base = m.compute_market_short_run(p, ss, pf['lqp'], pf, KQ=0.3)
    assert np.isclose(base['QP'] + p.qbar * base['D'], 0.3)
    # better model -> more delegation -> higher compute price, less automation, higher labour share
    p2 = m.EconParams(qbar=0.05, sig_A_R=0.02)
    pf2 = m.portfolio_lq_params(p2, ss)
    better = m.compute_market_short_run(p2, ss, pf2['lqp'], pf2, KQ=0.3)
    assert better['D'] > base['D']
    assert better['pQ'] > base['pQ']
    assert better['iota'] < base['iota']
    assert better['share_L'] > base['share_L']


def test_no_private_gain_no_delegation():
    p = m.EconParams(sig_A_R=0.10, sig_H_R=0.10)
    ss = m.steady_state(p, KQ=0.3)
    pf = m.portfolio_lq_params(p, ss)
    assert m.delegation_share(p, ss, pf['lqp'], pf, pQ=0.5) == 0.0