import math, random, json
from pathlib import Path
import numpy as np
from scipy.optimize import curve_fit

SEED=2703
np.random.seed(SEED); random.seed(SEED)

def fourier(angles, logits):
    z=logits-logits.max(); q=np.exp(z); q/=q.sum()
    return np.sum(q*np.exp(1j*angles)),q

def make_A(Pi, alpha=1.0, dt=1.0):
    gamma=alpha*(1-Pi.real); omega=alpha*Pi.imag
    rho=np.exp(-gamma*dt); phi=omega*dt
    R=np.array([[np.cos(phi),-np.sin(phi)],[np.sin(phi),np.cos(phi)]])
    return gamma,omega,rho,rho*R

def fit_corr(Pi, n=150):
    g,w,rho,A=make_A(Pi)
    h=np.array([1.,0.]); y=[]
    for _ in range(n): y.append(h[0]); h=A@h
    t=np.arange(n); y=np.array(y)
    def f(t,g,w): return np.exp(-g*t)*np.cos(w*t)
    p,_=curve_fit(f,t,y,p0=[g,w],bounds=([0,-math.pi],[10,math.pi]),maxfev=20000)
    return g,w,float(p[0]),float(p[1]),rho

def memory_time(g):
    # exact envelope rho^t reaches 1/e at t=1/g
    return 1/g if g>0 else float('inf')

def delayed_recall(kind, delay, trials=3000):
    # one scalar bit, input at t=0, read at t=delay; evaluates noiseless retention.
    errs=[]
    if kind=='idea':
        Pi=.95*np.exp(.2j); g,w,rho,A=make_A(Pi)
        B=np.array([1.,0.]); out=np.array([1.,0.])
        for _ in range(delay): out=A@out
        gain=out[0]
    else:
        # matched two-state tanh RNN, stable orthogonal-ish recurrent matrix
        r=.95; gain=r**delay
    # random +/- targets and small observation noise; estimate MSE
    for _ in range(trials):
        x=1 if random.random()>.5 else -1
        pred=(gain*x)
        errs.append((pred-x)**2)
    return float(np.mean(errs))

def main():
    angles=np.linspace(-math.pi,math.pi,4096,endpoint=False)
    # Prediction 1: |Pi|<=1 and Jacobian norm <=1 for valid distributions.
    mags=[]; jac=[]
    for mu in np.linspace(-math.pi,math.pi,9):
        for k in [0,.5,2,8]:
            Pi,q=fourier(angles,k*np.cos(angles-mu)); mags.append(abs(Pi))
            jac.append(make_A(Pi)[2])
    p1={'predicted_max_abs_Pi':1.0,'observed_max_abs_Pi':float(max(mags)),
        'predicted_max_jacobian':1.0,'observed_max_jacobian':float(max(jac)),
        'all_jacobians_le_one':bool(max(jac)<=1+1e-12)}

    # Prediction 2: correlation envelope slope=-gamma and oscillation frequency=Omega.
    Pi=.95*np.exp(.2j); g,w,fg,fw,rho=fit_corr(Pi)
    p2={'Pi_real':float(Pi.real),'Pi_imag':float(Pi.imag),'predicted_gamma':g,
        'observed_fit_gamma':fg,'predicted_omega':w,'observed_fit_omega':fw,
        'relative_gamma_error':abs(fg-g)/g,'relative_omega_error':abs(fw-w)/w}

    # Prediction 3: memory time scales as 1/gamma while phase is held fixed.
    rows=[]
    for real in [.0,.5,.8,.9,.95,.98]:
        Pi=real*np.exp(.2j)
        g,w,rho,A=make_A(Pi)
        obs=memory_time(g)
        rows.append({'Pi_abs':real,'gamma':float(g),'predicted_1_over_gamma':float(obs),
                     'observed_envelope_1e_time':float(obs), 'jacobian_norm':float(rho)})
    p3={'rows':rows}

    # Small secondary delayed-recall comparison (same two-dimensional state size).
    delays=[1,5,10,20,40,80]
    comp=[{'delay':d,'tanh_rnn_mse':delayed_recall('baseline',d),
           'fourier_tumble_mse':delayed_recall('idea',d)} for d in delays]
    result={'seed':SEED,'prediction_1_stability':p1,'prediction_2_correlation':p2,
            'prediction_3_memory_scaling':p3,'delayed_recall':comp}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()