"""Stationary equilibrium with free entry, posted prices, financing extension, planner, transition."""
import numpy as np
import pytest

from src.solver import (solve_equilibrium, partial_equilibrium, free_entry_residual, solve_equilibrium_posted,
                        solve_equilibrium_financing, transition, frictionless_price, welfare_flow)


def test_baseline_equilibrium_residuals(params):
    eq = solve_equilibrium(params)
    assert eq.converged and len(eq.theta_roots) == 1
    assert abs(eq.free_entry_resid) < 1e-10
    assert np.abs(eq.sellers.resid).max() < 1e-10
    # stock equation
    assert np.allclose(params.s * params.f, (eq.sellers.eta + params.delta) * eq.n)
    assert abs(eq.omega.sum() - 1) < 1e-12 and abs(eq.N.sum() - 1) < 1e-12
    m = eq.moments(params)
    assert m["participation_ok"]
    assert 0 < m["price_mean"] < params.G.hi
    assert m["D_urgent"] > 0 and m["mean_tom_sold"] > 0
    # free-entry residual is decreasing in theta around the root (uniqueness check, heterogeneous case)
    th = eq.theta * np.array([0.5, 0.8, 1.0, 1.25, 2.0])
    res = np.array([free_entry_residual(t, params) for t in th])
    assert np.all(np.diff(res) < 0)


def test_prop3_homogeneous_uniqueness_and_comparative_statics(params):
    p = params.copy(c=np.array([0.08]), f=np.array([1.0]))
    grid = np.geomspace(0.05, 50, 40)
    res = np.array([free_entry_residual(t, p) for t in grid])
    assert np.all(np.diff(res) < 0)                      # strictly decreasing => unique root
    th_phi = [solve_equilibrium(p.copy(phi=ph)).theta for ph in (0.6, 0.75, 0.9)]
    th_A = [solve_equilibrium(p.copy(A=a)).theta for a in (3.0, 4.5, 6.0)]
    assert np.all(np.diff(th_phi) > 0) and np.all(np.diff(th_A) > 0)
    # observational equivalence of A and phi in the baseline
    e1 = solve_equilibrium(p.copy(phi=0.5, A=9.0)); e2 = solve_equilibrium(p.copy(phi=1.0, A=4.5))
    assert abs(e1.theta - e2.theta) < 1e-8 and abs(e1.Lam_s - e2.Lam_s) < 1e-8
    # amplification: d ln Lam_s / d ln phi > 1
    lo, hi = solve_equilibrium(p.copy(phi=0.84)), solve_equilibrium(p.copy(phi=0.86))
    el = (np.log(hi.Lam_s) - np.log(lo.Lam_s)) / (np.log(0.86) - np.log(0.84))
    assert el > 1.0


def test_prop4_liquidity_shock(params):
    base = solve_equilibrium(params)
    shock = solve_equilibrium(params.copy(A=params.A * 0.7))
    assert shock.theta < base.theta
    assert np.all(shock.sellers.V < base.sellers.V) and np.all(shock.sellers.pbar < base.sellers.pbar)
    assert np.all(shock.sellers.eta < base.sellers.eta)
    assert shock.S > base.S and shock.Q < base.Q                     # delta > 0: steady-state volume falls
    # impact effect: stock predetermined
    Q_impact = float((shock.sellers.eta * base.n).sum())
    assert Q_impact < base.Q
    # delta = 0: steady-state volume equals inflow and composition equals f (Proposition 6(i))
    p0 = params.copy(delta=0.0)
    e0 = solve_equilibrium(p0); e0s = solve_equilibrium(p0.copy(A=p0.A * 0.7))
    assert abs(e0.Q - p0.s) < 1e-10 and abs(e0s.Q - p0.s) < 1e-10
    assert np.allclose(e0.omega, p0.f) and np.allclose(e0s.omega, p0.f)


def test_prop5_dispersion_bounds(params):
    eq = solve_equilibrium(params)
    m = eq.moments(params)
    assert m["price_var_between"] > 0
    rng = eq.sellers.pbar[0] - eq.sellers.pbar[-1]
    assert rng <= (params.c[-1] - params.c[0]) / (params.r + params.delta)
    # beta -> 0: price = V(c); between variance <= Var(c)/(r+delta)^2
    p = params.copy(beta=1e-6)
    e = solve_equilibrium(p)
    var_c = float((e.omega * (params.c - (e.omega * params.c).sum()) ** 2).sum())
    assert e.moments(p)["price_var_between"] <= var_c / (p.r + p.delta) ** 2 + 1e-12


def test_prop6_impact_composition(params):
    """Impact decomposition of Proposition 6(ii): the composition term has the sign of Cov_omega(pbar, eps) d ln Lam_s.
    For the truncated-normal baseline h_G*M is increasing, so eps(c) increases in c and selection DAMPENS the index
    response; for exponential G (eps decreasing in c) selection AMPLIFIES it."""
    from src.model import Exponential
    for G, expect_amplify in ((params.G, False), (Exponential(rate=8.0, lo=0.7), True)):
        p = params.copy(G=G)
        base = solve_equilibrium(p)
        pe = partial_equilibrium(base.theta, p.copy(A=p.A * 0.9))   # same theta, lower Lam_s, stock predetermined
        eta_new = pe.sellers.eta
        omega_imp = eta_new * base.n / (eta_new * base.n).sum()
        P_obs = float((omega_imp * pe.sellers.pbar).sum()); P_fix = float((base.omega * pe.sellers.pbar).sum())
        P0 = base.moments(p)["price_mean"]
        eps = (np.log(eta_new) - np.log(base.sellers.eta)) / np.log(0.9)
        cov = float((base.omega * (base.sellers.pbar - P_fix) * (eps - (base.omega * eps).sum())).sum())
        assert P_fix < P0                                            # within-type prices fall
        assert np.sign(P_obs - P_fix) == np.sign(-cov)               # composition term = Cov * d ln Lam (<0)
        if expect_amplify:
            assert np.all(np.diff(eps) < 0) and P_obs < P_fix
        else:
            assert np.all(np.diff(eps) > 0) and P_obs > P_fix
        # analytic elasticity formula eps = 1 - beta eta h_G M /(r+delta+beta eta) at the base point (small-shock check)
        sb = base.sellers
        an = 1 - p.beta * sb.eta * np.asarray(G.hazard(sb.vstar)) * sb.M / ((p.r + p.delta) + p.beta * sb.eta)
        assert np.allclose(an, eps, atol=0.02)


def test_posted_prices(params):
    eq = solve_equilibrium_posted(params)
    assert eq.converged and np.abs(eq.sellers.resid).max() < 1e-9 and abs(eq.free_entry_resid) < 1e-9
    assert np.all(np.diff(eq.sellers.pbar) < 0) and np.all(np.diff(eq.sellers.eta) > 0)
    # posted threshold exceeds Nash threshold at equal (V,U)
    assert np.all(eq.sellers.vstar > eq.sellers.V + params.k)
    # FOC: y - U - V = (1-G)/g
    foc = eq.sellers.vstar - params.k - eq.sellers.V - params.G.surv(eq.sellers.vstar) / params.G.pdf(eq.sellers.vstar)
    assert np.abs(foc).max() < 1e-7


@pytest.mark.slow
def test_financing_extension(params):
    fe = solve_equilibrium_financing(params)
    assert fe.converged and abs(fe.resid["free_entry"]) < 1e-8 and fe.resid["max_abs"] < 1e-8
    # Proposition 7(ii): the selection term has the sign of v*_m - v*_c (positive iff the pending wedge exceeds U_c - U_m)
    assert np.all(np.sign(fe.sel_term) == np.sign(fe.vstar_m - fe.vstar_c))
    wedge = (params.c + params.kappa_b + params.r * (fe.V + fe.U_m)) / (params.nu * params.phi_m)
    assert np.allclose(fe.vstar_m - fe.vstar_c, wedge - (fe.U_c - fe.U_m))
    assert np.all(np.diff(fe.Delta_b) >= -1e-12)                         # Delta_b nondecreasing in c
    assert fe.U_c > fe.U_m                                               # cash buyers have the better outside option
    assert np.all(np.diff(fe.pbar_m - fe.pbar_c) > 0)                    # gap increasing in urgency (numerical, baseline)
    assert abs(fe.Q_c.sum() / fe.Q - params.pi) < 1e-8                   # steady-state sales composition = entry composition
    assert 0 < fe.pi_tilde < 1
    # Delta_b formula check against the closed form of Proposition 7
    an = ((1 - params.beta) * (params.c + params.r * fe.V) - params.beta * (params.kappa_b + params.r * fe.U_m)) / (params.nu * params.phi_m)
    assert np.allclose(an, fe.Delta_b)
    # uniqueness: different starting points reach the same equilibrium
    f2 = solve_equilibrium_financing(params, x0=np.array([np.log(1.5), 0.0, 0.5]))
    assert abs(f2.theta - fe.theta) < 1e-6
    # nesting: nu -> infinity, pi -> 1 reproduces the baseline with phi = 1
    p2 = params.copy(nu=1e5, pi=0.999)
    fe2 = solve_equilibrium_financing(p2); eq = solve_equilibrium(params.copy(phi=1.0))
    assert abs(fe2.theta - eq.theta) / eq.theta < 1e-2 and np.abs(fe2.V - eq.sellers.V).max() < 1e-3


@pytest.mark.slow
def test_transition_converges_and_matches_steady_states(params):
    p_new = params.copy(A=params.A * 0.8)
    tr = transition(params, lambda t: p_new, T=6.0, dt=0.05)
    assert tr.converged
    e0 = solve_equilibrium(params); e1 = solve_equilibrium(p_new)
    assert abs(tr.P_obs[0] - 0) >= 0
    assert np.allclose(tr.n[0], e0.n)
    assert abs(tr.theta[-1] - e1.theta) / e1.theta < 2e-2
    assert np.abs(tr.n[-1] - e1.n).max() < 2e-2 * e1.n.max()
    # composition term: sign on impact follows Proposition 6 (dampening for the truncated normal: P_obs > P_fix after a
    # negative shock) and is an order of magnitude smaller than the within-type change; volume and inventory adjust slowly
    assert tr.P_obs[0] > tr.P_fix[0]                                   # impact: composition relative to the pre-shock weights
    assert abs(np.log(tr.P_obs[0] / tr.P_fix[0])) < 0.1 * abs(np.log(tr.P_fix[0] / e0.moments(params)["price_mean"]))
    assert tr.Q[0] < tr.Q[-1] and tr.S[0] < tr.S[-1]
    assert np.all(tr.Q >= 0) and np.all(tr.S > 0)
    # exponential G: amplification (P_obs < P_fix on impact)
    from src.model import Exponential
    pe = params.copy(G=Exponential(rate=8.0, lo=0.7)); tre = transition(pe, lambda t: pe.copy(A=pe.A * 0.8), T=3.0, dt=0.05)
    assert tre.converged and tre.P_obs[0] < tre.P_fix[0]


def test_frictionless_price_and_welfare(params):
    ps = frictionless_price(params)
    assert abs(float(params.G.H(ps)) - params.r * params.k) < 1e-10
    eq = solve_equilibrium(params)
    w = welfare_flow(eq, params)
    assert w["total"] > 0 and w["holding_costs"] > 0


def test_grid_and_tolerance_convergence(params):
    """Equilibrium objects stable across the seller-type grid and root tolerances."""
    m9 = solve_equilibrium(params).moments(params)
    p25 = params.with_types(0.02, 0.3, 25, 2.0, 4.0)
    m25 = solve_equilibrium(p25).moments(p25)
    for key in ("theta", "price_mean", "mean_tom_sold", "Q"):
        assert abs(m9[key] - m25[key]) / abs(m25[key]) < 1e-2, key
    mt = solve_equilibrium(params.copy(root_xtol=1e-8)).moments(params)
    assert abs(mt["theta"] - m9["theta"]) < 1e-6