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

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


def simulate_front(D=0.25, r=1.0, K=1.0, kappa=0.0, nx=1501, nt=3000, dx=1.0,
                   dt=0.2, seed=0, impulse_x=150, impulse_amplitude=None, noise=True):
    """Euler-Maruyama Fisher-KPP field with reflecting boundaries and clipping."""
    if D > 0:
        dt = min(dt, 0.8 * dx * dx / (2 * D))
    rng = np.random.default_rng(seed)
    u = np.zeros(nx, dtype=np.float64)
    u[impulse_x] = K if impulse_amplitude is None else impulse_amplitude
    history = [u.copy()]
    for _ in range(nt):
        lap = np.zeros_like(u)
        if D > 0:
            lap[1:-1] = (u[2:] - 2*u[1:-1] + u[:-2]) / dx**2
            lap[0] = 2*(u[1]-u[0]) / dx**2
            lap[-1] = 2*(u[-2]-u[-1]) / dx**2
        react = r * u * (1.0 - u / K)
        stochastic = 0.0
        if noise and kappa > 0 and r > 0:
            # Formula in the prompt, with spatial white-noise scaling 1/dx.
            var = np.maximum(0.0, 2*kappa*r*dt/dx * u*(1-u/K))
            stochastic = np.sqrt(var) * rng.standard_normal(nx)
        u = np.clip(u + dt*(D*lap + react) + stochastic, 0.0, K)
        history.append(u.copy())
    return np.asarray(history), dt


def front_position(profile, threshold=0.1, dx=1.0):
    ids = np.where(profile >= threshold)[0]
    if len(ids) == 0: return 0.0
    j = ids[-1]
    if j >= len(profile)-1: return j*dx
    if j == 0: return 0.0
    a,b=profile[j],profile[j+1]
    return (j + ((threshold-a)/(b-a) if b != a else 0.0))*dx


def estimate_speed(hist, dt, threshold=0.1, start_frac=.30, end_frac=.75):
    ts=np.arange(len(hist))*dt
    xs=np.array([front_position(p,threshold) for p in hist])
    lo,hi=int(len(ts)*start_frac),int(len(ts)*end_frac)
    return float(np.polyfit(ts[lo:hi],xs[lo:hi],1)[0]),ts,xs


def math_checks():
    # Prediction 1: v* = 2 sqrt(D r), tested after a long transient.
    speed_rows=[]
    for D in (0.10,0.25,0.50):
        h,dt=simulate_front(D=D,r=1.0,nt=3000)
        v,_,_=estimate_speed(h,dt)
        pred=2*math.sqrt(D)
        speed_rows.append({'D':D,'predicted':pred,'observed':v,'relative_error':abs(v-pred)/pred})
    # Prediction 2: linearized zero-state growth rate is r. Use a small perturbation.
    stability=[]
    for r in (-0.5,0.0,0.5):
        h,dt=simulate_front(D=0.0,r=r,K=1.0,nt=30,impulse_amplitude=1e-3)
        amp=h[:,150]
        slope=float(np.polyfit(np.arange(1,20)*dt,np.log(np.maximum(amp[1:20],1e-30)),1)[0])
        stability.append({'r':r,'predicted_log_slope':r,'observed_log_slope':slope,'absolute_error':abs(slope-r)})
    # Prediction 3: increasing kappa broadens the stochastic front.
    widths=[]; means=[]
    for kap in (0.0,0.15,0.5):
        vals=[]; early=[]
        for seed in range(24):
            h,dt=simulate_front(D=.25,r=.8,kappa=kap,nt=120,seed=seed,nx=401,impulse_x=40)
            p=h[-1]; right=np.where(p>=.1)[0]; core=np.where(p>=.9)[0]
            vals.append((right[-1]-core[0]) if len(right) and len(core) else 0)
            early.append(float(h[15].sum()))
        widths.append({'kappa':kap,'width_mean':float(np.mean(vals)),'width_std':float(np.std(vals))})
        means.append({'kappa':kap,'early_total_mean':float(np.mean(early))})
    return {'speed_scaling':speed_rows,'stability_boundary':stability,'noise_broadening':widths,'noise_early_mean':means}


def task_experiment(device='cpu'):
    import torch
    torch.manual_seed(SEED); np.random.seed(SEED)
    n,length,d,layers=96,16,24,6
    x=torch.randint(0,2,(n,length),dtype=torch.float32,device=device); y=(x.sum(1)%2).long()
    class Net(torch.nn.Module):
        def __init__(self,gated):
            super().__init__(); self.gated=gated
            self.inp=torch.nn.Linear(1,d)
            self.blocks=torch.nn.ModuleList([torch.nn.Sequential(torch.nn.Linear(d,d),torch.nn.Tanh(),torch.nn.Linear(d,d)) for _ in range(layers)])
            self.out=torch.nn.Linear(d,2)
            if gated: self.rho=torch.nn.Parameter(torch.tensor(-.5))
        def forward(self,x):
            h=self.inp(x.unsqueeze(-1))
            if self.gated:
                u=torch.clamp(x*.5+.1,0,1).mean(0).detach().clone()
                D=.12; r=torch.nn.functional.softplus(self.rho); dt=.4
                for b in self.blocks:
                    lap=torch.zeros_like(u); lap[1:-1]=u[2:]-2*u[1:-1]+u[:-2]; lap[0]=u[1]-u[0]; lap[-1]=u[-2]-u[-1]
                    u=torch.clamp(u+dt*(D*lap+r*u*(1-u)),0,1)
                    h=h+(u/(1+u))[None,:,None]*b(h)
            else:
                for b in self.blocks: h=h+b(h)
            return self.out(h.mean(1))
    out={}
    for gated in (False,True):
        torch.manual_seed(SEED); net=Net(gated).to(device); opt=torch.optim.Adam(net.parameters(),lr=.01); losses=[]
        for _ in range(180):
            opt.zero_grad(); loss=torch.nn.functional.cross_entropy(net(x),y); loss.backward(); opt.step(); losses.append(float(loss))
        out['gated' if gated else 'baseline']={'final_loss':losses[-1],'best_loss':min(losses),'params':sum(p.numel() for p in net.parameters())}
    return out


def main():
    result={'math':math_checks()}
    try:
        import torch
        device='cuda' if torch.cuda.is_available() else 'cpu'
        try: result['task']=task_experiment(device)
        except Exception as e: result['task_error']=str(e); result['task']=task_experiment('cpu')
        result['device']=device
    except Exception as e: result['task_error']=str(e)
    Path('results.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))

if __name__=='__main__': main()