"""Step-size rules for NLMS.  Everything else in the filter is shared.

    mu = 1        drives the a posteriori error of the current sample to zero
    mu < 1        hedges, because that sample carries measurement noise too

So mu is a safety factor on an exact step, exactly as in ODE step control.
The variable-step-size literature answers the difficulty of choosing it with
more constants; the rule at the bottom measures it instead.
"""
from __future__ import annotations

import math


class Rule:
    name = "?"
    constants = 0

    def reset(self):
        pass

    def step(self, e, x, nrm):
        raise NotImplementedError

    def observe(self, e_pre, e_post, mu):
        pass


class Fixed(Rule):
    def __init__(self, mu):
        self.mu = mu
        self.name = f"fixed mu={mu:g}"
        self.constants = 1

    def step(self, e, x, nrm):
        return self.mu


class KwongJohnston(Rule):
    """The classical variable step size: mu(n+1) = a*mu(n) + g*e(n)^2, clamped.

    Four constants, and the paper that introduced it spends its length on how
    to pick them."""

    name = "Kwong-Johnston VSS"
    constants = 4
    A, G, MU_MIN, MU_MAX = 0.97, 0.02, 0.001, 1.0

    def reset(self):
        self.mu = 0.5

    def step(self, e, x, nrm):
        return self.mu

    def observe(self, e_pre, e_post, mu):
        self.mu = self.A * self.mu + self.G * e_pre * e_pre
        self.mu = min(self.MU_MAX, max(self.MU_MIN, self.mu))


class InSampleMeasured(Rule):
    """REFUTED - kept so the mistake is reproducible.

    The first attempt measured the A POSTERIORI error of the sample the update
    had just been fitted to.  That observable is uninformative by
    construction: NLMS with mu = 1 drives it to exactly zero, so the rule
    reads "the model was perfect", proposes mu = 1 again, and is locked there.
    Measured, it scores 18.50 dB against fixed mu = 1's 18.49 - the same rule
    wearing a different name.

    It is the same error as the rho detector in the root-finding work:
    measuring a quantity in the one place where it is identically zero."""

    name = "in-sample (refuted)"
    constants = 0

    def reset(self):
        self.rho = 1.0

    def step(self, e, x, nrm):
        return min(1.0, 1.0 / self.rho) if self.rho > 0.0 else 1.0

    def observe(self, e_pre, e_post, mu):
        if mu <= 0.0 or e_pre == 0.0:
            return
        predicted = abs((1.0 - mu) * e_pre)
        observed = abs(e_post)
        if predicted <= 0.0:
            self.rho = 1.0 + observed / (abs(e_pre) + 1e-300)
        else:
            r = observed / predicted
            if math.isfinite(r) and r > 0.0:
                self.rho = math.sqrt(self.rho * max(r, 1e-6))


class Measured(Rule):
    """The safety factor measured out of sample.

    What mu is hedging against is the part of the error the filter cannot
    cancel - the measurement noise.  NLMS theory gives the step that minimises
    the mean square deviation as the SIGNAL FRACTION of the error,

        mu = 1 - p_noise / p_error,

    and both quantities are observable.  p_error is the running power of e.
    p_noise is what e cannot be driven below, so it is estimated as the
    smallest error power the run has sustained - the floor the filter keeps
    hitting.

    Nothing here is chosen.  The two smoothing weights are 1/2, an average of
    the current observation with the accumulated one, not a tuned forgetting
    factor; and the first samples, having no floor to compare against, take
    the exact step - the licence the root-finding safeguard grants.
    """

    name = "measured (no constants)"
    constants = 0

    def reset(self):
        self.p_err = None      # running error power
        self.floor = None      # the smallest sustained error power seen

    def step(self, e, x, nrm):
        if self.p_err is None or self.floor is None or self.p_err <= 0.0:
            return 1.0
        frac = 1.0 - self.floor / self.p_err
        return min(1.0, max(0.0, frac))

    def observe(self, e_pre, e_post, mu):
        p = e_pre * e_pre
        self.p_err = p if self.p_err is None else 0.5 * (self.p_err + p)
        if self.floor is None:
            self.floor = self.p_err
        elif self.p_err < self.floor:
            self.floor = self.p_err
        else:
            # the floor may rise when the environment changes; let it drift up
            # by the same averaging, so a regime change is not remembered for
            # ever as a floor that no longer exists
            self.floor = 0.5 * (self.floor + min(self.p_err, self.floor * 2.0))


class CosineSimilarity(Rule):
    """mu as a SIMILARITY, which is what this project was supposed to be about.

    The two rules above estimate mu from variances - a noise floor, a model
    discrepancy.  Neither is a similarity; both are statistics of the error.
    The similarity formulation asks a different question:

        how much does the error RESEMBLE something the filter can remove?

    The filter can only remove what is correlated with its input.  So the
    quantity is the cosine similarity between the error and the filtered
    reference,

        mu = |<e, x>| / (||e|| ||x||)   in [0, 1],

    which is the cosine of the original paper, applied to the two signals the
    filter actually has.  It equals 1 when the error is entirely explainable
    by the reference - the exact step is then right - and 0 when the error is
    orthogonal to it, which is measurement noise the update can only fit.

    The three running means use the weight 1/2, an average of the current
    observation with the accumulated one, and the first sample takes the exact
    step for want of anything to compare.  No constant is chosen.
    """

    name = "cosine similarity (no constants)"
    constants = 0
    S = None            # None means the identity: mu = |cos| directly

    def reset(self):
        self.ex = None      # running <e, x>
        self.ee = None      # running <e, e>
        self.xx = None      # running <x, x>

    def step(self, e, x, nrm):
        # decide from what has been accumulated SO FAR, then fold in the
        # current sample - so the step never uses the observation it is about
        # to be judged on
        if self.ex is None:
            mu = 1.0
        else:
            den = math.sqrt(self.ee * self.xx)
            if den <= 0.0 or not math.isfinite(den):
                mu = 1.0
            else:
                cos = min(1.0, abs(self.ex) / den)
                if self.S is None:
                    mu = cos
                else:
                    # the general family: mu = S(d) with d = 1 - |cos| the
                    # dissimilarity between the error and what the filter can
                    # remove.  d = 0 means fully removable, so S(0) = 1 and the
                    # exact step is taken; the identity above is S(d) = 1 - d.
                    mu = max(0.0, min(1.0, self.S(1.0 - cos)))
        xv = float(x[0])
        self.ex = e * xv if self.ex is None else 0.5 * (self.ex + e * xv)
        self.ee = e * e if self.ee is None else 0.5 * (self.ee + e * e)
        self.xx = xv * xv if self.xx is None else 0.5 * (self.xx + xv * xv)
        return mu


ALL = [Fixed(0.1), Fixed(0.5), Fixed(1.0), KwongJohnston(),
       InSampleMeasured(), Measured(), CosineSimilarity()]