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


def cubic_check(r=0.2, L=1.0, omega=1.0, rho=0.8):
    a1 = 2*r/L
    a2 = (r/L)**2 + omega**2
    kappa = 1.5*omega/L
    boundary = a1*a2/kappa
    rows = []
    for frac in [0.2, 0.8, 0.99, 1.0, 1.01, 1.5]:
        g = frac*boundary
        roots = np.roots([1., a1, a2, kappa*g])
        rows.append({'fraction': frac, 'chi': float(kappa*g/(a1*a2)),
                     'max_real_root': float(np.max(roots.real)),
                     'roots': [[float(z.real), float(z.imag)] for z in roots]})
    return {'a1': a1, 'a2': a2, 'kappa': kappa,
            'boundary_g': boundary, 'cap_g': rho*boundary, 'rows': rows,
            'crossing_observed': bool(rows[2]['max_real_root'] < 0 and rows[4]['max_real_root'] > 0)}


def run_momentum(lr, momentum, H, steps=250, cap=False, rho=0.8):
    # Two-dimensional quadratic, with one deliberately high-curvature direction.
    x = np.array([1., 1.])
    v = np.zeros(2)
    losses, gains, chis = [], [], []
    # Use the cubic coefficients as the specified local stability model.
    r, L, omega = .2, 1., 1.
    a1, a2, kappa = 2*r/L, (r/L)**2 + omega**2, 1.5*omega/L
    gmax = rho*a1*a2/kappa
    effective_lr = min(lr, gmax/max(H)) if cap else lr
    for _ in range(steps):
        grad = H*x
        g = effective_lr*max(H)
        chi = kappa*g/(a1*a2)
        v = momentum*v + grad
        x = x - effective_lr*v
        losses.append(0.5*float(x@(H*x))); gains.append(g); chis.append(chi)
        if not np.all(np.isfinite(x)) or losses[-1] > 1e100:
            break
    tail = losses[-50:]
    return {'final_loss': losses[-1] if losses else float('inf'),
            'min_loss': min(losses) if losses else float('inf'),
            'loss_std_tail': float(np.std(tail)) if tail else float('inf'),
            'steps': len(losses), 'effective_lr': effective_lr,
            'max_chi': max(chis) if chis else float('inf'),
            'diverged': len(losses) < steps or not np.isfinite(losses[-1])}


def main():
    random.seed(7); np.random.seed(7)
    check = cubic_check()
    # Sweep a range around the ordinary quadratic momentum stability boundary.
    H = np.array([1., 20.])
    settings = []
    for lr in [.12, .2, .4, .8, 1.2]:
        settings.append({'lr': lr,
                         'baseline': run_momentum(lr, .9, H, cap=False),
                         'routh_cap': run_momentum(lr, .9, H, cap=True)})
    out = {'cubic_check': check, 'quadratic_experiment': settings}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()