import numpy as np
import pytest
from scipy.stats import norm

from quantcore.models import sampling as smp


def test_next_power_of_two():
    assert [smp.next_power_of_two(n) for n in (1, 2, 3, 2000, 2048, 2049, 100_000)] == [1, 2, 4, 2048, 2048, 4096, 131072]


def test_standard_normal_sample_methods():
    for method in smp.SAMPLING_METHODS:
        z = smp.standard_normal_sample(1024, 5, method, seed=3)
        assert z.shape == (1024, 5) and np.isfinite(z).all()
    anti = smp.standard_normal_sample(1000, 3, "antithetic", seed=1)
    np.testing.assert_allclose(anti[:500], -anti[500:], atol=1e-15)  # exact zero mean by construction
    with pytest.raises(ValueError):
        smp.standard_normal_sample(1000, 3, "sobol")  # not a power of two
    with pytest.raises(ValueError):
        smp.standard_normal_sample(1024, 3, "bogus")
    # Sobol' points are deterministic for a seed and stratify the marginals far better than MC
    a = smp.standard_normal_sample(2048, 4, "sobol", seed=7)
    b = smp.standard_normal_sample(2048, 4, "sobol", seed=7)
    np.testing.assert_array_equal(a, b)
    lhs = smp.standard_normal_sample(4096, 2, "lhs", seed=7)
    counts = np.histogram(norm.cdf(lhs[:, 0]), bins=64, range=(0, 1))[0]
    assert counts.min() == counts.max() == 64  # exactly 64 points per stratum


def test_pca_root():
    rng = np.random.default_rng(0)
    L = rng.standard_normal((6, 6))
    cov = L @ L.T
    a = smp.pca_root(cov)
    np.testing.assert_allclose(a @ a.T, cov, atol=1e-10)
    vals, vecs = np.linalg.eigh(cov)
    # first column = largest-variance direction, columns orthogonal
    assert abs(abs(a[:, 0] @ vecs[:, -1]) - np.sqrt(vals[-1])) < 1e-9
    np.testing.assert_allclose(a.T @ a, np.diag(np.sort(vals)[::-1]), atol=1e-9)


def test_moment_match_exact_and_weighted():
    rng = np.random.default_rng(1)
    z = rng.standard_normal((500, 7)) * 1.3 + 0.2
    m = smp.moment_match(z)
    np.testing.assert_allclose(m.mean(axis=0), 0.0, atol=1e-13)
    np.testing.assert_allclose(m.T @ m / len(m), np.eye(7), atol=1e-12)
    p = rng.uniform(0.5, 2.0, 500)
    p /= p.sum()
    mw = smp.moment_match(z, p, target_variance=5.0)
    np.testing.assert_allclose(p @ mw, 0.0, atol=1e-12)
    np.testing.assert_allclose((mw * p[:, None]).T @ mw, 5.0 * np.eye(7), atol=1e-11)
    with pytest.raises(ValueError):
        smp.moment_match(rng.standard_normal((5, 7)))


def test_tilt_normal_sample_weights_bounded_and_tail_enriched():
    rng = np.random.default_rng(2)
    n, d = 4096, 8
    z = smp.standard_normal_sample(n, d, "sobol", seed=5)
    cov = np.eye(d) * 0.01
    root = smp.pca_root(cov)
    exposure = rng.uniform(1, 3, d) * 1000
    direction = smp.loss_direction(root, exposure)
    assert abs(np.linalg.norm(direction) - 1) < 1e-12
    beta = 0.95
    zt, p = smp.tilt_normal_sample(z, direction, norm.ppf(beta), mixture=0.5)
    assert p.shape == (n,) and abs(p.sum() - 1) < 1e-12 and (p > 0).all()
    assert p.max() <= 1.0 / (1 - 0.5) / n * 1.1  # likelihood ratios bounded by 1/(1-mixture), self-normalised
    ess = smp.effective_sample_size(p)
    assert 0.5 * n < ess < n
    # the loss of the reference exposure exceeds its VaR far more often than 5 % of the time
    loss_nominal = -(z @ root.T) @ exposure
    loss_tilted = -(zt @ root.T) @ exposure
    var = np.quantile(loss_nominal, beta)
    assert (loss_tilted > var).mean() > 0.2
    # ... while the weighted tail probability is still ~ 1 - beta (unbiased)
    assert abs(p[loss_tilted > var].sum() - (1 - beta)) < 0.01
    with pytest.raises(ValueError):
        smp.tilt_normal_sample(z, direction, 1.0, mixture=1.0)


def test_tilted_day_probabilities_and_resampling():
    rng = np.random.default_rng(3)
    scores = rng.standard_t(5, 2000)
    q, theta = smp.tilted_day_probabilities(scores, target_mean=0.7)
    assert theta > 0 and abs(q @ scores - 0.7) < 1e-8 and abs(q.sum() - 1) < 1e-12
    q0, theta0 = smp.tilted_day_probabilities(scores, target_mean=-1.0)
    assert theta0 == 0.0 and np.allclose(q0, 1 / 2000)
    with pytest.raises(ValueError):
        smp.tilted_day_probabilities(scores, target_mean=scores.max() + 1)
    idx, p, _theta = smp.tilted_resample(scores, 1000, 5, 0.7, seed=1, mixture=0.5)
    assert idx.shape == (1000, 5) and abs(p.sum() - 1) < 1e-12
    assert p.max() <= 2.0 / 1000 * 1.1  # likelihood ratios <= 1/(1-mixture), self-normalised
    # tilted half draws high scores more often
    assert scores[idx[:500]].mean() > scores[idx[500:]].mean() + 0.3


def test_resample_indices_methods():
    T, S, h = 300, 1024, 5
    mc = smp.resample_indices(T, S, h, "mc", seed=0)
    assert mc.shape == (S, h) and mc.min() >= 0 and mc.max() < T
    bal = smp.resample_indices(T, S, h, "balanced", seed=0)
    counts = np.bincount(bal.ravel(), minlength=T)
    assert counts.max() - counts.min() <= 1  # every observation used equally often
    order = np.argsort(np.random.default_rng(1).standard_normal(T))
    sob = smp.resample_indices(T, S, h, "sobol", seed=0, order=order)
    assert sob.shape == (S, h) and set(np.unique(sob)) <= set(range(T))
    ranks = smp.resample_indices(T, S, h, "sobol", seed=0)
    np.testing.assert_array_equal(order[ranks], sob)
    with pytest.raises(ValueError):
        smp.resample_indices(T, 1000, h, "sobol")


def test_scenario_hash():
    rng = np.random.default_rng(4)
    x = rng.standard_normal((50, 3))
    h = smp.scenario_hash(x)
    assert len(h) == 16 and h == smp.scenario_hash(x.copy())
    assert h != smp.scenario_hash(x + 1e-9)
    assert h != smp.scenario_hash(x, np.full(50, 0.02))