# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Checks of the theory module against the closed forms stated in paper/sections/theory.tex and welfare.tex."""
from __future__ import annotations

import itertools

import numpy as np
import pytest

from agentic_valuation.theory import aggregation_sim as ag
from agentic_valuation.theory import welfare as wf

SEED = 20261027


# ----------------------------------------------------------------------------------------- Propositions 1–3
@pytest.mark.parametrize("K,rho", [(1, 0.0), (3, 0.2), (12, 0.4), (12, 0.0), (48, 0.8), (6, -0.1)])
def test_closed_form_mse_matches_simulation(K, rho):
    rng = np.random.default_rng(SEED)
    closed = ag.mse_mean_closed_form(K, 1.0, rho)
    sim = ag.simulate_mean_mse(K, 1.0, rho, 200_000, rng)
    assert abs(sim - closed) < 0.02 * max(closed, 0.05) + 1e-3


def test_exchangeable_equal_weights_are_optimal_and_floor():
    for K, rho in itertools.product([2, 5, 12], [0.0, 0.3, 0.7]):
        S = ag.exchangeable_cov(K, 1.0, rho)
        assert ag.mse_optimal(S) == pytest.approx(ag.mse_equal(S), rel=1e-10)
        assert ag.mse_equal(S) >= rho  # floor rho * sigma^2 is never beaten by K
    assert ag.mse_mean_closed_form(10_000, 1.0, 0.5) == pytest.approx(0.5, abs=1e-3)


def test_general_sigma_optimal_dominates_equal_weights():
    rng = np.random.default_rng(SEED)
    for _ in range(20):
        A = rng.standard_normal((8, 8))
        S = A @ A.T + 0.1 * np.eye(8)
        assert ag.mse_optimal(S) <= ag.mse_equal(S) + 1e-12
    S = ag.exchangeable_cov(5, 2.0, 0.3)  # rows with equal sums: equality
    assert ag.mse_optimal(S) == pytest.approx(ag.mse_equal(S))


def test_nested_common_model_closed_form():
    for n, M in [(12, 1), (6, 2), (4, 3), (1, 12)]:
        S = ag.nested_cov(n, M, 0.3, 0.4, 0.5)
        assert ag.mse_equal(S) == pytest.approx(ag.mse_nested_closed_form(n, M, 0.3, 0.4, 0.5))
    # more providers at fixed K lowers the pooled MSE; the common component is a floor
    vals = [ag.mse_nested_closed_form(12 // M, M, 0.3, 0.4, 0.5) for M in (1, 2, 3, 4, 6, 12)]
    assert all(a > b for a, b in zip(vals, vals[1:]))
    assert min(vals) >= 0.3**2


# ----------------------------------------------------------------------------------------- Propositions 4, 7
def test_opinion_only_deliberation_contracts_dispersion_and_keeps_mean():
    rng = np.random.default_rng(SEED)
    x0 = rng.standard_normal((50_000, 12)) * 0.2
    lam = 0.4
    for anchor in ("mean", "median"):
        traj = ag.deliberate_opinion_only(x0, 4, lam, anchor)
        d0 = traj[0].std(axis=1)
        for r in range(5):
            np.testing.assert_allclose(traj[r].std(axis=1), d0 * (1 - lam) ** r, rtol=1e-10)
            # the anchoring statistic is invariant across rounds
            stat = np.median if anchor == "median" else np.mean
            np.testing.assert_allclose(stat(traj[r], axis=1), stat(traj[0], axis=1), atol=1e-12)
        if anchor == "mean":
            np.testing.assert_allclose(traj[4].mean(axis=1), traj[0].mean(axis=1), atol=1e-12)
        else:  # exact formula: MSE_r = MSE_0 + (1 - (1-lam)^r)^2 (MSE_med - MSE_mean)
            m0 = np.mean(traj[0].mean(axis=1) ** 2)
            md = np.mean(np.median(traj[0], axis=1) ** 2)
            for r in range(5):
                mr = np.mean(traj[r].mean(axis=1) ** 2)
                assert mr == pytest.approx(m0 + (1 - (1 - lam) ** r) ** 2 * (md - m0), rel=0.02)
            assert md >= m0  # the median is dominated by the mean under Gaussian errors


def test_information_conservation_and_over_shrinkage():
    sig2 = np.linspace(0.05, 0.4, 12) ** 2
    S0 = np.diag(sig2)
    best0 = ag.mse_optimal(S0)
    prev = best0
    for r in range(1, 6):
        Sr = ag.deliberation_cov(S0, r, 0.5)
        assert ag.mse_optimal(Sr) >= best0 - 1e-12  # no pool of round r beats the best round-0 pool
        iv = ag.mse_pool(ag.inverse_variance_weights(Sr), Sr)
        assert iv > prev  # strictly dominated and worsening (heterogeneous variances)
        prev = iv
    assert prev <= ag.mse_equal(S0) + 1e-12 and prev == pytest.approx(ag.mse_equal(S0), rel=0.05)


def test_degroot_consensus_is_round0_mean_for_complete_topology():
    K, lam = 7, 0.3
    W = (1 - lam) * np.eye(K) + lam * np.ones((K, K)) / K
    np.testing.assert_allclose(ag.degroot_consensus_weights(W), np.full(K, 1 / K), atol=1e-10)


# ----------------------------------------------------------------------------------------- Proposition 6
def test_evidence_exchange_closed_form_and_sign():
    rng = np.random.default_rng(SEED)
    K, J = 12, 28
    views = [np.array([0, 1, 2, 3, 4 + 2 * k, 5 + 2 * k]) for k in range(K)]
    comm = ag.EvidenceCommittee(J=J, views=views, providers=np.arange(K) % 3, tau2=np.full(K, 0.01),
                                sigma_c=0.05, sigma_u=0.12, sigma_g=0.04)
    full = ag.exchange_views(views)
    before, after = comm.mse_equal_closed_form(), comm.mse_equal_closed_form(full)
    assert after < before
    assert after == pytest.approx(0.05**2 + 0.12**2 / J + 0.04**2 / 3 + 0.01 / K)
    res = ag.simulate_deliberation(comm, 1, 0.0, 60_000, rng, exchange_at=1)
    assert res["mse_mean"][0] == pytest.approx(before, rel=0.03)
    assert res["mse_mean"][1] == pytest.approx(after, rel=0.03)
    # exchange restricted to the common view is vacuous
    same = ag.exchange_views(views, cited=[np.array([0, 1, 2, 3])] * K)
    assert comm.mse_equal_closed_form(same) == pytest.approx(before)
    # persuasion noise (independent of the errors) raises the MSE
    res2 = ag.simulate_deliberation(comm, 2, 0.3, 60_000, rng, anchor="mean", persuasion_sd=0.05)
    assert res2["mse_mean"][2] > res2["mse_mean"][0]


def test_partial_exchange_can_raise_mse():
    ex = ag.EvidenceCommittee(J=2, views=[np.array([0]), np.array([1]), np.array([1])], providers=np.zeros(3, int),
                              tau2=np.zeros(3), sigma_u=1.0)
    before = ex.mse_equal_closed_form()
    partial = ex.mse_equal_closed_form(views=[np.array([0, 1]), np.array([1]), np.array([1])])
    assert before == pytest.approx(5 / 9) and partial == pytest.approx(13 / 18) and partial > before


# ----------------------------------------------------------------------------------------- Proposition 8
def test_placebo_invariance_under_separability():
    rng = np.random.default_rng(SEED)
    K, J = 12, 28
    views = [np.array([0, 1, 2, 3, 4 + 2 * k, 5 + 2 * k]) for k in range(K)]
    tau2 = np.linspace(0.05, 0.3, K) ** 2
    comm = ag.EvidenceCommittee(J=J, views=views, providers=np.arange(K) % 3, tau2=tau2, sigma_c=0.05,
                                sigma_u=0.12, sigma_g=0.04)
    base = comm.mse_equal_closed_form()
    shocks = comm.draw_shocks(50_000, rng)
    x_id = comm.draw(50_000, rng, shocks=shocks)
    for _ in range(10):
        perm = rng.permutation(K)
        assert ag.placebo_mse(comm, perm) == pytest.approx(base)
        x_p = comm.draw(50_000, rng, tau2=tau2[perm], shocks=shocks)
        assert np.mean(x_p.mean(axis=1) ** 2) == pytest.approx(np.mean(x_id.mean(axis=1) ** 2), rel=0.03)
    # a prompt x view complementarity (variance depending on the match) breaks the invariance
    tau2_match = tau2.copy()
    tau2_match[0] = 1.0  # slot 0 mismatched prompt
    assert comm.mse_equal_closed_form(tau2=tau2_match) > base


# ----------------------------------------------------------------------------------------- Welfare
def test_failure_probability_monotone_for_financeable_matches():
    p = wf.WelfareParams()
    u_s = np.array([-0.10, -0.05, 0.0])  # financeable at the error-free appraisal (a < 0)
    prev = wf.failure_prob_given_seller(u_s, 0.0, 0.01, p)
    for s in np.linspace(0.02, 0.3, 15):
        cur = wf.failure_prob_given_seller(u_s, 0.0, float(s), p)
        assert np.all(cur >= prev - 1e-12)
        prev = cur
    # undervaluation bias raises, overvaluation lowers, the failure probability
    f_neg = wf.failure_prob_given_seller(u_s, -0.03, 0.1, p)
    f_pos = wf.failure_prob_given_seller(u_s, 0.03, 0.1, p)
    assert np.all(f_neg > f_pos)
    # a larger haircut raises it
    assert np.all(wf.failure_prob_given_seller(u_s, 0.0, 0.1, wf.with_params(p, h1=0.5)) >
                  wf.failure_prob_given_seller(u_s, 0.0, 0.1, p))
    # non-financeable matches (a > 0): noise lowers the failure probability
    u_hi = np.array([0.10])
    assert wf.failure_prob_given_seller(u_hi, 0.0, 0.2, p) < wf.failure_prob_given_seller(u_hi, 0.0, 0.05, p)


def test_market_simulation_failure_rate_increasing_in_s():
    rng = np.random.default_rng(SEED)
    p = wf.WelfareParams()
    shocks = wf.draw_market_shocks(100_000, rng)
    rates = [wf.simulate_market(0.0, s, p, 100_000, rng, shocks=shocks)["failure_rate"] for s in (0.02, 0.05, 0.1, 0.2)]
    assert all(a < b for a, b in zip(rates, rates[1:]))
    exact = wf.simulate_market(0.0, 0.1, p, 100_000, rng, shocks=shocks)
    assert exact["failure_rate"] == pytest.approx(exact["failure_rate_exact"], abs=0.01)


def test_expected_loss_increasing_in_bias_and_spread():
    p = wf.WelfareParams()
    els = [wf.expected_loss(b, 0.1, p) for b in (-0.05, 0.0, 0.05, 0.1)]
    assert all(a < b for a, b in zip(els, els[1:]))
    els = [wf.expected_loss(0.0, s, p) for s in (0.0, 0.05, 0.1, 0.2, 0.3)]
    assert all(a < b for a, b in zip(els, els[1:]))


def test_systemic_mean_flat_and_dispersion_increasing_in_rho():
    rng = np.random.default_rng(SEED)
    p = wf.WelfareParams()
    N, n = 30, 20_000
    shocks = {"C": rng.standard_normal(n), "E": rng.standard_normal((n, N))}
    res = [wf.systemic_loss(r, 0.0, 0.1, p, N, n, rng, shocks=shocks) for r in (0.0, 0.25, 0.5, 0.75, 1.0)]
    means = [x["mean"] for x in res]
    assert max(means) - min(means) < 0.03 * np.mean(means)
    sds = [x["sd"] for x in res]
    covs = [x["cov_pair"] for x in res]
    assert all(a < b for a, b in zip(sds, sds[1:]))
    assert all(a <= b + 1e-12 for a, b in zip(covs, covs[1:]))