"""Stage 20 - a prediction registered before the run, then the run.

The request was to find the domain where the method wins unambiguously.  The
way to get that wrong is to try domains until one of them comes out ahead;
that is fitting, spread over time.  So this file states the prediction and the
falsification condition FIRST, from what the quantity does, and then measures.

WHAT MU IS.  In NLMS the update w <- w + mu * e * x / ||x||^2 with mu = 1
drives the a posteriori error of the current sample to exactly zero.  Every
practical system uses mu < 1 because part of e is measurement noise the filter
cannot cancel, and fitting it is harmful.  So mu is the SIGNAL FRACTION of the
error, and its correct value is

    mu* = 1 - p_noise / p_error.

THE PREDICTION.  mu* is a property of the noise level.  Therefore:

  (A) CONTROL - when the noise level is constant, a constant mu can be exactly
      right.  An oracle that picks the best fixed mu for that scenario should
      MATCH OR BEAT the measured rule.  If the measured rule wins here, the
      reasoning is wrong and everything below is suspect.

  (B) PREDICTION - when the noise level CHANGES during the run, no constant is
      right throughout, while a measured quantity tracks.  The measured rule
      should BEAT the oracle fixed mu, and the margin should grow with the size
      of the change.

FALSIFICATION.  If (B) fails - if an oracle constant matches the measured rule
even when the noise level swings by 25 dB within a single run - then there is
no such domain on this evidence, and that is the finding.

The oracle is the strongest competitor available: the best of a mu grid chosen
separately for every scenario, with knowledge of that scenario, which no
deployable controller has.  Kwong-Johnston is given the same treatment.
"""
from __future__ import annotations

import math

import numpy as np

from anc.bench import (FS, gen_broadband, gen_impulse_train, gen_multitonal,
                       gen_tonal, make_primary_path, make_secondary_path)
from anc.rules import Fixed, KwongJohnston, Measured, Rule

N = 16000                     # two seconds
SEEDS = 8
MU_GRID = (0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 0.85, 1.0)
KJ_GRID = [(a, g) for a in (0.9, 0.97, 0.99)
           for g in (2e-4, 2e-3, 2e-2, 2e-1)]
FILTER_LEN = 64


class KJTuned(KwongJohnston):
    def __init__(self, a, g):
        self.A, self.G = a, g
        self.name = f"KJ({a:g},{g:g})"


def run(noise, prim, sec, rule, snr_profile, seed=0):
    """One run.  snr_profile(i) gives the instantaneous SNR in dB, so the
    stationary and the varying case go through exactly the same code."""
    rng = np.random.default_rng(seed)
    n = len(noise)
    d = np.convolve(noise, prim)[:n]
    xf = np.convolve(noise, sec)[:n]

    # the disturbance is added with a time-varying power, which is the only
    # thing that differs between the control and the prediction
    p_sig = float(np.mean(d * d)) + 1e-30
    snr = np.array([snr_profile(i / n) for i in range(n)])
    amp = np.sqrt(p_sig / (10.0 ** (snr / 10.0)))
    d = d + rng.standard_normal(n) * amp

    w = np.zeros(FILTER_LEN)
    buf = np.zeros(FILTER_LEN)
    err = np.zeros(n)
    rule.reset()
    for i in range(n):
        buf[1:] = buf[:-1]
        buf[0] = xf[i]
        e = d[i] - float(np.dot(w, buf))
        err[i] = e
        nrm = float(np.dot(buf, buf)) + 1e-6
        mu = rule.step(e, buf, nrm)
        if mu != 0.0:
            w += mu * e * buf / nrm
            rule.observe(e, d[i] - float(np.dot(w, buf)), mu)
        else:
            rule.observe(e, e, 0.0)
    # The measurement window matters, and the first version got it wrong.
    # Scoring only the second half made the varying-noise scenarios
    # indistinguishable from a constant run at their FINAL level - the ramp
    # and the step have both arrived by then, so the metric never saw the
    # transition and the prediction could not have been confirmed even if it
    # were true.  The window now excludes only the initial adaptation
    # transient, which every controller pays alike.
    warm = n // 10
    p_in = float(np.mean(d[warm:] ** 2))
    p_out = float(np.mean(err[warm:] ** 2))
    if p_out <= 0.0 or not math.isfinite(p_out):
        return float("nan")
    return 10.0 * math.log10(p_in / p_out)


def score(rule, gen, profile, seeds=SEEDS):
    vals = []
    for s in range(seeds):
        rng = np.random.default_rng(2000 + s)
        noise = gen(N, rng)
        nr = run(noise, make_primary_path(rng), make_secondary_path(rng),
                 rule, profile, seed=s)
        if math.isfinite(nr):
            vals.append(nr)
    return (sum(vals) / len(vals)) if vals else float("nan")


def oracle_fixed(gen, profile):
    """The ceiling of what any constant can do on this scenario: the best mu
    of the grid, chosen WITH knowledge of the scenario."""
    best, arg = -1e30, None
    for mu in MU_GRID:
        v = score(Fixed(mu), gen, profile)
        if math.isfinite(v) and v > best:
            best, arg = v, mu
    return best, arg


def oracle_kj(gen, profile):
    best, arg = -1e30, None
    for a, g in KJ_GRID:
        v = score(KJTuned(a, g), gen, profile)
        if math.isfinite(v) and v > best:
            best, arg = v, (a, g)
    return best, arg


# --------------------------------------------------------------- scenarios --
def sig_tonal(n, rng):
    return gen_tonal(n, rng)


def sig_cabin(n, rng):
    s = gen_multitonal(n, rng, f0=80.0) + 0.6 * gen_broadband(n, rng) \
        + 0.4 * gen_impulse_train(n, rng)
    return s / (np.max(np.abs(s)) + 1e-12)


SIGNALS = [("tonal", sig_tonal), ("car cabin", sig_cabin)]

# noise-level profiles, as functions of the fraction of the run elapsed
FLAT_30 = lambda u: 30.0
FLAT_15 = lambda u: 15.0
FLAT_5 = lambda u: 5.0
SWING_10 = lambda u: 25.0 - 10.0 * u
SWING_25 = lambda u: 30.0 - 25.0 * u
STEP_25 = lambda u: 30.0 if u < 0.5 else 5.0

PROFILES = [("flat 30 dB", FLAT_30, "control"),
            ("flat 15 dB", FLAT_15, "control"),
            ("flat 5 dB", FLAT_5, "control"),
            ("ramp 25->15 dB", SWING_10, "prediction"),
            ("ramp 30->5 dB", SWING_25, "prediction"),
            ("step 30->5 dB", STEP_25, "prediction")]


def main():
    print(__doc__.split("The oracle is")[0])
    print("=" * 92)
    print("THE RUN")
    print("=" * 92)
    print(f"{'signal':11s}{'noise profile':17s}{'kind':12s}"
          f"{'oracle mu':>11s}{'best mu':>9s}{'oracle KJ':>11s}"
          f"{'measured':>10s}{'verdict':>12s}")
    print("-" * 92)
    control_wins = pred_wins = control_n = pred_n = 0
    for signame, gen in SIGNALS:
        for pname, prof, kind in PROFILES:
            om, arg = oracle_fixed(gen, prof)
            okj, _a = oracle_kj(gen, prof)
            ms = score(Measured(), gen, prof)
            beats = ms > max(om, okj)
            if kind == "control":
                control_n += 1
                control_wins += beats
            else:
                pred_n += 1
                pred_wins += beats
            print(f"{signame:11s}{pname:17s}{kind:12s}{om:>11.2f}"
                  f"{arg:>9.2f}{okj:>11.2f}{ms:>10.2f}"
                  f"{'MEASURED' if beats else 'oracle':>12s}", flush=True)
    print("-" * 92)
    print(f"\n   control cells  (constant noise): measured wins "
          f"{control_wins} of {control_n}")
    print(f"   predicted cells (varying noise): measured wins "
          f"{pred_wins} of {pred_n}")
    ok = control_wins == 0 and pred_wins == pred_n
    print(f"\n   the prediction is {'CONFIRMED' if ok else 'NOT confirmed'}: "
          f"it required the measured rule to lose every\n   control cell and "
          f"win every predicted one.")


if __name__ == "__main__":
    main()