import json
import math
import numpy as np


def detailed_balance_check(seed=0):
    rng = np.random.default_rng(seed)
    beta, nu = 1.7, 0.8
    vals = []
    for _ in range(1000):
        work, du = rng.normal(size=2)
        kf = nu * np.exp(beta * (work - du) / 2)
        # Reverse work is -work and reverse energy change is -du.
        kr = nu * np.exp(beta * (-work + du) / 2)
        vals.append(math.log(kf / kr) - beta * (work - du))
    return float(np.max(np.abs(vals)))


def exact_response(eta, lam, T, delta, P):
    """Expected excess response for x'=a*x+sqrt(2 eta T) z, q=dx^2/eta."""
    a = 1.0 - eta * lam
    v0 = 2.0 * eta * T / (1.0 - a*a)
    v = v0
    q0 = 4.0 * T / (1.0 + a)
    excess = 0.0
    for _ in range(P):
        v = a*a * v + 2.0 * eta * (T + delta)
        q = (1.0 - a)**2 * v / eta + 2.0 * (T + delta)
        excess += q - q0
    return excess / delta


def monte_carlo_response(eta, lam, T, delta, P, ntraj=30000, seed=1):
    rng = np.random.default_rng(seed)
    a = 1.0 - eta * lam
    v0 = 2.0 * eta * T / (1.0 - a*a)
    x = rng.normal(0.0, math.sqrt(v0), size=ntraj)
    # Estimate baseline q from independent stationary transitions.
    xb = rng.normal(0.0, math.sqrt(v0), size=ntraj)
    qss = np.mean(((a * xb + math.sqrt(2*eta*T)*rng.normal(size=ntraj) - xb)**2) / eta)
    total = 0.0
    for _ in range(P):
        xn = a*x + math.sqrt(2*eta*(T+delta))*rng.normal(size=ntraj)
        total += np.mean((xn-x)**2 / eta) - qss
        x = xn
    return total / delta


def run():
    lam, T, P = 1.0, 1.0, 20
    # Prediction 1: explicit Euler SGD is stable iff 0 < eta*lambda < 2.
    etas = np.array([*np.linspace(0.2, 1.8, 9), 1.9, 1.95, 1.99, 2.01, 2.1])
    rows = []
    for eta in etas:
        a = 1 - eta*lam
        C = exact_response(eta, lam, T, 0.05*T, P)
        tau = -1.0 / math.log(abs(a)) if abs(a) not in (0, 1) else (0.0 if a == 0 else float('inf'))
        rows.append({'eta': float(eta), 'eta_lambda': float(eta*lam), 'stable': bool(abs(a)<1), 'C': float(C), 'tau': float(tau)})
    # Predicted pole C ~ const/(2-eta lambda): test scaling at the stable edge.
    edge = [r for r in rows if 1.8 <= r['eta_lambda'] < 2.0]
    scaled = [r['C'] * (2-r['eta_lambda']) for r in edge]
    edge_cv = float(np.std(scaled) / np.mean(scaled))
    response_monotone_stable = bool(np.all(np.diff([r['C'] for r in rows if r['stable']]) > 0))
    # Prediction 2: relaxation time increases toward the stability boundary.
    stable_near = [r for r in rows if r['stable'] and r['eta_lambda'] > 1.2]
    taus = np.array([r['tau'] for r in stable_near])
    monotone_tau = bool(np.all(np.diff(taus) > 0))
    # Prediction 3: differential response is linear in pulse size for small pulses.
    deltas = np.array([0.01, 0.02, 0.05, 0.1]) * T
    Cs = np.array([exact_response(1.5, lam, T, d, P) for d in deltas])
    # Compare to infinitesimal numerical derivative and coefficient of variation.
    C0 = exact_response(1.5, lam, T, 1e-5, P)
    linear_rel_error = float(np.max(np.abs(Cs/C0 - 1)))
    mc = monte_carlo_response(1.5, lam, T, 0.05, P, ntraj=100000, seed=4)
    exact = exact_response(1.5, lam, T, 0.05, P)

    # Small optimizer comparison: noisy quadratic, fixed eta vs calorimetric trigger.
    def train(adaptive, seed):
        rng = np.random.default_rng(seed); x = 4.0; eta = 1.5; losses=[]; triggers=0
        for step in range(300):
            # Probe every 30 updates, using a short temperature pulse around current state.
            if adaptive and step > 0 and step % 30 == 0:
                c = exact_response(eta, lam, T, 0.05*T, 20)
                if c > 35.0:  # response threshold selected before trajectory, near edge
                    eta *= 0.7; triggers += 1
            grad = lam*x + rng.normal(0, math.sqrt(2*T/eta))
            x -= eta*grad
            losses.append(0.5*lam*x*x)
            if abs(x) > 1e6: return float('inf'), triggers
        return float(np.mean(losses[-30:])), triggers
    base = [train(False, s)[0] for s in range(20)]
    idea = [train(True, s)[0] for s in range(20)]
    triggers = [train(True, s)[1] for s in range(20)]
    out = {
      'detailed_balance_max_log_error': detailed_balance_check(),
      'prediction_stability_boundary': {'predicted_eta_lambda': 2.0, 'sweep_last_stable': max(r['eta_lambda'] for r in rows if r['stable']), 'sweep_first_unstable': min(r['eta_lambda'] for r in rows if not r['stable']), 'all_stable_below_2': all(r['stable'] for r in rows if r['eta_lambda'] < 2.0), 'first_unstable_is_above_2': all(not r['stable'] for r in rows if r['eta_lambda'] > 2.0)},
      'prediction_response_pole': {'predicted': 'C*(2-eta*lambda) approximately constant near edge', 'edge_cv': edge_cv, 'response_monotone_on_stable_sweep': response_monotone_stable, 'edge_rows': edge},
      'prediction_critical_slowing': {'predicted': 'tau increases as eta*lambda approaches 2 from below', 'monotone_on_1.2_to_edge': monotone_tau, 'taus': [float(x) for x in taus]},
      'prediction_pulse_linearity': {'deltas_over_T': [float(x) for x in deltas], 'C_values': [float(x) for x in Cs], 'max_relative_error_vs_infinitesimal': linear_rel_error},
      'monte_carlo_check': {'exact_C': exact, 'mc_C': mc, 'relative_error': abs(mc-exact)/abs(exact)},
      'mini_experiment': {'baseline_final_loss_mean': float(np.mean(base)), 'idea_final_loss_mean': float(np.mean(idea)), 'baseline_divergences': int(sum(np.isinf(base))), 'idea_divergences': int(sum(np.isinf(idea))), 'idea_trigger_count_mean': float(np.mean(triggers))},
      'rows': rows
    }
    with open('results.json','w') as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__': run()