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

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


def regulator(q, alpha, s):
    x = np.asarray(q, dtype=float) ** 2 / (s * s)
    return alpha * s * s / np.expm1(np.minimum(x, 700.0))


def crossover_root(alpha, s):
    # Unique positive solution of R_s(q)=q.
    lo, hi = 1e-12, max(10.0 * s, 10.0 * alpha * s * s + 1.0)
    for _ in range(100):
        mid = (lo + hi) / 2.0
        if regulator(mid, alpha, s) > mid:
            lo = mid
        else:
            hi = mid
    return (lo + hi) / 2.0


def stability_boundary(q, alpha, s):
    # For x_{t+1}=(1-eta*q/(q+R))x_t, stability is eta < 2/max(rate).
    rates = q / (q + regulator(q, alpha, s))
    return 2.0 / np.max(rates)


def exact_relaxation(q, alpha, s, eta):
    rate = eta * q / (q + regulator(q, alpha, s))
    multiplier = np.abs(1.0 - rate)
    t = -1.0 / np.log(np.maximum(multiplier, 1e-300))
    return t


def local_exponent(q, t):
    return -np.gradient(np.log(t), np.log(q))


def run_math_checks():
    alpha, s = 0.7, 1.0
    q = np.logspace(-4, 2, 1000)
    # Prediction A: R ~ alpha*s^4/q^2 at q << s; R/q -> 0 at q >> s.
    low = q < 0.01 * s
    high = q > 5.0 * s
    low_ratio = float(np.median(regulator(q[low], alpha, s) /
                               (alpha * s**4 / q[low]**2)))
    high_ratio = float(np.median(regulator(q[high], alpha, s) / q[high]))

    # Prediction B: crossover solves R=q and has scale q*=s*y(alpha*s).
    cases = [(0.2, 0.7), (0.7, 1.0), (1.5, 1.4), (3.0, 0.8)]
    cross = []
    for a, ss in cases:
        pred = crossover_root(a, ss)
        grid = np.logspace(-7, 4, 50000) * max(ss, a * ss * ss, 1.0)
        obs = float(grid[np.argmin(np.abs(regulator(grid, a, ss) - grid))])
        cross.append({'alpha': a, 's': ss, 'predicted_q': pred,
                      'observed_q': obs, 'relative_error': abs(obs-pred)/pred})
    # Dimensionless scaling check: q*/s is a function only of alpha*s.
    scaling = []
    for u in [0.15, 0.5, 1.0, 2.0]:
        vals = []
        for ss in [0.5, 1.0, 2.0]:
            aa = u / ss
            vals.append(crossover_root(aa, ss) / ss)
        scaling.append({'alpha_times_s': u, 'qstar_over_s': vals,
                        'spread': max(vals)-min(vals)})

    # Prediction C: Euler stability boundary is eta*=2/max_q q/(q+R).
    qs = np.logspace(-5, 3, 50000)
    eta_star = stability_boundary(qs, alpha, s)
    rates = qs / (qs + regulator(qs, alpha, s))
    test_etas = [0.99 * eta_star, 1.01 * eta_star]
    stability = []
    for eta in test_etas:
        multiplier = np.max(np.abs(1.0 - eta * rates))
        stability.append({'eta': eta, 'predicted_stable': bool(eta < eta_star),
                          'observed_max_multiplier': float(multiplier),
                          'observed_stable': bool(multiplier < 1.0)})

    # Claimed exponents: directly measure t(q) from the proposed update.
    # This is a falsification check; no exponent is imposed.
    t = exact_relaxation(qs, alpha, s, eta=0.5)
    z = local_exponent(qs, t)
    bands = [(1e-4, 3e-3), (0.03, 0.2), (5.0, 30.0)]
    exponents = [{'q_band': list(b), 'observed_z': float(np.median(z[(qs >= b[0]) & (qs <= b[1])]))}
                 for b in bands]
    return {'asymptotics': {'low_ratio_R_over_alpha_s4_over_q2': low_ratio,
                            'high_ratio_R_over_q': high_ratio},
            'crossovers': cross, 'dimensionless_scaling': scaling,
            'stability': {'eta_star': eta_star, 'tests': stability},
            'relaxation_exponents': exponents}


def adam_scalar(H, steps=400, lr=0.08):
    x = np.ones_like(H) * 2.0
    m = np.zeros_like(x); v = np.zeros_like(x)
    for t in range(1, steps + 1):
        g = H * x
        m = 0.9*m + 0.1*g
        v = 0.999*v + 0.001*g*g
        x -= lr * (m/(1-0.9**t)) / (np.sqrt(v/(1-0.999**t)) + 1e-8)
    return float(0.5*np.sum(H*x*x)), float(np.linalg.norm(x))


def smooth_rg(H, alpha=0.7, s=1.0, eta=0.5, steps=400):
    # Modewise curvature q=H and exact local damping a=q for this quadratic.
    x = np.ones_like(H) * 2.0
    for _ in range(steps):
        g = H*x
        x -= eta * g / (np.abs(H) + regulator(H, alpha, s) + 1e-8)
    return float(0.5*np.sum(H*x*x)), float(np.linalg.norm(x))


def mini_experiment():
    # Same diagonal quadratic, spanning shells; compare fixed AdamW-like and RG rule.
    H = np.logspace(-2, 2, 64)
    rows = []
    for s in [0.3, 1.0, 3.0]:
        b_loss, b_norm = adam_scalar(H)
        i_loss, i_norm = smooth_rg(H, s=s)
        rows.append({'s': s, 'adam_loss': b_loss, 'smooth_rg_loss': i_loss,
                     'adam_norm': b_norm, 'smooth_rg_norm': i_norm})
    return rows


def main():
    result = {'seed': SEED, 'math_checks': run_math_checks(),
              'mini_experiment': mini_experiment()}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()