"""Tests for post-hoc E2 analysis (analysis.py) and the tail-band metric.

Encodes the work-order's verified numbers: the operational g* on a synthetic E2 set, the
stationary reducer, and the band-wise survival ordering on a grounded Zipf tail.
"""

import numpy as np
import pandas as pd
import pytest

from knowledge.analysis import reduce_to_stationary, critical_grounding
from knowledge.metrics import tail_band_metrics


def _H_eq(n, m, Hs):
    return Hs * m * (2 * n + m - 1) / (n + 2 * n * m + m * m)


def _synthetic_E2(n=200, Hs=0.95, reps=100, seed=1):
    gs = [0.0, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.4]
    rng = np.random.default_rng(seed)
    rows = []
    for g in gs:
        m = 0 if g == 0 else int(round(g * n / (1 - g)))
        Htrue = 0.10 if g == 0 else _H_eq(n, m, Hs)   # g=0: finite-time artifact
        for s in range(reps):
            rows.append({"g": g, "seed": s, "heterozygosity": Htrue + rng.normal(0, 0.004)})
    return pd.DataFrame(rows)


def test_critical_grounding_matches_known_crossing():
    stat = _synthetic_E2()
    r = critical_grounding(stat, H_star=0.95, frac=0.95, seed=7)
    assert r["status"] == "ok"
    assert 0.03 < r["g_star"] < 0.07                   # ~0.047 for frac=0.95
    assert r["ci_low"] <= r["g_star"] <= r["ci_high"]
    r90 = critical_grounding(stat, H_star=0.95, frac=0.90, seed=7)
    assert r90["g_star"] < r["g_star"]                 # lower bar -> smaller g*


def test_reduce_to_stationary_recovers_plateau():
    # constant plateau + noise -> mean ~ plateau
    rng = np.random.default_rng(0)
    rows = []
    for g, plateau in [(0.02, 0.843), (0.05, 0.908)]:
        for s in range(20):
            for t in range(300):
                v = (0.95 if t < 50 else plateau) + rng.normal(0, 0.003)
                rows.append({"g": g, "seed": s, "generation": t, "heterozygosity": v})
    red = reduce_to_stationary(pd.DataFrame(rows), replicate_col="seed", last_frac=0.33)
    means = red.groupby("g")["heterozygosity"].mean()
    assert means[0.02] == pytest.approx(0.843, abs=0.01)
    assert means[0.05] == pytest.approx(0.908, abs=0.01)


def test_tail_band_deep_below_shallow():
    # grounded dynamics on a Zipf tail: deepest band <= shallowest band, replicate-avg
    K = 1000
    s = 1.1
    w = 1.0 / np.arange(1, K + 1) ** s
    p_star = w / w.sum()
    tail_mask = p_star < 1e-3
    n, m = 200, 10
    FA = np.zeros(4)
    for r in range(20):
        rng = np.random.default_rng(1000 + r)
        p = p_star.copy()
        for _ in range(600):
            c = rng.multinomial(n, p) + rng.multinomial(m, p_star)
            p = c / c.sum()
        fa, _ = tail_band_metrics(p, p_star, tail_mask, n_bands=4)
        FA += fa
    FA /= 20
    assert FA[0] <= FA[-1]                              # deepest no better than shallowest
    assert FA[-1] > FA[0]                               # and strictly worse on average