import json
import math
from pathlib import Path
import numpy as np
from scipy.optimize import brentq


def g_filter(s, nu, t):
    d = float(s - nu)
    if d == 0.0:
        return float(t)
    return -math.expm1(-d * t) / d


def f_filter(s, nu, t):
    return float(s) * g_filter(s, nu, t)


def discrete_g(s, nu, eta, K):
    """Coefficient filter of K exact Euler updates from beta_0=0."""
    q = 1.0 - eta * (s - nu)
    if abs(q - 1.0) < 1e-12:
        return eta * K
    return eta * (1.0 - q ** K) / (1.0 - q)


def root_filter(nu, t, lo=1e-8, hi=5.0):
    grid = np.linspace(lo, hi, 4001)
    vals = [f_filter(s, nu, t) - 1.0 for s in grid]
    for a, b, fa, fb in zip(grid[:-1], grid[1:], vals[:-1], vals[1:]):
        if fa * fb < 0:
            return float(brentq(lambda s: f_filter(s, nu, t) - 1.0, a, b))
    return None


def core_checks():
    nu, t = 1.0, 1.5
    # Prediction 1: the pole is removable, with f(nu)=nu*t, and error is O(|s-nu|).
    eps = 10.0 ** (-np.arange(2, 9))
    errors = np.array([abs(f_filter(nu + e, nu, t) - nu * t) for e in eps])
    slopes = np.diff(np.log(errors)) / np.diff(np.log(eps))
    pole = {
        "predicted_limit": nu * t,
        "observed_at_pole": f_filter(nu, nu, t),
        "exact_error": abs(f_filter(nu, nu, t) - nu * t),
        "errors_eps_1e-2_to_1e-8": errors.tolist(),
        "median_loglog_slope": float(np.median(slopes[-4:])),
        "prediction": "limit is nu*t and off-pole error scales linearly in |s-nu|",
    }

    # Prediction 2: for fixed nu,t, f(s)-1 changes sign once; compare analytic
    # root with a grid observation, and sweep t to verify the predicted transition.
    crossover_rows = []
    for tt in [0.5, 1.0, 1.5, 2.0, 3.0]:
        pred = root_filter(nu, tt)
        grid = np.linspace(1e-5, 5.0, 100001)
        obs = float(grid[np.argmin(np.abs(np.array([f_filter(s, nu, tt) for s in grid]) - 1))])
        crossover_rows.append({"t": tt, "predicted_root": pred, "observed_grid_root": obs,
                               "absolute_error": abs(pred - obs) if pred is not None else None,
                               "f_at_nu": f_filter(nu, nu, tt)})
    mixed = {"rows": crossover_rows,
             "example_low_f": f_filter(0.2, nu, t),
             "example_high_f": f_filter(2.0, nu, t),
             "mixed_sign_observed": f_filter(0.2, nu, t) < 1 < f_filter(2.0, nu, t),
             "prediction": "crossover moves with t and f(nu)=nu*t crosses one at t=1/nu"}

    # Prediction 3: for a positive-curvature mode s>nu, Euler stability ends at
    # eta*=2/(s-nu). Verify by sweeping eta and checking the exact multiplier.
    s = 2.0
    predicted_eta = 2.0 / (s - nu)
    etas = np.linspace(0.01, 3.0, 10000)
    factors = np.abs(1.0 - etas * (s - nu))
    stable = etas[factors <= 1.0 + 1e-12]
    observed_eta = float(stable.max())
    stability = {"s": s, "nu": nu, "predicted_boundary": predicted_eta,
                 "observed_sweep_boundary": observed_eta,
                 "absolute_error": abs(predicted_eta - observed_eta),
                 "factor_below_boundary": float(1 - (observed_eta - .001) * (s - nu)),
                 "factor_above_boundary": float(1 - (observed_eta + .001) * (s - nu)),
                 "prediction": "|1-eta*(s-nu)|<=1; instability beyond eta=2/(s-nu)"}

    # Small-step Euler-vs-flow check: discrete filter converges to continuous filter.
    flow_rows = []
    for eta in [0.1, 0.05, 0.02, 0.01]:
        K = int(round(t / eta))
        ss = 0.7
        flow_rows.append({"eta": eta, "K": K,
                          "flow_g": g_filter(ss, nu, t),
                          "euler_g": discrete_g(ss, nu, eta, K),
                          "abs_error": abs(discrete_g(ss, nu, eta, K) - g_filter(ss, nu, t))})
    return {"removable_pole": pole, "mixed_sign_crossover": mixed,
            "euler_stability_boundary": stability, "euler_to_flow": flow_rows}


def regression_experiment():
    rng = np.random.default_rng(1285)
    d, n_train, n_test = 30, 120, 3000
    Q, _ = np.linalg.qr(rng.normal(size=(d, d)))
    eig = np.geomspace(0.05, 3.0, d)
    A = Q @ np.diag(np.sqrt(eig))
    X = rng.normal(size=(n_train, d)) @ A.T
    Xt = rng.normal(size=(n_test, d)) @ A.T
    beta_eig = np.zeros(d); beta_eig[:5] = [2.0, 1.5, 1.0, .8, .6]
    beta_true = Q @ beta_eig
    y = X @ beta_true + rng.normal(scale=.35, size=n_train)
    yt = Xt @ beta_true + rng.normal(scale=.35, size=n_test)
    S, b = X.T @ X / n_train, X.T @ y / n_train
    K, eta, nu = 30, .05, 1.0
    def run(kind):
        beta = np.zeros(d)
        for _ in range(K):
            grad = S @ beta - b
            if kind == 'gd': beta -= eta * grad
            elif kind == 'ns-gd': beta = (1 + eta * nu) * beta - eta * grad
            else: beta -= eta * (grad + .1 * beta)
        return beta
    out = {}
    for kind in ('gd', 'positive-ridge', 'ns-gd'):
        beta = run(kind)
        out[kind] = {'test_mse': float(np.mean((Xt @ beta - yt)**2)),
                     'train_mse': float(np.mean((X @ beta - y)**2)),
                     'parameter_norm': float(np.linalg.norm(beta))}
    out['settings'] = {'K': K, 'eta': eta, 'nu': nu, 'n_train': n_train, 'dimension': d}
    return out


def main():
    out = {'core_checks': core_checks(), 'regression': regression_experiment()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()