import json, math, random
from pathlib import Path
import numpy as np

SEED = 2037
np.random.seed(SEED)
random.seed(SEED)

class BoundedPredictiveGain:
    """Euclidean mirror descent gain, one scalar per group."""
    def __init__(self, alpha0, eta, amin=None, amax=None, rho=0.0, qmax=None):
        self.alpha0 = float(alpha0)
        self.eta = float(eta)
        self.amin = 0.1 * alpha0 if amin is None else float(amin)
        self.amax = 10.0 * alpha0 if amax is None else float(amax)
        self.rho = float(rho)
        self.qmax = qmax
        self.a = float(alpha0)
        self.ref = float(alpha0)
        self.clip_count = 0
        self.steps = 0
    def step(self, q):
        q0 = float(np.clip(q, -self.qmax, self.qmax)) if self.qmax is not None else float(q)
        raw = (self.a - self.eta * q0 + self.rho * self.ref) / (1.0 + self.rho)
        new = float(np.clip(raw, self.amin, self.amax))
        self.clip_count += int(new != raw)
        self.a = new
        self.steps += 1
        return new

def toy_prediction_1_sign_and_slope():
    # At a=aref the exact update predicts delta a = eta*r*g^2/(1+rho).
    eta, rho, g, a0 = 0.07, 0.4, 1.3, 0.8
    rs = np.array([-1.0, -0.5, 0.0, 0.5, 1.0])
    observed = []
    predicted = []
    for r in rs:
        opt = BoundedPredictiveGain(a0, eta, amin=0.0, amax=100.0, rho=rho)
        old = opt.a
        opt.step(-r * g * g)
        observed.append(opt.a - old)
        predicted.append(eta * r * g * g / (1 + rho))
    slope = float(np.polyfit(rs, observed, 1)[0])
    expected_slope = eta * g * g / (1 + rho)
    max_err = float(np.max(np.abs(np.array(observed) - predicted)))
    return {
        "name": "sign and magnitude of consecutive-direction product",
        "prediction": {"delta_a": "eta*r*g^2/(1+rho)", "slope": expected_slope},
        "observed": {"r": rs.tolist(), "delta_a": observed, "slope": slope, "max_abs_error": max_err},
        "pass": bool(max_err < 1e-12 and slope > 0 and observed[0] < 0 and observed[-1] > 0)
    }

def toy_prediction_2_persistence():
    # With q=0, a'=(a+rho*ref)/(1+rho); error contracts by 1/(1+rho).
    a0, ref, rho = 2.0, 0.7, 0.5
    opt = BoundedPredictiveGain(ref, eta=0.1, amin=0, amax=100, rho=rho)
    opt.a = a0
    errors = [abs(opt.a - ref)]
    for _ in range(5):
        opt.step(0.0)
        errors.append(abs(opt.a - ref))
    ratios = np.array(errors[1:]) / np.array(errors[:-1])
    expected = 1 / (1 + rho)
    return {
        "name": "Bregman persistence contraction when predictive product is zero",
        "prediction": {"error_ratio": expected},
        "observed": {"errors": errors, "ratios": ratios.tolist(), "max_abs_ratio_error": float(np.max(abs(ratios-expected)))},
        "pass": bool(np.max(abs(ratios-expected)) < 1e-12)
    }

def toy_prediction_3_quadratic_boundary():
    # For f=.5*lam*theta^2 and fixed a: theta sign first reverses at a*lam>1,
    # and linear dynamics is stable iff 0<a*lam<2.
    lam = 1.0
    gains = np.linspace(0.2, 2.4, 12)
    sign_reversal = []
    stable = []
    for a in gains:
        theta = 1.0
        signs = []
        for _ in range(80):
            signs.append(np.sign(theta))
            theta = (1-a*lam)*theta
        # Strict reversal means consecutive nonzero states have opposite signs;
        # at a*lam=1 the iterate lands exactly at zero rather than overshooting.
        vals = np.array(signs)
        sign_reversal.append(bool(np.any(vals[1:] * vals[:-1] < 0)))
        stable.append(abs(1-a*lam) < 1)
    first_observed = float(gains[np.where(sign_reversal)[0][0]])
    first_grid_above = float(gains[np.where(gains > 1/lam)[0][0]])
    stable_mismatch = int(np.sum(np.array(stable) != ((gains > 0) & (gains < 2/lam))))
    return {
        "name": "quadratic first sign reversal and stability boundary",
        "prediction": {"first_reversal_threshold": 1/lam, "stability_interval": [0.0, 2/lam]},
        "observed": {"first_grid_reversal": first_observed, "first_grid_point_above_threshold": first_grid_above, "stable_boundary_test_mismatches": stable_mismatch, "grid": gains.tolist(), "reversal": sign_reversal},
        "pass": bool(first_observed > 1/lam and first_observed <= first_grid_above and stable_mismatch == 0)
    }

def make_data(n=800):
    rng = np.random.RandomState(SEED)
    x = rng.randn(n, 2).astype(np.float32)
    y = ((x[:,0] * x[:,1] + 0.25*x[:,0] - 0.15*x[:,1]) > 0).astype(np.int64)
    return x, y

def mini_training():
    # Small deterministic numpy MLP keeps the comparison transparent and cheap.
    x, y = make_data()
    rng = np.random.RandomState(SEED+1)
    d, h = 2, 24
    w1 = rng.randn(d,h).astype(np.float64)*0.5; b1=np.zeros(h)
    w2 = rng.randn(h,2).astype(np.float64)*0.5; b2=np.zeros(2)
    init = [z.copy() for z in (w1,b1,w2,b2)]
    def run(adaptive):
        ww = [z.copy() for z in init]; gains=[0.8,0.8]; prev_u=[None,None,None,None]
        losses=[]; accs=[]; clips=0; reversals=0
        for t in range(160):
            ix = np.arange((t*64)% (len(x)-64), (t*64)% (len(x)-64)+64)
            X=x[ix].astype(np.float64); Y=y[ix]
            z=X@ww[0]+ww[1]; ah=np.maximum(z,0); logits=ah@ww[2]+ww[3]
            logits-=logits.max(1,keepdims=True); p=np.exp(logits); p/=p.sum(1,keepdims=True)
            loss=-np.log(p[np.arange(len(Y)),Y]+1e-12).mean(); losses.append(float(loss))
            dl=p; dl[np.arange(len(Y)),Y]-=1; dl/=len(Y)
            grads=[X.T@(dl@ww[2].T*(z>0)), (dl@ww[2].T*(z>0)).sum(0), ah.T@dl, dl.sum(0)]
            us=[grads[0],grads[1],grads[2],grads[3]]
            if adaptive and t > 0:
                for gi, inds in enumerate(([0,1],[2,3])):
                    dot=sum(float(np.mean(prev_u[j]*us[j])) for j in inds)/2
                    q=-dot
                    old=gains[gi]
                    gains[gi]=float(np.clip((old-0.12*q+0.15*0.8)/(1.15),0.08,8.0))
                    clips += int(gains[gi] in (0.08,8.0)); reversals += int(dot < 0)
            step_g = gains if adaptive else [0.8,0.8]
            for j in [0,1]: ww[j] -= step_g[0]*us[j]
            for j in [2,3]: ww[j] -= step_g[1]*us[j]
            prev_u=[u.copy() for u in us]
            # full-data accuracy checkpoint
            zz=x@ww[0]+ww[1]; pp=np.maximum(zz,0)@ww[2]+ww[3]
            accs.append(float((pp.argmax(1)==y).mean()))
        return {"final_loss":losses[-1],"best_loss":min(losses),"final_accuracy":accs[-1],"loss_at_40":losses[39],"gain_trajectories":gains,"clip_events":clips,"reversal_events":reversals}
    return run(False), run(True)

def main():
    checks=[toy_prediction_1_sign_and_slope(),toy_prediction_2_persistence(),toy_prediction_3_quadratic_boundary()]
    baseline, idea=mini_training()
    out={"seed":SEED,"checks":checks,"mini_experiment":{"baseline":baseline,"bounded_predictive_gain":idea}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()