"""
Unit tests for src/diversity.py.  Run: python3 -m pytest src/test_diversity.py -q
                                  or: python3 src/test_diversity.py
"""
import numpy as np

from diversity import (
    cosine_kernel,
    greedy_diverse_subset,
    l2_normalize,
    logdet_volume,
    marginal_contributions,
    pairwise_deviation,
    zscore,
)

RNG = np.random.default_rng(0)


def _orth(G, d):
    """G mutually orthonormal rows in R^d."""
    Q, _ = np.linalg.qr(RNG.standard_normal((d, G)))
    return Q.T[:G]


# ---------------------------------------------------------------- deviation
def test_deviation_duplicates_are_zero():
    """Exact duplicates: every pair distance is 0, so d_i == 0 for all."""
    e = l2_normalize(RNG.standard_normal((1, 32)))
    E = np.repeat(e, 8, axis=0)
    d = pairwise_deviation(E)
    assert np.allclose(d, 0.0, atol=1e-9), d


def test_deviation_orthogonal_is_one():
    """Orthogonal rows: cos = 0 for every pair, so d_i == 1 exactly."""
    E = _orth(8, 64)
    d = pairwise_deviation(E)
    assert np.allclose(d, 1.0, atol=1e-9), d


def test_deviation_singleton_and_empty():
    assert pairwise_deviation(np.zeros((1, 8))).shape == (1,)
    assert pairwise_deviation(np.zeros((1, 8)))[0] == 0.0
    assert pairwise_deviation(np.zeros((0, 8))).shape == (0,)


def test_deviation_flags_the_odd_one_out():
    """7 near-identical + 1 far: the outlier must have the highest d_i."""
    base = l2_normalize(RNG.standard_normal((1, 64)))
    tight = l2_normalize(np.repeat(base, 7, axis=0) + 0.01 * RNG.standard_normal((7, 64)))
    far = l2_normalize(RNG.standard_normal((1, 64)))
    E = np.vstack([tight, far])
    d = pairwise_deviation(E)
    assert d.argmax() == 7, d


# ------------------------------------------------------- marginal / logdet
def test_marginal_duplicate_is_large_negative():
    """A duplicated direction is already spanned -> dropping one costs ~nothing,
    so the PRESENT duplicate's marginal is driven to ~log(eps), very negative."""
    e = l2_normalize(RNG.standard_normal((1, 32)))
    E = np.vstack([np.repeat(e, 2, axis=0), _orth(4, 32)])
    m = marginal_contributions(E)
    # the two duplicates (rows 0,1) are the least valuable members
    assert m[0] < -3.0 and m[1] < -3.0, m
    assert m[:2].max() < m[2:].min(), m


def test_marginal_orthogonal_is_near_zero_and_uniform():
    """Orthonormal rows: L = (1+eps)I, dropping any row costs log(1+eps) ~ 0.
    'High m_i for all' in the sense of at-ceiling and symmetric."""
    E = _orth(8, 64)
    m = marginal_contributions(E)
    assert np.allclose(m, m[0], atol=1e-9), m
    assert abs(m[0] - np.log1p(1e-3)) < 1e-6, m[0]


def test_marginal_is_bounded_above_by_zero_ish():
    """logdet is monotone under adding a row with unit norm + jitter, so
    m_i can never exceed log(1+eps)."""
    E = l2_normalize(RNG.standard_normal((12, 64)))
    m = marginal_contributions(E)
    assert m.max() <= np.log1p(1e-3) + 1e-9, m.max()


def test_logdet_ordering_dup_lt_spread_lt_orthogonal():
    e = l2_normalize(RNG.standard_normal((1, 64)))
    dup = np.repeat(e, 8, axis=0)
    spread = l2_normalize(RNG.standard_normal((8, 64)))
    orth = _orth(8, 64)
    assert logdet_volume(dup) < logdet_volume(spread) < logdet_volume(orth)


def test_kernel_is_psd_even_with_duplicates():
    e = l2_normalize(RNG.standard_normal((1, 16)))
    L = cosine_kernel(np.repeat(e, 6, axis=0))
    assert np.linalg.eigvalsh(L).min() > 0, "jitter failed to make L PD"
    assert np.isfinite(logdet_volume(np.repeat(e, 6, axis=0)))


# --------------------------------------- THE E1-vs-E2 HYPOTHESIS, AS A TEST
def test_two_clusters_fool_deviation_but_not_logdet():
    """This is the claim E2 rests on, so it gets asserted rather than assumed.

    Config A: two tight antipodal clusters of 4 (rank ~2, 'diverse' only in
              the sense that half the samples are far from the other half).
    Config B: 8 genuinely spread directions (rank ~8).

    Pairwise deviation cannot tell these apart -- mean pairwise distance for A
    is actually HIGHER, because antipodal pairs sit at cos = -1. Log-det sees
    straight through it: A occupies a 2-dimensional subspace.
    """
    u, v = _orth(2, 64)
    jit = 0.01
    A = l2_normalize(np.vstack([
        np.repeat(u[None], 4, axis=0) + jit * RNG.standard_normal((4, 64)),
        np.repeat(-u[None], 4, axis=0) + jit * RNG.standard_normal((4, 64)),
    ]))
    B = _orth(8, 64)

    dev_A, dev_B = pairwise_deviation(A).mean(), pairwise_deviation(B).mean()
    vol_A, vol_B = logdet_volume(A), logdet_volume(B)

    # deviation RANKS THE DEGENERATE SET HIGHER -- the failure mode, reproduced
    assert dev_A > dev_B, (dev_A, dev_B)
    # log-det correctly ranks the spread set far higher
    assert vol_B > vol_A + 10.0, (vol_A, vol_B)

    # and per-sample: in A every member is redundant (its twin covers it),
    # so marginal contributions are uniformly terrible
    assert marginal_contributions(A).max() < -3.0
    assert marginal_contributions(B).min() > -1e-3
    _ = v  # second basis vector unused, kept for clarity of construction


# ------------------------------------------------------------------ zscore
def test_zscore_constant_input_is_zeros_not_nan():
    """A constant reward column must degrade to 0, never NaN -- otherwise it
    poisons the whole advantage tensor."""
    z = zscore(np.full(8, 3.7))
    assert np.all(np.isfinite(z)) and np.allclose(z, 0.0)


def test_zscore_standardizes():
    z = zscore(RNG.standard_normal(64))
    assert abs(z.mean()) < 1e-9 and abs(z.std() - 1.0) < 1e-9


# ------------------------------------------------------------ greedy subset
def test_greedy_avoids_duplicates_when_quality_is_flat():
    """With flat quality, selection is pure logdet: must not pick both dupes."""
    e = l2_normalize(RNG.standard_normal((1, 64)))
    E = np.vstack([np.repeat(e, 3, axis=0), _orth(3, 64)])  # rows 0,1,2 identical
    sel = greedy_diverse_subset(np.ones(6), E, k=3, lam=1.0)
    assert len(sel) == 3
    assert len(set(sel) & {0, 1, 2}) <= 1, sel


def test_greedy_respects_quality_when_lambda_is_zero():
    E = l2_normalize(RNG.standard_normal((10, 64)))
    q = np.arange(10, dtype=float)
    sel = greedy_diverse_subset(q, E, k=3, lam=0.0)
    assert sorted(sel) == [7, 8, 9], sel


def test_greedy_k_larger_than_pool():
    E = l2_normalize(RNG.standard_normal((3, 32)))
    assert len(greedy_diverse_subset(np.ones(3), E, k=8)) == 3


if __name__ == "__main__":
    import sys, traceback
    fns = [(n, f) for n, f in sorted(globals().items())
           if n.startswith("test_") and callable(f)]
    bad = 0
    for n, f in fns:
        try:
            f()
            print(f"  PASS  {n}")
        except Exception:
            bad += 1
            print(f"  FAIL  {n}")
            traceback.print_exc()
    print(f"\n{len(fns)-bad}/{len(fns)} passed")
    sys.exit(1 if bad else 0)