"""Observer-Corrected Robust Optimizer: scalar quadratic MVP.

For f(theta)=lambda*theta and disturbed gradient lambda*theta+d, the
feedforward controller applies -d_hat, hence
 theta_next = theta - h*(lambda*theta + d - d_hat).
The measured transition identifies d exactly in this toy plant.  d_hat is
then an EMA, which is the minimal discrete state observer for slowly varying
bias.  We test: (1) rho=|1-h*lambda| stability boundary, (2) EMA residual
frequency response, and (3) correction versus ordinary SGD on a biased
quadratic.
"""
import json
import pathlib
import numpy as np

SEED = 1330


def spectral_radius(h, lam):
    return abs(1.0 - h * lam)


def ema_residual_ratio(alpha, omega):
    """Predicted RMS ratio of applied residual d-d_hat for sinusoid d."""
    a = 1.0 - alpha
    # d_hat[k] = alpha*d[k-1] + a*d_hat[k-1]; applied estimate is old d_hat.
    # Transfer from d[k] to d[k]-d_hat[k] has this exact magnitude.
    z = np.exp(-1j * omega)
    hhat = alpha * z / (1.0 - a * z)
    return float(abs(1.0 - hhat))


def run_observer(lam, h, alpha, disturbance, theta0=1.0):
    theta, d_hat = float(theta0), 0.0
    thetas, dhats, applied_residuals, estimates = [], [], [], []
    for d in disturbance:
        d = float(d)
        # Correction cancels the estimated component of gradient disturbance.
        applied_residuals.append(d - d_hat)
        theta_next = theta - h * (lam * theta + d - d_hat)
        # Transition residual: d_obs = -(delta)/h-lambda*theta+d_hat.
        d_obs = -(theta_next - theta) / h - lam * theta + d_hat
        d_hat = (1.0 - alpha) * d_hat + alpha * d_obs
        theta, thetas = theta_next, thetas + [theta_next]
        dhats.append(d_hat)
        estimates.append(d_obs)
    return {"theta": np.asarray(thetas), "dhat": np.asarray(dhats),
            "applied_residual": np.asarray(applied_residuals),
            "d_obs": np.asarray(estimates)}


def run_baseline(lam, h, disturbance, theta0=1.0):
    theta, vals = float(theta0), []
    for d in disturbance:
        theta = theta - h * (lam * theta + float(d))
        vals.append(theta)
    return np.asarray(vals)


def stability_sweep(lam):
    hs = np.linspace(0.02, 2.35 / lam, 240)
    stable = []
    for h in hs:
        x = run_baseline(lam, h, np.zeros(300))[-1]
        stable.append(spectral_radius(h, lam) < 1.0 and abs(x) < 1e8)
    last = hs[np.where(stable)[0][-1]]
    first = hs[np.where(~np.asarray(stable))[0][0]]
    predicted = 2.0 / lam
    observed = 0.5 * (last + first)
    return {"lambda": lam, "predicted_h_critical": predicted,
            "observed_h_critical": float(observed),
            "relative_error": float(abs(observed-predicted)/predicted),
            "bracket": [float(last), float(first)]}


def disturbance_sweep(lam=3.0, h=None):
    if h is None: h = 0.35 / lam
    n, burn = 3000, 1000
    t = np.arange(n)
    cases = {"constant": (np.ones(n), 0.0),
             "slow_sine": (np.sin(2*np.pi*t/300.0), 2*np.pi/300.0),
             # Non-grid frequency avoids the accidental zero of sin(pi*k).
             "high_sine": (np.sin(2*np.pi*t/2.3), 2*np.pi/2.3)}
    rows = []
    for name, (d, omega) in cases.items():
        base = run_baseline(lam, h, d)
        input_rms = float(np.sqrt(np.mean(d[burn:]**2)))
        for alpha in [0.05, 0.2, 0.5, 0.9, 1.0]:
            out = run_observer(lam, h, alpha, d)
            actual = float(np.sqrt(np.mean(out["applied_residual"][burn:]**2)))
            pred = 0.0 if omega == 0 else ema_residual_ratio(alpha, omega)
            rows.append({"case": name, "alpha": alpha,
              "input_rms": input_rms, "residual_rms": actual,
              "residual_ratio": actual/input_rms,
              "predicted_ratio": pred,
              "ratio_relative_error": (None if omega == 0 else abs(actual/input_rms-pred)/pred),
              "baseline_theta_rms": float(np.sqrt(np.mean(base[burn:]**2))),
              "observer_theta_rms": float(np.sqrt(np.mean(out["theta"][burn:]**2)))})
    return rows


def main():
    np.random.seed(SEED)
    stability = [stability_sweep(lam) for lam in [0.5, 1.0, 3.0, 10.0]]
    rows = disturbance_sweep()
    headline = {}
    for case, alpha in [("constant", .5), ("slow_sine", .5),
                        ("high_sine", .05), ("high_sine", 1.0)]:
        headline[f"{case}_alpha_{alpha}"] = next(r for r in rows if r["case"] == case and r["alpha"] == alpha)
    summary = {"seed": SEED,
      "predictions": {
        "P1_stability": "h_critical=2/lambda from rho=|1-h*lambda|=1",
        "P2_frequency_response": "applied residual/input RMS follows |1-alpha exp(-iw)/(1-(1-alpha)exp(-iw))|",
        "P3_slow_bias": "constant/slow residual is strongly reduced, while high-frequency residual remains or grows"},
      "stability": stability, "disturbance_rows": rows, "headline": headline}
    pathlib.Path("observer_results.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))

if __name__ == "__main__": main()