"""Tests for the integer drift monitor that models the embedded C.

These pin the claims the embedded page makes about the firmware: the
accuracy of the integer logarithm, the fact that Q8 costs nothing at the
word lengths used, the range the accumulator actually reaches (so the
int16 choice is justified rather than hoped), and the two guards around
a zero-energy block.
"""

import numpy as np
import pytest

from intcusum import (LN2_Q12, LN_MAX_ABS_ERR_NATS, Q8_MAX,
                      IntDriftMonitor, adc_block_energies, int_ln_q12)


class TestIntegerLog:
    def test_accuracy_over_the_adc_range(self):
        # 64 samples of a 12-bit ADC: energies from a few counts to ~1e11.
        vals = np.unique(np.rint(np.exp(np.linspace(np.log(64.0),
                                                    np.log(1e11), 20000))))
        err = np.array([int_ln_q12(int(v)) / 4096.0 - np.log(float(v))
                        for v in vals])
        assert np.abs(err).max() < LN_MAX_ABS_ERR_NATS
        # And the constant is not slack: it is close to the real worst case.
        assert np.abs(err).max() > 0.5 * LN_MAX_ABS_ERR_NATS

    def test_error_is_small_against_the_noise_it_measures(self):
        # The CFAR's energy removes the block mean, so the in-control
        # log-energy spread is sqrt(trigamma((M-1)/2)) with M = 64.
        from scipy.special import polygamma
        sd = np.sqrt(polygamma(1, 31.5))
        assert LN_MAX_ABS_ERR_NATS / sd < 0.02, "well under 2% of a sigma"

    def test_exact_at_powers_of_two(self):
        for p in range(1, 40):
            got = int_ln_q12(2 ** p) / 4096.0
            assert abs(got - p * np.log(2)) < 2 * LN_MAX_ABS_ERR_NATS

    def test_ln2_constant_is_right(self):
        assert LN2_Q12 == round(np.log(2) * 4096)

    def test_zero_energy_is_floored_not_nonsense(self):
        # ln(0) has no representation; a muted mic produces it anyway.
        assert int_ln_q12(0) == int_ln_q12(1)
        assert abs(int_ln_q12(1)) <= 1, "ln(1) = 0 to within a Q12 count"

    def test_rejects_negative(self):
        with pytest.raises(ValueError):
            int_ln_q12(-1)

    def test_monotone_to_within_one_count(self):
        # Truncation in the Q12 Horner steps makes the result dither by a
        # single count, so it is NOT strictly monotone; it is monotone to
        # within 1/4096 nat, which is a quarter of the polynomial's own
        # error and 0.14% of a log-energy sigma.  Pin the real property.
        vals = np.array([int_ln_q12(v) for v in range(1, 200000)])
        d = np.diff(vals)
        assert d.min() == -1, "one count of truncation dither, no more"
        assert np.mean(d < 0) < 0.05


def _design(sd, k, h):
    return dict(k_q8=int(round(k * 256)), h_q8=int(round(h * 256)),
                inv_sd_q8=int(round(256.0 / sd)))


class TestIntDriftMonitor:
    def setup_method(self):
        from scipy.special import polygamma
        self.sd = float(np.sqrt(polygamma(1, 31.5)))
        rng = np.random.default_rng(1)
        train = adc_block_energies(4000, 1000.0, rng)
        self.mu_q12 = int(round(np.mean([int_ln_q12(int(e)) for e in train])))

    def test_quiet_input_rarely_alarms(self):
        rng = np.random.default_rng(2)
        d = IntDriftMonitor(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 6.0))
        e = adc_block_energies(20000, 1000.0, rng)
        alarms = d.run(int(v) for v in e)
        assert len(alarms) < 40, "h = 6 must be much quieter than 1-in-500"

    def test_detects_a_1db_drift_in_both_directions(self):
        for direction, gain in ((1, 10 ** (1 / 20)), (-1, 10 ** (-1 / 20))):
            rng = np.random.default_rng(3)
            d = IntDriftMonitor(mu_q12=self.mu_q12,
                                **_design(self.sd, 0.65, 6.0))
            for v in adc_block_energies(300, 1000.0, rng):
                d.push(int(v))
            d.hi = d.lo = 0
            hit = None
            for i, v in enumerate(adc_block_energies(400, 1000.0 * gain, rng)):
                if d.push(int(v)):
                    hit = (i, d)
                    break
            assert hit is not None, f"missed a 1 dB drift ({direction})"
            assert hit[0] < 120, "and it must be quick"

    def test_accumulator_stays_far_inside_int16(self):
        # The justification for int16: measure the high-water mark.
        rng = np.random.default_rng(4)
        d = IntDriftMonitor(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 9.3))
        for v in adc_block_energies(60000, 1000.0, rng):
            d.push(int(v))
        assert d.peak < Q8_MAX // 2, "int16 must not be marginal"
        assert d.peak > d.h_q8 // 2, "and the test must actually exercise it"

    def test_a_zero_energy_block_cannot_overflow_the_accumulator(self):
        # Without the saturating add this is where int16 wraps.
        rng = np.random.default_rng(5)
        d = IntDriftMonitor(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 9.3))
        for v in adc_block_energies(200, 1000.0, rng):
            d.push(int(v))
        for _ in range(50):
            d.push(0)
        assert 0 <= d.hi <= Q8_MAX and 0 <= d.lo <= Q8_MAX
        assert d.peak <= Q8_MAX

    def test_a_zero_energy_block_raises_the_downward_alarm(self):
        rng = np.random.default_rng(6)
        d = IntDriftMonitor(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 9.3))
        for v in adc_block_energies(200, 1000.0, rng):
            d.push(int(v))
        assert d.push(0) == -1, "a dead input is a downward change"

    def test_rebaselining_stops_the_endless_alarm_after_a_step(self):
        rng = np.random.default_rng(7)
        cfg = dict(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 6.0))
        post = adc_block_energies(3000, 1000.0 * 10 ** (1 / 20), rng)

        adapting = IntDriftMonitor(rebase_n=128, rebaseline=True, **cfg)
        frozen = IntDriftMonitor(rebase_n=128, rebaseline=False, **cfg)
        n_adapting = len(adapting.run(int(v) for v in post))
        n_frozen = len(frozen.run(int(v) for v in post))
        # A stale baseline means the statistic re-crosses h forever.
        assert n_frozen > 5 * n_adapting, "re-baselining must pay"
        assert n_frozen > 200, "and the frozen detector must really be stuck"

    def test_the_energy_interface_is_immune_to_the_adc_pedestal(self):
        # The CFAR removes the block mean, so its energy numerator is
        # EXACTLY invariant to the mid-rail offset.  That is the whole
        # reason it is safe to feed this monitor.
        a = adc_block_energies(400, 100.0, np.random.default_rng(24),
                               dc_lsb=2048)
        b = adc_block_energies(400, 100.0, np.random.default_rng(24),
                               dc_lsb=0)
        assert np.array_equal(a, b), "mean removal cancels the offset exactly"

        # And the raw sum of squares, which is what a naive hook-up would
        # pass instead, is destroyed by that same pedestal.
        rng = np.random.default_rng(24)
        x = np.rint(100.0 * rng.standard_normal((400, 64))).astype(np.int64)
        raw_at_midrail = np.sum((x + 2048) ** 2, axis=1)
        sd_raw = float(np.std(np.log(raw_at_midrail.astype(float))))
        sd_ok = float(np.std(np.log(a.astype(float))))
        assert sd_raw < 0.1 * sd_ok, \
            "the pedestal must swamp the raw path, or this test proves nothing"

    def test_the_spread_is_the_mean_removed_one(self):
        # M - 1 degrees of freedom, not M: the block mean was estimated.
        # Getting this wrong is a 0.8% error in sigma and about 17% in ARL0.
        from scipy.special import polygamma
        e = adc_block_energies(60000, 1000.0, np.random.default_rng(25))
        sd = float(np.std(np.log(e.astype(float))))
        good = float(np.sqrt(polygamma(1, 31.5)))     # (M-1)/2
        bad = float(np.sqrt(polygamma(1, 32.0)))      # M/2, the trap
        assert abs(sd / good - 1) < 0.01, "must match the M-1 closed form"
        assert abs(sd - good) < abs(sd - bad), "and must be closer to it than to M/2"

    def test_the_ring_holds_q8_because_raw_q12_would_not_fit(self):
        # The C declares int16_t hist[REBASE_N].  A raw Q12 log-energy is
        # about 2.3x too big for that, which is why drift_push stores
        # ln_q12 >> 4.  Both halves of that claim are pinned here.
        rng = np.random.default_rng(21)
        e = [int(v) for v in adc_block_energies(300, 1000.0, rng)]
        assert max(int_ln_q12(v) for v in e) > 32767, \
            "if raw Q12 fitted int16 the shift would be pointless"
        d = IntDriftMonitor(mu_q12=self.mu_q12, rebase_n=128,
                            **_design(self.sd, 0.65, 6.0))
        d.run(e)
        assert all(-32768 <= v <= 32767 for v in d.hist), \
            "every stored entry must fit the int16 ring the firmware has"

    def test_a_rebaselined_baseline_is_quantized_to_the_q8_grid(self):
        # (sum >> REBASE_SHIFT) << 4 can only land on multiples of 16.
        rng = np.random.default_rng(22)
        cfg = dict(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 4.0))
        d = IntDriftMonitor(rebase_n=16, rebaseline=True, **cfg)
        post = adc_block_energies(3000, 1000.0 * 10 ** (2 / 20), rng)
        seen = set()
        for v in post:
            d.push(int(v))
            seen.add(d.mu_q12)
        assert len(seen) > 1, "the baseline must actually have moved"
        assert all(m % 16 == 0 for m in seen if m != self.mu_q12), \
            "a re-baselined mu is reconstructed from Q8 and so is a multiple of 16"

    def test_refuses_to_rebaseline_until_the_ring_is_clean(self):
        # drift_push sets d->filled = 0 after re-baselining, so a ring
        # still holding pre-change blocks cannot pull the baseline back.
        # Without it the detector re-baselines on EVERY alarm, which is
        # the policy the embedded page explicitly disclaims.
        rng = np.random.default_rng(23)
        n = 16
        cfg = dict(mu_q12=self.mu_q12, **_design(self.sd, 0.65, 4.0))
        d = IntDriftMonitor(rebase_n=n, rebaseline=True, **cfg)
        quiet = adc_block_energies(n, 1000.0, rng)
        loud = adc_block_energies(2000, 1000.0 * 10 ** (3 / 20), rng)

        for v in quiet:
            d.push(int(v))
        assert d.filled == n, "the ring must start full"

        changes, mu_prev = [], d.mu_q12
        for i, v in enumerate(loud):
            d.push(int(v))
            if d.mu_q12 != mu_prev:
                changes.append(i)
                mu_prev = d.mu_q12
        assert len(changes) >= 2, "the test needs at least two re-baselines"
        gaps = np.diff(changes)
        assert gaps.min() >= n, \
            "no two re-baselines may be closer than a full ring of fresh blocks"

    def test_rebase_n_must_be_a_power_of_two(self):
        with pytest.raises(ValueError):
            IntDriftMonitor(rebase_n=100, mu_q12=0, **_design(self.sd, 0.6, 5))

    def test_rejects_bad_design_constants(self):
        with pytest.raises(ValueError):
            IntDriftMonitor(mu_q12=0, **_design(self.sd, 0.6, 0.0))


class TestIntegerMatchesFloat:
    def test_q8_costs_nothing_against_a_float_cusum(self):
        # The embedded page's claim that fixed point is a non-event at
        # these word lengths: same stream, same design, same alarms.
        from changedet import cusum_detect, standardize
        from scipy.special import polygamma

        sd_exact = float(np.sqrt(polygamma(1, 31.5)))
        rng = np.random.default_rng(8)
        train = adc_block_energies(4000, 1000.0, rng)
        ln_train = np.array([int_ln_q12(int(e)) / 4096.0 for e in train])
        mu, sd = ln_train.mean(), ln_train.std()
        assert abs(sd / sd_exact - 1) < 0.05

        e = adc_block_energies(30000, 1000.0, rng)
        k, h = 0.65, 6.0
        # Re-baselining OFF: this test isolates the ARITHMETIC.  With it on
        # the integer detector alarms about twice as often, and that is
        # adaptation, not fixed point -- see the re-baselining test below
        # and the embedded page's window-length table.
        d = IntDriftMonitor(mu_q12=int(round(mu * 4096)),
                            rebaseline=False, **_design(sd, k, h))
        int_alarms = [i for i, _ in d.run(int(v) for v in e)]

        ln_f = np.array([int_ln_q12(int(v)) / 4096.0 for v in e])
        f_alarms, _ = cusum_detect(standardize(ln_f, mu, sd), k, h,
                                   two_sided=True)
        assert len(f_alarms) > 3, "the comparison must have something to compare"
        # Same stream, same design: Q8 must reproduce the float alarms.
        assert int_alarms == [int(a) for a in f_alarms]

    def test_rebaselining_is_what_moves_the_alarm_count_not_q8(self):
        from changedet import cusum_detect, standardize
        rng = np.random.default_rng(9)
        from scipy.special import polygamma
        sd = float(np.sqrt(polygamma(1, 31.5)))
        train = adc_block_energies(4000, 1000.0, rng)
        mu = np.mean([int_ln_q12(int(v)) / 4096.0 for v in train])
        e = [int(v) for v in adc_block_energies(30000, 1000.0, rng)]
        cfg = dict(mu_q12=int(round(mu * 4096)), **_design(sd, 0.65, 6.0))
        on = len(IntDriftMonitor(rebaseline=True, **cfg).run(e))
        off = len(IntDriftMonitor(rebaseline=False, **cfg).run(e))
        assert on > off, "adapting the baseline costs false alarms"