#!/usr/bin/env python3
"""Reproducible classical calculations for the explicitly chosen UFT candidate.

All numbers are dimensionless (c=hbar=m0=f=Lambda=1 unless stated).
No parameter has been fitted to experiment. This is not a device simulation.
Run: python reproduce.py
"""
from pathlib import Path
import json, platform, sys, hashlib
import numpy as np
import scipy
from scipy.integrate import solve_ivp, solve_bvp, quad, simpson
from scipy.linalg import eigh_tridiagonal
from scipy.special import expit
import sympy as sp
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

ROOT = Path(__file__).resolve().parent
DATA = ROOT / 'data'; FIG = ROOT / 'figures'
DATA.mkdir(exist_ok=True); FIG.mkdir(exist_ok=True)
plt.rcParams.update({'font.size': 10, 'figure.dpi': 130, 'axes.spines.top': False,
                     'axes.spines.right': False, 'savefig.bbox': 'tight'})
R = {'units': 'dimensionless; illustrative, not empirical',
     'software': {'python': platform.python_version(), 'numpy': np.__version__,
                  'scipy': scipy.__version__, 'sympy': sp.__version__,
                  'matplotlib': matplotlib.__version__}}

def csv(name, a, header):
    np.savetxt(DATA/name, a, delimiter=',', header=header, comments='')

def savefig(name):
    plt.savefig(FIG/(name+'.png')); plt.savefig(FIG/(name+'.svg')); plt.close()

def checked_ivp(*args, **kwargs):
    sol=solve_ivp(*args, **kwargs)
    if not sol.success: raise RuntimeError(sol.message)
    return sol

# 1. Independent symbolic checks, not a claim of automated tensor variation.
x,z,p,eps,f,Lam,m0=sp.symbols('x z p eps f Lam m0', positive=True)
U=Lam**4*x**2*(1-x)**2
m=m0*(1-x)**2
H=sp.sqrt(p**2+m**2)
mu=sp.sqrt(2)*Lam**2/f
kink=1/(1+sp.exp(-mu*z))
checks={
 'potential_first':sp.diff(U,x)-2*Lam**4*x*(1-x)*(1-2*x),
 'potential_second':sp.diff(U,x,2)-2*Lam**4*(1-6*x+6*x*x),
 'hamilton_velocity':sp.diff(H,p)-p/H,
 'hamilton_chi_source':sp.diff(H,x)-m*sp.diff(m,x)/H,
 'mass_squared_derivative':sp.diff(m*m,x)+4*m0*m0*(1-x)**3,
 'wall_equation':f*f*sp.diff(kink,z,2)-sp.diff(U,x).subs(x,kink),
 'wall_first_integral':f*f*sp.diff(kink,z)**2/2-U.subs(x,kink),
}
R['symbolic']={key: bool(sp.simplify(val)==0) for key,val in checks.items()}
assert all(R['symbolic'].values())
R['symbolic_scope']='Scalar derivatives, Hamiltonian derivatives and wall identities only; tensor variations are derived in the report.'
print('Stage 1 symbolic identities passed', flush=True)

# 2. Independent boundary value solve and fluctuation spectrum.
muval=np.sqrt(2.0); B=12.0
zz=np.linspace(-B,B,501)
def wall_rhs(z,y): return np.vstack([y[1],2*y[0]*(1-y[0])*(1-2*y[0])])
bc_left=expit(-muval*B); bc_right=expit(muval*B)
wall=solve_bvp(wall_rhs,lambda a,b:np.array([a[0]-bc_left,b[0]-bc_right]),
               zz,np.vstack([expit(zz),expit(zz)*(1-expit(zz))]),
               tol=1e-9,max_nodes=15000)
assert wall.success,wall.message
zz=np.linspace(-B,B,6001); wy=wall.sol(zz)
wall_exact=expit(muval*zz)
sigma=simpson(.5*wy[1]**2+wy[0]**2*(1-wy[0])**2,x=zz)
sigma_exact=np.sqrt(2)/6
spectra=[]
for n in [600,1200,2400]:
    zb=np.linspace(-B,B,n+2); h=zb[1]-zb[0]; ch=expit(muval*zb[1:-1])
    v=2*(1-6*ch+6*ch*ch)
    vals=eigh_tridiagonal(2/h**2+v,-np.ones(n-1)/h**2,
                         select='i',select_range=(0,3),eigvals_only=True)
    spectra.append({'N':n,'eigenvalues':vals.tolist(),
                    'zero_error':float(abs(vals[0])), 'shape_error':float(abs(vals[1]-1.5))})
R['wall']={'max_profile_error':float(np.max(abs(wy[0]-wall_exact))),
 'tension':float(sigma),'tension_exact':float(sigma_exact),
 'tension_abs_error':float(abs(sigma-sigma_exact)),
 'spectrum':spectra,'boundary':'Dirichlet at exact tail values, z in [-12,12]'}
assert R['wall']['max_profile_error']<1e-7
assert R['wall']['tension_abs_error']<1e-8
assert spectra[-1]['zero_error']<spectra[0]['zero_error']/10
csv('wall.csv',np.c_[zz,wy[0],wall_exact,wy[1]],'z,chi_bvp,chi_exact,dchi_bvp')
fig,ax=plt.subplots(1,2,figsize=(10,3.5))
ax[0].plot(zz,wy[0],label='BVP');ax[0].plot(zz,wall_exact,'--',label='analytic');ax[0].set(xlabel='z',ylabel='chi',xlim=(-6,6));ax[0].legend()
ax[1].plot(zz,.5*wy[1]**2+wy[0]**2*(1-wy[0])**2);ax[1].set(xlabel='z',ylabel='wall energy density',xlim=(-6,6))
savefig('01_wall')
print('Stage 2 wall and spectrum passed',flush=True)

# 3. Static weak-field smooth Gaussian source; finite difference refinement.
# (d_z^2 - mu^2) dchi = s(z), where s=S/f^2, f=1.
def yukawa_gaussian(zv,s0=-.02):
    fun=lambda q:np.exp(-muval*abs(zv-q))*s0*np.exp(-q*q/2)
    return -(quad(fun,-12,zv,epsabs=1e-12)[0]+quad(fun,zv,12,epsabs=1e-12)[0])/(2*muval)
weak=[]
from scipy.linalg import solve_banded
for n in [400,800,1600]:
    grid=np.linspace(-10,10,n+2);h=grid[1]-grid[0];zi=grid[1:-1]
    band=np.zeros((3,n));band[0,1:]=1/h**2;band[1,:]=-2/h**2-2;band[2,:-1]=1/h**2
    rhs=-.02*np.exp(-zi*zi/2)
    # Boundary values use the independent infinite-line Green function.
    rhs[0]-=yukawa_gaussian(-10)/h**2;rhs[-1]-=yukawa_gaussian(10)/h**2
    sol=solve_banded((1,1),band,rhs)
    idx=np.linspace(0,n-1,81,dtype=int); exact=np.array([yukawa_gaussian(v) for v in zi[idx]])
    weak.append({'N':n,'max_sample_error':float(np.max(abs(sol[idx]-exact))),
                 'peak':float(sol.max())})
R['weak_response']={'gaussian_source':{'s0':-.02,'width':1.0},'refinement':weak,
                    'mu':float(muval),'sign':'negative electric-dominated source raises chi for beta>0'}
assert weak[-1]['max_sample_error']<weak[0]['max_sample_error']/10
csv('weak_response.csv',np.c_[zi,sol,-sol],'z,electric_response,opposite_source_response')
plt.figure(figsize=(7,3.5));plt.plot(zi,sol,label='s = -0.02 exp(-z^2/2)');plt.plot(zi,-sol,label='opposite source');plt.xlabel('z');plt.ylabel('delta chi');plt.legend();savefig('02_weak_response')
print('Stage 3 weak response passed',flush=True)

# 4 and 6. Prescribed homogeneous conversion -> exact zero plateau -> recovery.
def smooth(u): return 6*u**5-15*u**4+10*u**3
def dsmooth(u): return 30*u*u*(u-1)**2
def control(t):
    if t<10: return smooth(t/10),dsmooth(t/10)/10
    if t<=15:return 1.0,0.0
    if t<25:return 1-smooth((t-15)/10),-dsmooth((t-15)/10)/10
    return 0.0,0.0
def particle_rhs(t,y):
    ch,ct=control(t);mass=(1-ch)**2;mt=-2*(1-ch)*ct
    E=np.hypot(y[1],mass)
    return [y[1]/E,0,mass*mt/E]
tt=np.linspace(0,25,2501)
cycle_runs=[]
for tol in [1e-8,1e-10,1e-12]:
    cyc=checked_ivp(particle_rhs,(0,25),[0,.3,0],method='DOP853',rtol=tol,
                    atol=tol*.01,t_eval=tt,max_step=.1)
    ch=np.array([control(t)[0] for t in tt]);mass=(1-ch)**2
    energy=np.hypot(cyc.y[1],mass);v=cyc.y[1]/energy
    cycle_runs.append({'rtol':tol,'max_work_balance_error':float(np.max(abs(energy-energy[0]-cyc.y[2]))),
                      'x_final':float(cyc.y[0,-1]),'p_final':float(cyc.y[1,-1])})
R['cycle']={'p0':.3,'initial_energy':float(energy[0]),'plateau_energy':.3,
 'initial_speed':float(v[0]),'plateau_speed':float(v[1200]),
 'energy_to_controller_on_entry':float(energy[0]-.3),'convergence':cycle_runs,
 'max_endpoint_error':float(abs(energy[-1]-energy[0]))}
assert cycle_runs[-1]['max_work_balance_error']<1e-9
csv('particle_cycle.csv',np.c_[tt,ch,mass,cyc.y[0],cyc.y[1],energy,v,cyc.y[2]],
    't,chi,mass,x,p,energy,velocity,work_on_particle')
fig,ax=plt.subplots(2,2,figsize=(10,6))
ax[0,0].plot(tt,ch,label='chi');ax[0,0].plot(tt,mass,label='m');ax[0,0].legend()
ax[0,1].plot(tt,v);ax[0,1].set_ylabel('v/c')
ax[1,0].plot(tt,energy,label='particle');ax[1,0].plot(tt,-cyc.y[2],label='controller gain');ax[1,0].legend()
ax[1,1].plot(tt,energy-energy[0]-cyc.y[2]);ax[1,1].set_ylabel('work balance residual')
for a in ax.flat:a.set_xlabel('t');a.axvspan(10,15,alpha=.10,color='green')
fig.tight_layout();savefig('03_conversion_cycle')

# Static spatial mass well. This is an external background, not a closed thruster.
def masswell(x):
    ch=.35*np.exp(-x*x/(2*1.5**2));cx=-x/1.5**2*ch
    return (1-ch)**2,-2*(1-ch)*cx
def well_rhs(t,y):
    mm,mx=masswell(y[0]);ee=np.hypot(y[1],mm)
    return [y[1]/ee,-mm*mx/ee]
tw=np.linspace(0,25,2001)
w=checked_ivp(well_rhs,(0,25),[-6,.6],method='DOP853',rtol=1e-12,atol=1e-14,t_eval=tw,max_step=.05)
mw=np.array([masswell(v)[0] for v in w.y[0]]);ew=np.hypot(w.y[1],mw)
R['mass_well']={'initial_energy':float(ew[0]),'max_energy_error':float(np.max(abs(ew-ew[0]))),
 'max_speed':float(np.max(w.y[1]/ew)),'final_position':float(w.y[0,-1]),
 'final_momentum':float(w.y[1,-1]),'note':'finite start/end positions, not identical asymptotic endpoints'}
assert R['mass_well']['max_energy_error']<1e-9
csv('mass_well.csv',np.c_[tw,w.y[0],w.y[1],mw,ew,w.y[1]/ew],'t,x,p,mass,energy,velocity')
print('Stages 4/6 particle cycle and spatial well passed',flush=True)

# 5. Dynamical controller: periodic 1+1D EM + chi + two smoothed neutral probes.
# Full closed reduced Hamiltonian. No prescribed time dependence or external drive.
# The Gaussian averaging of chi is an explicit UV regularization, not a claim of
# 3+1D Lorentz covariance at the regulator scale. L=2pi, f=Lambda=1, beta=.4.
def closed_run(N,tol,max_step=.03):
    box=2*np.pi; dx=box/N; xx=np.arange(N)*dx
    kk=np.fft.fftfreq(N,d=dx)*2*np.pi
    # Keep Nyquist derivative zero for a real collocation grid.
    kd=kk.copy();kd[N//2]=0
    beta=.4;sigma=.35;mass0=np.array([.03,.05])
    def deriv(a):return np.fft.ifft(1j*kd*np.fft.fft(a)).real
    def kernel(X):
        coef=np.exp(-.5*(sigma*kk)**2)*np.exp(-1j*kk*X)/dx
        coef[N//2]=0
        return np.fft.ifft(coef).real
    def split(y):return y[:N],y[N:2*N],y[2*N:3*N],y[3*N:4*N],y[4*N:4*N+2],y[4*N+2:]
    def components(y):
        aa,PP,ch,pi,X,pm=split(y);zz=np.exp(beta*ch);ax=deriv(aa);cx=deriv(ch)
        kernels=np.array([kernel(q) for q in X]);bar=kernels@ch*dx
        mm=mass0*(1-bar)**2;mp=-2*mass0*(1-bar);ee=np.hypot(pm,mm)
        return aa,PP,ch,pi,X,pm,zz,ax,cx,kernels,bar,mm,mp,ee
    def rhs(t,y):
        aa,PP,ch,pi,X,pm,zz,ax,cx,K,bar,mm,mp,ee=components(y)
        q=mm*mp/ee
        chi_force=deriv(cx)-2*ch*(1-ch)*(1-2*ch)+.5*beta*(PP*PP/zz-zz*ax*ax)-q@K
        pdot=-q*(K@cx*dx)
        return np.r_[PP/zz,deriv(zz*ax),pi,chi_force,pm/ee,pdot]
    def diag(y):
        aa,PP,ch,pi,X,pm,zz,ax,cx,K,bar,mm,mp,ee=components(y)
        Eem=dx*np.sum(.5*PP*PP/zz+.5*zz*ax*ax)
        Ech=dx*np.sum(.5*pi*pi+.5*cx*cx+ch*ch*(1-ch)**2)
        Pem=-dx*np.sum(PP*ax);Pch=-dx*np.sum(pi*cx)
        return np.r_[Eem+Ech+sum(ee),Pem+Pch+sum(pm),Pem,Pch,pm,Eem,Ech,ee]
    ch=.015*np.cos(xx-.4);aa=.15*np.sin(xx)+.05*np.sin(2*xx+.2)
    PP=-np.exp(beta*ch)*deriv(aa)
    y0=np.r_[aa,PP,ch,np.zeros(N),[1.,4.],[0.,0.]]
    times=np.linspace(0,10,501)
    sol=checked_ivp(rhs,(0,10),y0,method='DOP853',rtol=tol,atol=tol*.01,
                    max_step=max_step,t_eval=times)
    d=np.array([diag(y) for y in sol.y.T]);e0=d[0,0];P0=d[0,1]
    row={'N':N,'rtol':tol,'max_step':max_step,'initial_energy':float(e0),'initial_total_momentum':float(P0),
         'max_relative_energy_error':float(np.max(abs(d[:,0]-e0))/e0),
         'max_absolute_total_momentum_error':float(np.max(abs(d[:,1]-P0))),
         'final_probe_momenta':d[-1,4:6].tolist(),'final_em_momentum':float(d[-1,2]),
         'final_scalar_momentum':float(d[-1,3]),'max_abs_chi':float(np.max(abs(sol.y[2*N:3*N]))),
         'max_particle_vs_field_impulse_residual':float(np.max(abs((d[:,4]+d[:,5]-d[0,4]-d[0,5])+(d[:,2]+d[:,3]-d[0,2]-d[0,3])))),
         'nfev':sol.nfev}
    return row,times,d,sol

closed=[]
for N,tol in [(64,1e-8),(128,1e-10),(256,1e-11)]:
    row,tc,dc,solc=closed_run(N,tol);closed.append(row)
    csv(f'closed_system_N{N}.csv',np.c_[tc,dc],
        't,total_energy,total_momentum,em_momentum,scalar_momentum,probe1_momentum,probe2_momentum,em_energy,scalar_energy,probe1_energy,probe2_energy')
    print('Stage 5 closed system',json.dumps(row),flush=True)
R['closed_system']={'parameters':{'box_length':float(2*np.pi),'beta':.4,'sigma':.35,
 'probe_m0':[.03,.05],'duration':10.,'boundary':'periodic','method':'Fourier collocation + DOP853'},
 'runs':closed,'endpoint_probe_resolution_difference':float(np.max(abs(np.array(closed[-1]['final_probe_momenta'])-np.array(closed[-2]['final_probe_momenta']))))}
assert closed[-1]['max_relative_energy_error']<1e-7
assert closed[-1]['max_absolute_total_momentum_error']<1e-7
assert R['closed_system']['endpoint_probe_resolution_difference']<1e-7
fig,ax=plt.subplots(1,2,figsize=(11,4))
for j,label in [(1,'total'),(2,'EM'),(3,'scalar'),(4,'probe 1'),(5,'probe 2')]:
    ax[0].plot(tc,dc[:,j],label=label)
ax[0].set(xlabel='t',ylabel='momentum');ax[0].legend(fontsize=8)
ax[1].plot(tc,dc[:,4]+dc[:,5],label='probe momentum gain')
ax[1].plot(tc,dc[:,2]+dc[:,3]-dc[0,2]-dc[0,3],label='field momentum change')
ax[1].set(xlabel='t',ylabel='momentum exchange');ax[1].legend(fontsize=8)
fig.tight_layout();savefig('04_closed_momentum')

# With the main small maximum step, tolerance refinement already reaches the
# floating-point floor. Separately vary step limits with relaxed tolerances.
step_study=[]
for hmax in [.12,.06,.03]:
    sr,_,_,_=closed_run(64,1e-5,max_step=hmax)
    sr['probe_endpoint_difference_from_main']=float(np.max(abs(np.array(sr['final_probe_momenta'])-np.array(closed[-1]['final_probe_momenta']))))
    step_study.append(sr)
R['closed_system']['independent_step_refinement']=step_study
print('Step refinement:',json.dumps(step_study),flush=True)

# 6. Noncommuting endpoint limits; force switch-off and controlled recovery.
dd=np.logspace(-5,0,500)
fig,ax=plt.subplots(1,2,figsize=(10,3.5))
for a in [.1,1,10]:
    pm=a*dd**2;mass=dd**2
    ax[0].semilogx(dd,pm/np.hypot(pm,mass),label=f'p = {a} delta^2')
ax[0].semilogx(dd,.3/np.hypot(.3,dd**2),'k--',label='p = 0.3')
ax[0].set(xlabel='delta = 1-chi',ylabel='v/c');ax[0].legend(fontsize=8)
force=2*dd**3/np.hypot(.3,dd**2)
ax[1].loglog(dd,force,label='|F| for |grad chi|=1')
ax[1].loglog(dd,2*dd**3/.3,'--',label='asymptotic delta^3')
ax[1].set(xlabel='delta',ylabel='force');ax[1].legend(fontsize=8)
fig.tight_layout();savefig('05_endpoint')
R['endpoint']={'fixed_nonzero_p_limit':1.,'p_equals_a_delta_squared_limits':{str(a):float(a/np.sqrt(a*a+1)) for a in [.1,1,10]},
 'zero_state':'p=0 and m=0 has zero energy and is not a regular propagating particle state',
 'recovery':'demonstrated for prescribed homogeneous control; not an autonomous device solution',
 'direction_locking':'absent: no direct matter-null-flow interaction in the action'}
R['script_sha256']=hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
(DATA/'results.json').write_text(json.dumps(R,ensure_ascii=False,indent=2)+'\n')
(DATA/'environment.json').write_text(json.dumps(R['software'],indent=2)+'\n')
print('All programmed checks passed. Results:',DATA/'results.json',flush=True)
