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

EPS = 1e-12

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -40.0, 40.0)))

def binary_logloss(p, y):
    p = float(np.clip(p, 1e-12, 1.0 - 1e-12))
    return float(-y * np.log(p) - (1.0 - y) * np.log1p(-p))

class LagNewton:
    def __init__(self, r, horizon, eta=1.0, eps=1e-8):
        self.r = np.asarray(r, dtype=float)
        self.T = len(r)
        n = np.maximum(np.arange(horizon, horizon - self.T, -1), 1.0)
        self.prior = 1.0 / (n * self.r * self.r + eps)
        self.H = self.prior.copy()
        self.theta = np.zeros(self.T)
        self.eta = eta

    def step(self, x, y):
        p = float(sigmoid(np.dot(self.theta, x)))
        g = (p - y) * x
        h = p * (1.0 - p) * x * x + self.prior
        self.H += h
        self.theta = np.clip(self.theta - self.eta * g / self.H, -self.r, self.r)
        return p

class AdaGrad:
    def __init__(self, r, lr=0.25, eps=1e-8):
        self.r = np.asarray(r, dtype=float)
        self.theta = np.zeros(len(r))
        self.G = np.zeros(len(r))
        self.lr, self.eps = lr, eps

    def step(self, x, y):
        p = float(sigmoid(np.dot(self.theta, x)))
        g = (p - y) * x
        self.G += g * g
        self.theta = np.clip(self.theta - self.lr * g / (np.sqrt(self.G) + self.eps), -self.r, self.r)
        return p

class AdamW:
    def __init__(self, r, lr=0.03, wd=1e-4, eps=1e-8):
        self.r = np.asarray(r, dtype=float)
        self.theta = np.zeros(len(r))
        self.m = np.zeros(len(r)); self.v = np.zeros(len(r)); self.t = 0
        self.lr, self.wd, self.eps = lr, wd, eps

    def step(self, x, y):
        p = float(sigmoid(np.dot(self.theta, x)))
        g = (p - y) * x
        self.t += 1
        self.m = 0.9 * self.m + 0.1 * g
        self.v = 0.999 * self.v + 0.001 * g * g
        mh = self.m / (1.0 - 0.9 ** self.t)
        vh = self.v / (1.0 - 0.999 ** self.t)
        self.theta *= (1.0 - self.lr * self.wd)
        self.theta = np.clip(self.theta - self.lr * mh / (np.sqrt(vh) + self.eps), -self.r, self.r)
        return p

def make_stream(seed, T, kind):
    rng = np.random.default_rng(seed)
    u = rng.choice([-1.0, 1.0], size=T + 40)
    j = np.arange(1, 21, dtype=float)
    if kind == 'exponential':
        true = 0.75 * np.exp(-j / 5.0)
    else:
        true = 0.95 / (j ** 0.8)
    true *= min(1.0, 2.0 / true.sum())
    y = np.zeros(T, dtype=int)
    for t in range(T):
        x = u[t + 20 - j.astype(int)]
        y[t] = int(rng.random() < sigmoid(np.dot(true, x)))
    return u, y, true

def run_one(seed, kind, T=3000, d=20):
    u, y, true = make_stream(seed, T, kind)
    r = np.maximum(0.08, 1.35 * np.exp(-np.arange(1, d + 1) / 7.0))
    # Ensure the envelope contains the generating filter, while retaining lag weighting.
    r = np.maximum(r, np.abs(true) + 0.03)
    methods = {
        'lag_newton': LagNewton(r, T, eta=1.0),
        'adagrad': AdaGrad(r, lr=0.25),
        'adamw': AdamW(r, lr=0.03),
    }
    losses = {k: [] for k in methods}
    for t in range(T):
        x = u[t + 20 - np.arange(1, d + 1)]
        for name, opt in methods.items():
            p = opt.step(x, y[t])
            losses[name].append(binary_logloss(p, y[t]))
    out = {}
    for name, opt in methods.items():
        a = np.asarray(losses[name])
        out[name] = {
            'cum_loss': float(a.sum()),
            'mean_last_1000': float(a[-1000:].mean()),
            'mean_first_500': float(a[:500].mean()),
            'theta_mse': float(np.mean((opt.theta - true) ** 2)),
            'max_abs_theta': float(np.max(np.abs(opt.theta))),
        }
    out['true_l1'] = float(np.sum(np.abs(true)))
    out['newton_old_lag_abs_mean'] = float(np.mean(np.abs(methods['lag_newton'].theta[10:])))
    return out

def math_checks():
    T, d = 100, 10
    r = np.linspace(0.1, 0.5, d)
    n = np.arange(T, T - d, -1.0)
    gamma = np.sum(np.log1p(n * r * r))
    # The claimed prior scale decreases as available rounds increase; verify exact formula.
    prior = 1.0 / (n * r * r + 1e-8)
    # For x^2 <= 1 and logistic curvature <= 1/4, each increment is bounded as stated.
    p = np.linspace(0.01, 0.99, 50)
    hmax = np.max(p * (1 - p) + prior[0])
    H0 = prior.copy()
    H = H0 + np.ones(d) * 0.25
    return {
        'gamma': float(gamma),
        'gamma_nonnegative': bool(gamma >= 0),
        'prior_positive': bool(np.all(prior > 0)),
        'prior_matches_formula': bool(np.allclose(prior, 1.0 / (n * r * r + 1e-8))),
        'curvature_upper_bound_check': bool(hmax <= 0.25 + prior[0] + 1e-12),
        'H_monotone_check': bool(np.all(H >= H0)),
        'recent_prior': float(prior[0]),
        'old_prior': float(prior[-1]),
        'recent_has_more_rounds': bool(n[0] > n[-1]),
    }

def main():
    random.seed(7); np.random.seed(7)
    result = {'math_checks': math_checks(), 'runs': {}}
    for kind in ['exponential', 'polynomial']:
        result['runs'][kind] = [run_one(100 + i, kind) for i in range(3)]
    # Aggregate means for direct comparison.
    result['means'] = {}
    for kind, rows in result['runs'].items():
        result['means'][kind] = {}
        for method in ['lag_newton', 'adagrad', 'adamw']:
            result['means'][kind][method] = {k: float(np.mean([x[method][k] for x in rows])) for k in rows[0][method]}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result['math_checks'], indent=2))
    print(json.dumps(result['means'], indent=2))

if __name__ == '__main__':
    main()