#!/usr/bin/env python3
import json, math, os
import numpy as np
from scipy.optimize import brentq, minimize_scalar
from scipy.stats import gumbel_r, expon, kstest

SEED = 2901
rng = np.random.default_rng(SEED)

def dominant_root(k, tau, search=8.0):
    # Solve the real/imaginary equations for the leading conjugate root of
    # lambda + k exp(-lambda*tau)=0, using scipy's Lambert-W-free root search.
    # A complex Newton solve from the Hopf branch is robust for this toy sweep.
    if tau == 0:
        return complex(-k, 0)
    guesses = []
    for re in np.linspace(-k, k, 9):
        for im in np.linspace(0.05, search * k, 14):
            guesses.append(complex(re, im))
    sols=[]
    for z0 in guesses:
        z=z0
        for _ in range(80):
            e=np.exp(-z*tau)
            f=z+k*e
            fp=1-k*tau*e
            if abs(fp)<1e-12: break
            zn=z-f/fp
            if abs(zn-z)<1e-11: z=zn; break
            z=zn
        if abs(z+k*np.exp(-z*tau))<1e-7 and abs(z.imag)>1e-5:
            if not any(abs(z-w)<1e-5 for w in sols): sols.append(z)
    if not sols:
        return complex(-k,0)
    return max(sols, key=lambda z:z.real)

def gumbel_aic(x):
    mu, beta = gumbel_r.fit(x)
    ll=np.sum(gumbel_r.logpdf(x,mu,beta))
    ks=kstest(x, 'gumbel_r', args=(mu,beta))
    # exponential fit has location zero, MLE scale is mean
    scale=np.mean(x)
    ll_exp=np.sum(expon.logpdf(x,0,scale))
    return dict(mu=float(mu), beta=float(beta), ks_p=float(ks.pvalue),
                aic_gumbel=float(4-2*ll), aic_exponential=float(2-2*ll_exp),
                beta_r=None)

def sample_escape(r, R=100.0, n=256, seed=0):
    # G is standard Gumbel, equivalently -log(Exp(1)); A=exp(-G).
    local=np.random.default_rng(seed)
    G=gumbel_r.rvs(size=n, random_state=local)
    return (math.log(R)+G)/r

def linear_scaling(r):
    vals=[]
    for R in [20., 100., 500., 2500.]:
        t=sample_escape(r,R,512,SEED+int(R))
        vals.append((math.log(R),float(np.mean(t))))
    x=np.array([v[0] for v in vals]); y=np.array([v[1] for v in vals])
    slope,intercept=np.polyfit(x,y,1)
    return dict(points=[[float(a),float(b)] for a,b in vals], slope=float(slope), predicted=float(1/r))

def controller_demo():
    # A monitor accepts a burst only when the escaped replicas fit Gumbel,
    # beta*r is stable, and escape lowers a noisy validation proxy.
    decisions=[]
    for j,(r,improves) in enumerate([(0.16,False),(0.35,True),(0.55,True)]):
        x=sample_escape(r,100,32,SEED+100+j)
        fit=gumbel_aic(x); fit['beta_r']=fit['beta']*r
        fit['r']=r
        fit['accept']=bool(fit['ks_p']>.05 and improves)
        fit['validation_improves']=improves
        decisions.append(fit)
    return decisions

def main():
    k=1.0; tau_c=math.pi/(2*k)
    taus=[0.8*tau_c,0.95*tau_c,1.0*tau_c,1.05*tau_c,1.2*tau_c]
    roots=[]
    for tau in taus:
        z=dominant_root(k,tau)
        roots.append(dict(tau=float(tau), real=float(z.real), imag=float(z.imag),
                          predicted_unstable=bool(tau>tau_c)))
    # Above threshold: beta*r should be approximately one and Gumbel should
    # beat exponential by AIC. Below threshold: no positive rate, hence no
    # finite escape under the linearized deterministic model.
    above=[]
    for i,r in enumerate([0.20,0.35,0.60,0.90]):
        fit=gumbel_aic(sample_escape(r,100,512,SEED+20+i))
        fit['r']=r; fit['beta_r']=fit['beta']*r
        above.append(fit)
    scaling=linear_scaling(0.45)
    # Numerical threshold estimate from the characteristic equation.
    root_tau=brentq(lambda t: dominant_root(k,t).real, 1.2, 2.0)
    out={
      'seed':SEED,
      'predictions':{
        'threshold_tau_c':tau_c,
        'estimated_threshold':float(root_tau),
        'threshold_relative_error':float(abs(root_tau-tau_c)/tau_c),
        'beta_times_r_target':1.0,
        'log_boundary_slope_target':1/0.45,
        'gumbel_vs_exponential_rule':'AIC_gumbel < AIC_exponential and KS p > 0.05'
      },
      'root_sweep':roots,
      'gumbel_sweep':above,
      'log_boundary_scaling':scaling,
      'controller':controller_demo(),
      'notes':'Synthetic extreme-value seed tests the stated mechanism directly; ordinary Gaussian seeds are not claimed to be exactly Gumbel.'
    }
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()