import json, math
import numpy as np

SEED = 3100
rng = np.random.default_rng(SEED)


def monomials(z, degree):
    """[1, linear terms, all total-degree 2..degree] in deterministic order."""
    z = np.asarray(z, dtype=float)
    out = [1.0]
    n = len(z)
    # recursive exponent enumeration, grouped by total degree
    def comps(total, dims, prefix=()):
        if dims == 1:
            yield prefix + (total,)
        else:
            for a in range(total + 1):
                yield from comps(total-a, dims-1, prefix+(a,))
    for d in range(1, degree+1):
        for alpha in comps(d, n):
            v = 1.0
            for zi, ai in zip(z, alpha): v *= zi ** ai
            out.append(v)
    return np.asarray(out)


class OnlineTaylorRLS:
    def __init__(self, in_dim, out_dim, degree=1, lam=0.98, p0=100.0):
        self.in_dim, self.out_dim, self.degree = in_dim, out_dim, degree
        self.lam = lam
        self.W = np.zeros((out_dim, len(monomials(np.zeros(in_dim), degree))))
        self.P = np.eye(self.W.shape[1]) * p0

    def update(self, z, target):
        phi = monomials(z, self.degree)
        Pphi = self.P @ phi
        K = Pphi / (self.lam + phi @ Pphi)
        err = np.asarray(target) - self.W @ phi
        self.W += np.outer(err, K)
        self.P = (self.P - np.outer(K, phi @ self.P)) / self.lam
        self.P = (self.P + self.P.T) / 2
        return err

    def predict_residual(self, z):
        return self.W @ monomials(z, self.degree)


def rls_exactness_check():
    # RLS with lambda=1 must equal ordinary least squares after every prefix.
    local = np.random.default_rng(SEED + 1)
    Z = local.normal(size=(30, 3)); Phi = np.array([monomials(z, 2) for z in Z])
    true_w = local.normal(size=(2, Phi.shape[1]))
    Y = Phi @ true_w.T + .01 * local.normal(size=(30, 2))
    r = OnlineTaylorRLS(3, 2, degree=2, lam=1.0, p0=1e8)
    errors=[]
    for i in range(len(Z)):
        r.update(Z[i], Y[i])
        # large p0 is effectively unregularized; compare after enough samples
        batch = np.linalg.lstsq(Phi[:i+1], Y[:i+1], rcond=None)[0].T
        errors.append(float(np.max(np.abs(r.W-batch))))
    # Verify the covariance-weighted identity more robustly via predictions on heldout points.
    pred_err = np.max(np.abs((r.W @ Phi[-5:].T) - (batch @ Phi[-5:].T)))
    return {"max_prefix_coef_error": max(errors[10:]), "final_prediction_error": float(pred_err)}


def forgetting_check():
    # Constant residual step: adaptation should be faster for lower lambda,
    # while noisy stationary variance should be larger (the stated tradeoff).
    def run(lam, noise, n=500):
        r=OnlineTaylorRLS(1,1,degree=0,lam=lam,p0=1.0)
        vals=[]; rg=np.random.default_rng(55)
        for k in range(n):
            y = (1.0 if k >= 30 else 0.0) + noise*rg.normal()
            r.update([0.0], [y]); vals.append(float(r.W[0,0]))
        return np.asarray(vals)
    fast=run(.90,.08); slow=run(.99,.08)
    def recovery(a):
        return int(np.argmax(a[30:] >= .9)+30)
    return {"lambda_.90_recovery_step": recovery(fast), "lambda_.99_recovery_step": recovery(slow),
            "lambda_.90_post_std": float(np.std(fast[-150:])), "lambda_.99_post_std": float(np.std(slow[-150:]))}


def dynamics_experiment():
    """Frozen nominal MLP versus online Taylor residual after a dynamics shift."""
    from sklearn.neural_network import MLPRegressor
    local = np.random.default_rng(SEED + 2)

    def F_nom(x, u):
        return .82*x + .18*u + .08*x*x
    def F_changed(x, u):
        return .82*x + .18*u + .08*x*x + .32*u + .10

    # Train only on nominal transitions; this is the frozen global prior.
    tx = local.uniform(-1, 1, 4000)
    tu = local.uniform(-1, 1, 4000)
    prior = MLPRegressor(hidden_layer_sizes=(24, 24), activation='tanh',
                         solver='lbfgs', alpha=1e-5, max_iter=300,
                         random_state=SEED)
    prior.fit(np.c_[tx, tu], F_nom(tx, tu))
    def base(x, u):
        return float(prior.predict(np.array([[x, u]]))[0])

    # Same shifted trajectory and inputs for every method.
    n = 260
    us = local.uniform(-.8, .8, n)
    xs = np.zeros(n + 1)
    for k in range(n):
        xs[k+1] = np.clip(F_changed(xs[k], us[k]), -2, 2)

    results = {}
    for degree in [0, 1, 2, 3]:
        r = None if degree == 0 else OnlineTaylorRLS(2, 1, degree=degree,
                                                     lam=.94, p0=20.)
        one_step, early, horizon = [], [], []
        for k in range(n):
            b = base(xs[k], us[k])
            if r is None:
                pred = b
                predict = lambda xx, uu: base(xx, uu)
            else:
                predict = lambda xx, uu: base(xx, uu) + float(
                    r.predict_residual(np.array([xx, uu]))[0])
                pred = predict(xs[k], us[k])
            if k >= 30:
                one_step.append(abs(pred - xs[k+1]))
                if k < 80:
                    early.append(abs(pred - xs[k+1]))
            if k >= 30 and k + 10 <= n:
                xx = xs[k]
                for j in range(10):
                    xx = predict(xx, us[k+j])
                    horizon.append(abs(xx - xs[k+j+1]))
            if r is not None:
                r.update(np.array([xs[k], us[k]]), [xs[k+1] - b])
        results[str(degree)] = {
            "mean_abs_one_step_error": float(np.mean(one_step)),
            "early_postshift_error": float(np.mean(early)),
            "mean_abs_10_step_rollout_error": float(np.mean(horizon)),
        }
    return results

def main():
    out={"seed":SEED,"rls_exactness":rls_exactness_check(),
         "forgetting_tradeoff":forgetting_check(),"dynamics":dynamics_experiment()}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()