"""Toy verification of position-only active-noise cancellation.

The quadratic is f(theta)=k theta^2/2.  Observed minibatch gradient is
 g=k*theta+a+epsilon, where a is an OU process.  The optimizer estimates a
from the residual g-k*theta (position and observed gradient only), then uses
 theta <- theta-eta*(g-c*ahat).

This is the minimal implementable part of the proposed controller.  The
script also computes the exact deterministic augmented transition matrix
for [theta,a,ahat] and checks rho<1 against simulated stability.
"""
import json
import numpy as np


def kf_run(seed, tau, eta, c, n=12000, k=1.0, sigma_a=1.0,
           sigma_g=0.20, burn=2000):
    rng = np.random.default_rng(seed)
    q = np.exp(-1.0 / tau)
    Q = sigma_a**2 * (1.0-q*q)
    # Stationary scalar Kalman covariance for measurement a + epsilon.
    P = sigma_a**2
    h = 0.0
    theta = 0.0
    xs, hs, aa = [], [], []
    a = rng.normal(0.0, sigma_a)
    for t in range(n):
        a = q*a + rng.normal(0.0, np.sqrt(Q))
        eps = rng.normal(0.0, sigma_g)
        g = k*theta + a + eps
        # Prediction and measurement update using residual g - k*theta.
        hp = q*h
        Pp = q*q*P + Q
        R = sigma_g*sigma_g
        gain = Pp/(Pp+R)
        h = hp + gain*((g-k*theta)-hp)
        P = (1.0-gain)*Pp
        theta = theta - eta*(g-c*h)
        if t >= burn:
            xs.append(theta); hs.append(h); aa.append(a)
    return float(np.mean(np.square(xs))), float(np.mean(np.square(np.array(aa)-np.array(hs)))), float(np.mean(np.array(xs)))


def transition(eta, k, tau, c):
    """Noise-free expected transition for [theta,a,ahat]."""
    q = np.exp(-1.0/tau)
    # In the zero-observation-noise filter, the update gain is 1.
    # For the actual noisy Kalman filter use its stationary gain below.
    R = .20**2
    Q = 1-q*q
    P = 1.0
    # converge scalar Riccati to get fixed gain
    for _ in range(10000):
        Pp=q*q*P+Q
        L=Pp/(Pp+R)
        P=(1-L)*Pp
    # theta' = (1-eta*k)theta -eta*a + eta*c*h
    # h' = L*a + q*(1-L)*h (residual is a in expectation)
    return np.array([[1-eta*k, -eta + eta*c*L, eta*c*q*(1-L)],
                     [0, q, 0],
                     [0, L, q*(1-L)]], dtype=float)


def feedback_transition(dt, k, G, Ga, tau, L):
    """A_cl for e=theta-lambda, a, ahat under the stated lambda feedback."""
    q=np.exp(-dt/tau)
    return np.array([[1+dt*(G-k), dt, dt*Ga],
                     [0,q,0], [0,L, q*(1-L)]], float)


def empirical_boundary():
    # Deterministic trajectories establish the observed boundary, while the
    # matrix supplies the claimed prediction.  Noise is deliberately absent.
    tau, k, c = 5.0, 1.0, 0.7
    last_stable = None; first_unstable = None
    for eta in np.linspace(.01, 2.5, 500):
        A=transition(eta,k,tau,c); rho=max(abs(np.linalg.eigvals(A)))
        x=np.array([1., 1., 0.])
        for _ in range(300): x=A@x
        stable=(rho < 1.0)
        if stable: last_stable=(eta,rho)
        elif first_unstable is None: first_unstable=(eta,rho)
    return last_stable, first_unstable


def main():
    taus=[0.2,1.0,5.0,20.0]
    etas=[0.15,0.30,0.50]
    rows=[]
    for tau in taus:
        for eta in etas:
            base=kf_run(10,tau,eta,0.0)
            idea=kf_run(10,tau,eta,0.85)
            rows.append({'tau':tau,'eta':eta,'sgd_mse':base[0],
                         'controller_mse':idea[0],
                         'ratio':idea[0]/base[0],
                         'residual_mse':idea[1]})
    lb, fu=empirical_boundary()
    # A representative feedback-loop stability check, matching the formula.
    fb=[]
    for G in [0.1,0.5,1.0,2.0]:
        A=feedback_transition(0.1,1.0,G,0.5,5.0,0.7)
        fb.append({'G':G,'rho':float(max(abs(np.linalg.eigvals(A))))})
    out={'rows':rows,'predicted_vs_empirical_boundary':{
        'last_stable_eta_rho':lb,'first_unstable_eta_rho':fu},
        'feedback_rho_scan':fb}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()