"""Units, bounds and closed forms of the match-value distributions."""
import numpy as np
from scipy import integrate

from src.model import TruncNormal, Exponential, Pareto, Weibull, seller_type_grid


def _check_dist(G, vs):
    for v in vs:
        # H(v) = int_v (1-G) dw  (numerical cross-check of the closed form)
        hi = G.hi if np.isfinite(G.hi) else np.inf
        num = integrate.quad(lambda w: float(G.surv(w)), v, hi, limit=300)[0]
        assert abs(float(G.H(v)) - num) < 1e-6, (G.name, v, float(G.H(v)), num)
        assert 0.0 <= float(G.cdf(v)) <= 1.0
        assert float(G.pdf(v)) >= 0.0
        assert float(G.H(v)) >= 0.0
        if float(G.surv(v)) > 1e-8:
            assert float(G.M(v)) >= 0.0
            assert float(G.cond_mean(v)) >= v - 1e-12
    # probability mass integrates to one
    lo = G.lo
    hi = G.hi if np.isfinite(G.hi) else np.inf
    mass = integrate.quad(lambda w: float(G.pdf(w)), lo, hi, limit=400)[0]
    assert abs(mass - 1.0) < 1e-5, (G.name, mass)
    # mean consistent
    mean_num = integrate.quad(lambda w: w * float(G.pdf(w)), lo, hi, limit=400)[0]
    assert abs(mean_num - G.mean()) < 1e-5


def test_truncnorm_closed_forms():
    G = TruncNormal(1.0, 0.1, 0.7, 1.3)
    _check_dist(G, [0.5, 0.7, 0.85, 1.0, 1.1, 1.29, 1.3, 1.5])
    assert G.is_ifr()


def test_exponential_pareto_weibull():
    _check_dist(Exponential(rate=10.0, lo=0.7), [0.5, 0.7, 0.9, 1.5])
    _check_dist(Pareto(a=3.0, lo=0.7), [0.5, 0.7, 0.9, 1.5])
    _check_dist(Weibull(shape=2.0, scale=0.2, lo=0.7), [0.5, 0.7, 0.9, 1.2])


def test_ifr_inequality_truncnorm():
    """Under IFR: g H <= (1-G)^2  (fact F2 used in Proposition 2)."""
    G = TruncNormal(1.0, 0.1, 0.7, 1.3)
    v = np.linspace(0.7, 1.29, 200)
    assert np.all(G.pdf(v) * G.H(v) <= G.surv(v) ** 2 + 1e-12)


def test_pareto_violates_ifr_inequality():
    G = Pareto(a=2.0, lo=0.7)
    v = np.linspace(0.8, 3.0, 50)
    assert np.all(G.pdf(v) * G.H(v) > G.surv(v) ** 2)


def test_seller_type_grid_mass_and_order():
    c, f = seller_type_grid(0.02, 0.3, 9, 2.0, 4.0)
    assert abs(f.sum() - 1.0) < 1e-12
    assert np.all(np.diff(c) > 0) and c[0] >= 0.02 and c[-1] <= 0.3
    # mean of the discretisation close to the Beta mean
    assert abs((c * f).sum() - (0.02 + 0.28 * 2.0 / 6.0)) < 1e-9


def test_samples_match_cdf():
    rng = np.random.default_rng(1)
    G = TruncNormal(1.0, 0.1, 0.7, 1.3)
    x = G.sample(rng, 200_000)
    assert abs(x.mean() - G.mean()) < 2e-3
    assert abs((x >= 1.0).mean() - float(G.surv(1.0))) < 5e-3