"""Closed forms of the LQ delegation game vs Monte Carlo, and Proposition 3 (zero wedge)."""
import numpy as np
import pytest
from hypothesis import given, settings, strategies as st

from age.equilibrium import lq_game as lq

P = lq.LQParams(kappa0=1.0, alpha=-0.3, tau_theta=1.0, tau_H=1.0, tau_A=4.0, M=2)


def test_best_response_consistency():
    """Strategies satisfy x_i = E_i[kappa0 theta + alpha X] for both agent types."""
    D = 0.4
    s = lq.strategies(P, D)
    bH, bA = P.beta_H, P.beta_A
    # non-delegator: a_H = (kappa0 + alpha Abar) bH
    assert np.isclose(s['aH'], (P.kappa0 + P.alpha * s['Abar']) * bH)
    # delegator: a_A = bA (kappa0 + alpha Abar) + alpha (D/M)(1-bA) a_A
    assert np.isclose(s['aA'], bA * (P.kappa0 + P.alpha * s['Abar']) + P.alpha * (D / P.M) * (1 - bA) * s['aA'])


def test_full_info_limits():
    """tau_A -> inf with D=1 reproduces the full-information response kappa0/(1-alpha)."""
    p = lq.LQParams(kappa0=1.0, alpha=-0.3, tau_theta=1.0, tau_H=1.0, tau_A=1e9, M=1)
    s = lq.strategies(p, 1.0)
    assert np.isclose(s['Abar'], p.kappa0 / (1 - p.alpha), rtol=1e-6)
    assert lq.moments(p, 1.0)['V_noise'] < 1e-8


def test_moments_vs_monte_carlo():
    D = 0.5
    m = lq.moments(P, D)
    sim = lq.simulate(P, D, n_agents=4000, n_draws=3000, seed=1)
    assert np.isclose(sim['EX2'], m['EX2'], rtol=0.08)
    assert np.isclose(sim['Ex2'], m['Ex2'], rtol=0.05)
    assert np.isclose(sim['V_disp'], m['V_disp'], rtol=0.05)
    assert np.isclose(sim['Abar_hat'], m['Abar'], rtol=0.05)


def test_lln_breaks_only_for_delegators():
    """Aggregate noise is zero at D=0 and grows with D^2/M."""
    assert lq.moments(P, 0.0)['V_noise'] == 0.0
    v1 = lq.moments(lq.LQParams(**{**P.__dict__, 'M': 1}), 0.6)['V_noise']
    v4 = lq.moments(lq.LQParams(**{**P.__dict__, 'M': 4}), 0.6)['V_noise']
    assert v4 < v1


def test_welfare_decomposition_matches_quadratic_form():
    """Under CRS (upsilon=1, U_Xtheta=0) the AP decomposition equals the exact quadratic welfare up to a constant."""
    Ds = np.linspace(0, 1, 11)
    exact = np.array([lq.welfare_gross(P, D) for D in Ds])
    dec = np.array([lq.welfare_decomposed(P, D)['W_rel'] for D in Ds])
    diff = exact - dec
    assert np.allclose(diff, diff[0], atol=1e-10)


@settings(max_examples=60, deadline=None)
@given(alpha=st.floats(-0.95, 0.95), D=st.floats(0.0, 1.0), M=st.integers(1, 6),
       tau_H=st.floats(0.1, 10), tau_A=st.floats(0.1, 20), tau_t=st.floats(0.1, 10))
def test_proposition3_zero_wedge(alpha, D, M, tau_H, tau_A, tau_t):
    p = lq.LQParams(kappa0=1.0, alpha=alpha, tau_theta=tau_t, tau_H=tau_H, tau_A=tau_A, M=M)
    assert abs(lq.wedge(p, D)) < 1e-5


def test_proposition4a_nonzero_wedge_without_crs():
    p = lq.LQParams(kappa0=1.0, alpha=-1 / 3, tau_theta=1, tau_H=1, tau_A=4, M=1, upsilon=0.0)
    assert lq.wedge(p, 0.5) > 1e-3      # over-delegation when the aggregate externality is weaker than pecuniary
    p2 = lq.LQParams(**{**p.__dict__, 'upsilon': 2.0})
    assert lq.wedge(p2, 0.5) < -1e-3


def test_private_value_alpha0_closed_form():
    p = lq.LQParams(kappa0=1.3, alpha=0.0, tau_theta=2.0, tau_H=0.5, tau_A=6.0, M=3)
    pv = lq.private_value(p, 0.3)
    closed = p.kappa0 ** 2 / 2 * (p.tau_A - p.tau_H) / ((p.tau_A + p.tau_theta) * (p.tau_H + p.tau_theta))
    assert np.isclose(pv, closed)


def test_delegation_equilibrium_threshold_and_comparative_statics():
    from scipy.stats import lognorm
    cdf = lambda c: lognorm(s=1.0, scale=np.exp(-2.5)).cdf(np.maximum(c, 0))
    D1 = lq.delegation_equilibrium(P, cdf, scale=1.0)
    D2 = lq.delegation_equilibrium(lq.LQParams(**{**P.__dict__, 'tau_A': 8.0}), cdf, scale=1.0)
    assert 0 < D1 < 1 and D2 > D1       # better model -> more delegation
    # planner agrees with the market under CRS (Proposition 3)
    q = lambda u: lognorm(s=1.0, scale=np.exp(-2.5)).ppf(np.clip(u, 1e-9, 1 - 1e-9))
    Dp = lq.planner_delegation(P, cdf, q, scale=1.0)
    assert abs(Dp - D1) < 0.01