"""Tests for the sequential change-detection module.

The load-bearing ones are the cross-route pins.  Every run-length
number this topic publishes exists in three versions -- a closed form,
a Markov-chain quadrature, and a direct simulation of the recursion --
and the tests require them to agree, because a run-length formula
quoted from memory is exactly as treacherous as a Cramer-Rao bound
quoted from memory.

The GLR half carries the other standing rule: its threshold is
calibrated on the statistic the code actually computes (a maximum over
W correlated candidates, evaluated every sample), not on the textbook
marginal of one candidate.  ``test_naive_threshold_is_a_trap`` measures
the size of the gap so the page cannot quietly stop being true.
"""

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

from changedet import (arl_markov, arl_mc, arl_siegmund, cusum_detect,
                       cusum_stat, design_cusum, detection_delay_mc,
                       glr_arl0_mc, glr_detect, glr_stat,
                       glr_threshold_mc, glr_threshold_naive,
                       llr_increment, run_length_quantile, standardize)


# ----------------------------------------------------------------------
# Standardization and the LLR increment
# ----------------------------------------------------------------------

class TestStandardizeAndIncrement:
    def test_standardize_centres_and_scales(self):
        y = np.array([1.0, 3.0, 5.0])
        z = standardize(y, 3.0, 2.0)
        assert np.allclose(z, [-1.0, 0.0, 1.0])

    def test_standardize_rejects_nonpositive_sigma(self):
        with pytest.raises(ValueError):
            standardize([1.0], 0.0, 0.0)

    def test_increment_is_the_gaussian_log_ratio(self):
        # Computed independently from the two densities.
        rng = np.random.default_rng(0)
        z = rng.standard_normal(200)
        delta = 0.7
        direct = np.log(norm.pdf(z, delta, 1.0) / norm.pdf(z, 0.0, 1.0))
        assert np.allclose(llr_increment(z, delta), direct)

    def test_increment_mean_is_the_kl_divergence(self):
        # E_1[s] = delta^2/2 = KL(N(delta,1) || N(0,1)); E_0[s] = -delta^2/2.
        delta = 0.8
        rng = np.random.default_rng(1)
        z0 = rng.standard_normal(400000)
        z1 = z0 + delta
        assert abs(llr_increment(z0, delta).mean() + delta**2 / 2) < 0.01
        assert abs(llr_increment(z1, delta).mean() - delta**2 / 2) < 0.01

    def test_llr_units_are_delta_times_sigma_units(self):
        # The conversion the module refuses to apply silently.
        rng = np.random.default_rng(2)
        z = rng.standard_normal(500)
        delta = 1.4
        s = llr_increment(z, delta)
        # CUSUM on LLR increments, run by hand.
        g, acc = [], 0.0
        for v in s:
            acc = max(0.0, acc + v)
            g.append(acc)
        assert np.allclose(np.asarray(g), delta * cusum_stat(z, delta / 2))


# ----------------------------------------------------------------------
# The CUSUM recursion
# ----------------------------------------------------------------------

class TestCusumStat:
    def test_matches_the_naive_recursion(self):
        rng = np.random.default_rng(3)
        z = rng.standard_normal(500)
        s, ref = 0.0, []
        for v in z:
            s = max(0.0, s + v - 0.4)
            ref.append(s)
        assert np.allclose(cusum_stat(z, 0.4), ref)

    def test_page_identity_recursion_equals_max_over_change_times(self):
        # The claim that makes the CUSUM cheap: the O(1) recursion IS
        # the maximization over every candidate change time.
        rng = np.random.default_rng(4)
        z = rng.standard_normal(300)
        k = 0.5
        s = cusum_stat(z, k)
        cs = np.concatenate(([0.0], np.cumsum(z - k)))
        brute = np.array([max(0.0, max(cs[i + 1] - cs[i + 1 - j]
                                       for j in range(1, i + 2)))
                          for i in range(len(z))])
        assert np.allclose(brute, s)

    def test_lower_arm_mirrors_the_upper_one(self):
        rng = np.random.default_rng(5)
        z = rng.standard_normal(400)
        hi, lo = cusum_stat(z, 0.5, two_sided=True)
        hi_flipped = cusum_stat(-z, 0.5)
        assert np.allclose(lo, hi_flipped)

    def test_statistic_is_non_negative(self):
        rng = np.random.default_rng(6)
        assert np.all(cusum_stat(rng.standard_normal(1000), 0.5) >= 0.0)

    def test_rejects_negative_k(self):
        with pytest.raises(ValueError):
            cusum_stat([0.0], -0.1)


class TestCusumDetect:
    def test_no_alarm_when_the_statistic_stays_low(self):
        z = np.zeros(100)
        alarms, stat = cusum_detect(z, 0.5, 5.0)
        assert len(alarms) == 0
        assert np.all(stat == 0.0)

    def test_alarms_on_a_sustained_shift(self):
        z = np.concatenate((np.zeros(50), np.full(50, 2.0)))
        alarms, _ = cusum_detect(z, 0.5, 5.0)
        assert len(alarms) > 0
        # 2.0 - 0.5 = 1.5 per sample from a zero start: ceil(5/1.5) = 4.
        assert alarms[0] == 50 + 3

    def test_reset_produces_repeated_alarms(self):
        z = np.full(200, 2.0)
        with_reset, _ = cusum_detect(z, 0.5, 5.0, reset=True)
        without, _ = cusum_detect(z, 0.5, 5.0, reset=False)
        assert len(with_reset) < len(without)
        assert len(with_reset) == pytest.approx(200 / 4, rel=0.1)

    def test_two_sided_catches_a_downward_shift_one_sided_misses(self):
        z = np.concatenate((np.zeros(50), np.full(50, -2.0)))
        one, _ = cusum_detect(z, 0.5, 5.0)
        two, _ = cusum_detect(z, 0.5, 5.0, two_sided=True)
        assert len(one) == 0
        assert len(two) > 0

    def test_rejects_nonpositive_threshold(self):
        with pytest.raises(ValueError):
            cusum_detect([0.0], 0.5, 0.0)


# ----------------------------------------------------------------------
# Run length: the three routes must agree
# ----------------------------------------------------------------------

class TestArlRoutesAgree:
    @pytest.mark.parametrize("k,h", [(0.5, 4.0), (0.25, 8.0), (0.5, 3.0)])
    def test_arl0_markov_matches_simulation(self, k, h):
        rng = np.random.default_rng(11)
        mc, se = arl_mc(k, h, 0.0, trials=4000, rng=rng)
        ref = arl_markov(k, h, 0.0)
        assert abs(ref - mc) < 4 * se + 0.01 * mc

    @pytest.mark.parametrize("k,h", [(0.5, 4.0), (0.25, 8.0)])
    def test_arl1_markov_matches_simulation(self, k, h):
        rng = np.random.default_rng(12)
        mc, se = arl_mc(k, h, 2 * k, trials=4000, rng=rng)
        ref = arl_markov(k, h, 2 * k)
        assert abs(ref - mc) < 4 * se + 0.01 * mc

    def test_siegmund_matches_markov_in_the_small_shift_regime(self):
        # The regime a CUSUM is actually for.
        for k in (0.1, 0.25, 0.5):
            _, h = design_cusum(2 * k, 1000.0)
            rel = abs(arl_siegmund(k, h, 0.0) / arl_markov(k, h, 0.0) - 1)
            assert rel < 0.01, f"k={k}: {rel:.3%}"

    def test_siegmund_degrades_for_large_k(self):
        # The docstring's honesty claim, pinned in both directions:
        # small error at k = 0.5, materially wrong by k = 1.5.
        _, h_small = design_cusum(1.0, 1000.0)       # k = 0.5
        _, h_large = design_cusum(3.0, 1000.0)       # k = 1.5
        e_small = abs(arl_siegmund(0.5, h_small) / arl_markov(0.5, h_small) - 1)
        e_large = abs(arl_siegmund(1.5, h_large) / arl_markov(1.5, h_large) - 1)
        assert e_small < 0.02
        assert e_large > 0.10
        assert e_large > 10 * e_small

    def test_overshoot_correction_is_what_makes_it_work(self):
        # Same expression with b = h instead of h + 1.166.
        def wald(k, h):
            d = -k
            return (np.exp(-2 * d * h) + 2 * d * h - 1) / (2 * d * d)

        for k in (0.25, 0.5, 1.0):
            _, h = design_cusum(2 * k, 1000.0)
            uncorrected = wald(k, h) / arl_markov(k, h, 0.0)
            assert uncorrected < 0.6, f"k={k}: {uncorrected:.2f}"

    def test_markov_is_converged_at_the_default_cell_count(self):
        # The default (m = 200, Richardson on) against a 16x finer grid.
        # The loose case is h = 20, where the cells are coarsest
        # relative to the unit-variance increments; the docstring
        # quotes it, so pin it here.
        for k, h, tol in ((0.5, 5.0, 1e-4), (1.0, 2.9, 1e-4),
                          (0.25, 9.0, 5e-4), (0.1, 20.0, 2e-3)):
            default = arl_markov(k, h, 0.0)
            fine = arl_markov(k, h, 0.0, m=1600)
            assert abs(default / fine - 1) < tol, f"k={k}, h={h}"

    def test_richardson_beats_the_raw_quadrature(self):
        raw = arl_markov(0.5, 5.0, 0.0, m=200, richardson=False)
        extrap = arl_markov(0.5, 5.0, 0.0, m=200, richardson=True)
        truth = arl_markov(0.5, 5.0, 0.0, m=1600, richardson=True)
        assert abs(extrap - truth) < abs(raw - truth) / 5

    def test_arl_grows_with_threshold_and_falls_with_shift(self):
        assert arl_markov(0.5, 6.0) > arl_markov(0.5, 4.0)
        assert arl_markov(0.5, 5.0, 2.0) < arl_markov(0.5, 5.0, 1.0)
        assert arl_markov(0.5, 5.0, 1.0) < arl_markov(0.5, 5.0, 0.0)

    def test_siegmund_zero_drift_limit_is_continuous(self):
        # D -> 0 is the (h + 1.166)^2 branch; approach it from both sides.
        k = 0.5
        near_lo = arl_siegmund(k, 4.0, k - 1e-7)
        near_hi = arl_siegmund(k, 4.0, k + 1e-7)
        at = arl_siegmund(k, 4.0, k)
        assert abs(near_lo - at) < 1e-3 * at
        assert abs(near_hi - at) < 1e-3 * at
        assert abs(at - (4.0 + 1.166) ** 2) < 1e-9

    def test_rejects_bad_parameters(self):
        with pytest.raises(ValueError):
            arl_siegmund(0.0, 4.0)
        with pytest.raises(ValueError):
            arl_markov(0.5, -1.0)
        with pytest.raises(ValueError):
            arl_markov(0.5, 4.0, m=4)


class TestRunLengthDistribution:
    def test_h0_run_length_is_close_to_geometric(self):
        # ARL0 is a mean, not a safe horizon: the simulated quantiles
        # must follow -ARL0 ln(1-q).
        rng = np.random.default_rng(13)
        k, h, trials = 0.5, 4.0, 8000
        lengths = np.empty(trials)
        s = np.zeros(trials)
        alive = np.ones(trials, dtype=bool)
        n = 0
        while alive.any() and n < 100000:
            n += 1
            s = np.where(alive, np.maximum(0.0, s + rng.standard_normal(trials)
                                           - k), 0.0)
            lengths = np.where(alive, n, lengths)
            alive &= s < h
        arl0 = lengths.mean()
        assert abs(lengths.std() / arl0 - 1.0) < 0.05
        for q in (0.25, 0.5, 0.9):
            emp = np.quantile(lengths, q)
            assert abs(emp / run_length_quantile(arl0, q) - 1.0) < 0.06

    def test_quantile_helper_rejects_bad_q(self):
        with pytest.raises(ValueError):
            run_length_quantile(100.0, 1.0)


class TestDesignCusum:
    def test_round_trips_through_the_markov_arl(self):
        for delta, target in ((1.0, 500.0), (0.5, 2000.0), (2.0, 1000.0)):
            k, h = design_cusum(delta, target)
            assert k == pytest.approx(delta / 2)
            assert arl_markov(k, h, 0.0) == pytest.approx(target, rel=1e-4)

    def test_design_holds_up_under_simulation(self):
        rng = np.random.default_rng(14)
        k, h = design_cusum(1.0, 500.0)
        mc, se = arl_mc(k, h, 0.0, trials=4000, rng=rng)
        assert abs(mc - 500.0) < 4 * se + 5.0

    def test_rejects_impossible_requests(self):
        with pytest.raises(ValueError):
            design_cusum(0.0, 100.0)
        with pytest.raises(ValueError):
            design_cusum(1.0, 0.5)


class TestDetectionDelay:
    def test_warm_started_delay_matches_the_markov_arl1(self):
        # arl_markov(delta) is a zero-start quantity; a warm detector is
        # slightly faster, so the simulation should land at or just
        # below it, never above.
        rng = np.random.default_rng(15)
        k, h = design_cusum(1.0, 500.0)

        def detector(z):
            a, _ = cusum_detect(z, k, h, reset=False)
            return int(a[0]) if len(a) else None

        mean, se = detection_delay_mc(detector, 1.0, trials=600,
                                      burn_in=100, rng=rng, max_samples=500)
        predicted = arl_markov(k, h, 1.0)
        assert mean < predicted + 4 * se
        assert mean > 0.6 * predicted

    def test_delay_falls_as_the_shift_grows(self):
        rng = np.random.default_rng(16)
        k, h = design_cusum(1.0, 500.0)

        def detector(z):
            a, _ = cusum_detect(z, k, h, reset=False)
            return int(a[0]) if len(a) else None

        d1, _ = detection_delay_mc(detector, 0.8, trials=300, burn_in=50,
                                   rng=rng, max_samples=500)
        d2, _ = detection_delay_mc(detector, 2.0, trials=300, burn_in=50,
                                   rng=rng, max_samples=500)
        assert d2 < d1


# ----------------------------------------------------------------------
# GLR
# ----------------------------------------------------------------------

class TestGlrStat:
    def test_matches_the_brute_force_definition(self):
        rng = np.random.default_rng(20)
        z = rng.standard_normal(400)
        w = 25
        cs = np.concatenate(([0.0], np.cumsum(z)))
        brute = np.array([max((cs[i + 1] - cs[i + 1 - j]) ** 2 / (2 * j)
                              for j in range(1, min(w, i + 1) + 1))
                          for i in range(len(z))])
        assert np.allclose(brute, glr_stat(z, w))

    def test_window_one_is_the_half_squared_sample(self):
        rng = np.random.default_rng(21)
        z = rng.standard_normal(300)
        assert np.allclose(glr_stat(z, 1), z ** 2 / 2)

    def test_window_one_tail_is_exactly_chi2_1(self):
        # The independent special case: with W = 1 the statistic has a
        # closed form that owes nothing to the machinery above.
        rng = np.random.default_rng(22)
        g = glr_stat(rng.standard_normal(400000), 1)
        for h in (1.0, 2.0, 3.0):
            emp = np.mean(g > h)
            closed = chi2.sf(2 * h, 1)
            assert abs(emp - closed) < 4 * np.sqrt(closed / 400000) + 1e-4
            assert closed == pytest.approx(2 * norm.sf(np.sqrt(2 * h)))

    def test_is_two_sided(self):
        rng = np.random.default_rng(23)
        z = rng.standard_normal(300)
        assert np.allclose(glr_stat(z, 20), glr_stat(-z, 20))

    def test_grows_with_the_window(self):
        rng = np.random.default_rng(24)
        z = rng.standard_normal(2000)
        assert np.all(glr_stat(z, 40) >= glr_stat(z, 10) - 1e-12)

    def test_handles_empty_input_and_rejects_bad_window(self):
        assert len(glr_stat([], 5)) == 0
        with pytest.raises(ValueError):
            glr_stat([1.0], 0)


class TestGlrDetect:
    def test_finds_a_change_of_unannounced_size(self):
        rng = np.random.default_rng(25)
        z = np.concatenate((rng.standard_normal(200),
                            rng.standard_normal(200) + 1.5))
        alarms, _ = glr_detect(z, 50, 8.0)
        assert len(alarms) > 0
        assert 200 <= alarms[0] < 230

    def test_reset_restarts_the_history(self):
        rng = np.random.default_rng(26)
        z = rng.standard_normal(600) + 1.0
        with_reset, _ = glr_detect(z, 30, 8.0, reset=True)
        without, _ = glr_detect(z, 30, 8.0, reset=False)
        assert len(with_reset) < len(without)

    def test_rejects_nonpositive_threshold(self):
        with pytest.raises(ValueError):
            glr_detect([0.0], 5, 0.0)


class TestGlrThreshold:
    def test_naive_threshold_matches_its_own_closed_form(self):
        h = glr_threshold_naive(1e-3)
        assert h == pytest.approx(chi2.isf(1e-3, 1) / 2)
        assert 2 * norm.sf(np.sqrt(2 * h)) == pytest.approx(1e-3)

    def test_naive_threshold_is_a_trap(self):
        # Gotcha 85 in its natural habitat: the single-candidate
        # marginal is exact, and calibrating with it still runs the
        # deployed detector several times hot.  Both the direction and
        # the rough size are pinned.
        rng = np.random.default_rng(27)
        h = glr_threshold_naive(1e-3)
        arl_w1, se1 = glr_arl0_mc(h, 1, trials=300, rng=rng)
        arl_w50, se50 = glr_arl0_mc(h, 50, trials=300, rng=rng)
        assert abs(arl_w1 - 1000.0) < 4 * se1 + 20.0      # W = 1 is honest
        assert 2.0 < 1000.0 / arl_w50 < 6.0               # W = 50 is not

    def test_naive_threshold_rejects_bad_p(self):
        with pytest.raises(ValueError):
            glr_threshold_naive(0.0)

    def test_monte_carlo_calibration_hits_its_target(self):
        rng = np.random.default_rng(28)
        target, w = 500.0, 40
        h = glr_threshold_mc(target, w, trials=150, rng=rng)
        check = np.random.default_rng(29)
        arl, se = glr_arl0_mc(h, w, trials=500, rng=check)
        assert abs(arl / target - 1.0) < 0.25
        # And it must land above the naive threshold, not below.
        assert h > glr_threshold_naive(1.0 / target)


class TestCusumVersusGlr:
    def test_glr_beats_a_misdesigned_cusum_on_an_unexpected_shift(self):
        # The trade the page is built on: the CUSUM is faster at the
        # shift it was told to expect, the GLR is steadier across
        # shifts it was not.
        rng = np.random.default_rng(30)
        k, h_c = design_cusum(0.5, 500.0)     # designed for a small shift
        h_g = 6.5                             # roughly ARL0-matched, W = 50

        def cusum_det(z):
            a, _ = cusum_detect(z, k, h_c, reset=False)
            return int(a[0]) if len(a) else None

        def glr_det(z):
            a, _ = glr_detect(z, 50, h_g, reset=False)
            return int(a[0]) if len(a) else None

        arl0_glr, _ = glr_arl0_mc(h_g, 50, trials=200, rng=rng)
        assert 200.0 < arl0_glr < 1500.0, "the comparison must be near-fair"

        big = 3.0
        d_cusum, _ = detection_delay_mc(cusum_det, big, trials=300,
                                        burn_in=50, rng=rng, max_samples=400)
        d_glr, _ = detection_delay_mc(glr_det, big, trials=300,
                                      burn_in=50, rng=rng, max_samples=400)
        assert d_glr < d_cusum