import json, math
import numpy as np

SEED = 7


def rls_update(theta, P, r, y, lam):
    """Forgetting-factor RLS, theta has shape (outputs, inputs)."""
    Pr = P @ r
    Pn = (P - np.outer(Pr, Pr) / (lam + r @ Pr)) / lam
    en = y - theta @ r
    return theta + np.outer(en, r @ Pn), Pn


def lift(x):
    return np.array([x, x * x], dtype=float)


def dynamics(x, regime):
    return (0.62 * x + 0.22 * x*x - 0.10 * x**3) if regime == 0 else (0.35 * x + 0.42 * x*x - 0.10 * x**3)


def stream(n=500, change=250):
    rng = np.random.default_rng(SEED)
    x, data = 0.15, []
    for t in range(n):
        u = .10 * math.sin(.17*t) + .025 * rng.normal()
        reg = int(t >= change)
        y = dynamics(x, reg) + u
        data.append((x, u, y, reg))
        x = np.clip(y, -1.4, 1.4)
    return data


def train(data, lam, lifted, updates=True):
    # Linear baseline models x' from [x,u,1]; lifted model z' from [x,x^2,u,1].
    def inp(x, u): return np.r_[lift(x) if lifted else x, u, 1.0]
    def out(y): return lift(y) if lifted else np.array([y])
    din, dout = (4, 2) if lifted else (3, 1)
    th, P = np.zeros((dout, din)), np.eye(din) * 100.
    for x,u,y,_ in data[:250]:
        if updates: th,P = rls_update(th,P,inp(x,u),out(y),lam)
    before = th.copy()
    errs, rhos = [], []
    for x,u,y,_ in data[250:]:
        if updates: th,P = rls_update(th,P,inp(x,u),out(y),lam)
        pred = th @ inp(x,u)
        errs.append((pred[0] - y)**2)
        rhos.append(float(max(abs(np.linalg.eigvals(th[:,:2])))) if lifted else abs(float(th[0,0])))
    return np.asarray(errs), np.asarray(rhos), before, th


def rollout(th, x0, horizon, regime):
    z = lift(x0); truth = x0; pred = []
    for _ in range(horizon):
        # no external input; last column is the intercept
        zp = th @ np.r_[z, 0.0, 1.0]
        z = zp
        truth = dynamics(truth, regime)
        pred.append((z[0], truth))
    return float(np.mean([(a-b)**2 for a,b in pred]))


def math_checks():
    # Exact finite-horizon identity: prior precision is lambda^N P0^-1.
    rng = np.random.default_rng(11); d, n, lam = 3, 40, .93
    th, P = np.zeros((1,d)), np.eye(d)*4.0
    X, Y = [], []
    for _ in range(n):
        r = rng.normal(size=d); y = np.array([.4*r[0]-.2*r[1]+.1*r[2]])
        th,P = rls_update(th,P,r,y,lam); X.append(r); Y.append(y)
    X, Y = np.asarray(X), np.asarray(Y).reshape(-1,1)
    W = np.diag(lam**np.arange(n-1,-1,-1))
    batch = (Y.T@W@X) @ np.linalg.inv(X.T@W@X + (lam**n)*np.eye(d)/4.0)
    # 63% exponential forgetting mass and geometric contraction.
    timescale = 1/(1-lam)
    measured = int(np.argmin(abs((1-lam**np.arange(1,200))-(1-1/math.e)))+1)
    return {"rls_batch_max_error": float(np.max(abs(th-batch))),
            "adaptation_timescale": timescale, "measured_63pct_steps": measured,
            "rho_0.8_power29": .8**29}


def stability_check():
    # Scaling a stable scalar lifted operator gives the claimed rho=1 boundary.
    vals=[]
    for s in [.8, .99, 1.01, 1.2]:
        rho=abs(s); norm=abs(s**40)
        vals.append({"scale":s,"rho":rho,"power40_norm":norm})
    return vals


def main():
    data=stream(); result={"seed":SEED,"math":math_checks(),"stability_boundary":stability_check()}
    for lam in (.98,.95,.90):
        for lifted in (False,True):
            e,r,_,_=train(data,lam,lifted,True)
            result[f"adaptive_{'lifted' if lifted else 'linear'}_lambda_{lam}"]={
                "post_shift_rmse_first20":float(np.sqrt(np.mean(e[:20]))),
                "post_shift_rmse_last100":float(np.sqrt(np.mean(e[-100:]))),
                "max_rho":float(np.max(r))}
    for lifted in (False,True):
        e,_,_,_=train(data,.95,lifted,False)
        result[f"frozen_{'lifted' if lifted else 'linear'}"]={
            "post_shift_rmse_first20":float(np.sqrt(np.mean(e[:20]))),
            "post_shift_rmse_last100":float(np.sqrt(np.mean(e[-100:]))) }
    _,_,frozen,adapted=train(data,.95,True,True)
    result["50step_post_shift_rollout_mse"]={"adaptive_lifted":rollout(adapted,.45,50,1),"frozen_lifted":rollout(frozen,.45,50,1)}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == '__main__': main()