#!/usr/bin/env python3
"""NetPhysics model 0.3 candidates. Python stdlib; simulations, not observations.

Run: python model_variants.py > expected-results.json
All units normalized. The earlier browser lab remains model 0.2.
"""
import json
import math

PARAMS = dict(L=1.0, alpha=0.1, beta=1.5, tau=15.0, zeta=0.12)

def drive(t, phase=None):
    phase = phase or ('train' if t < 30 else 'rest' if t < 50 else 'probe')
    return math.sin(t) if phase == 'train' else (0.0 if phase == 'rest' else 0.6*math.sin(1.25*(t-50)))

def rhs(t, y, writing='dissipation', damping='coefficient', phase=None):
    q, v, m = y[:3]
    p = PARAMS
    k = (1+p['beta']*m)/p['L']**2
    c = 2*p['zeta']*(math.sqrt(k) if damping == 'ratio' else 1/p['L'])
    u = drive(t, phase)
    intensity = {'drive': u*u, 'strain': q*q, 'dissipation': c*v*v}[writing]
    mdot = p['alpha']*intensity*(1-m)-m/p['tau']
    return [v, u-c*v-k*q, mdot, u*v, c*v*v, 0.5*q*q*p['beta']*mdot/p['L']**2]

def rk4(fun, t, y, h):
    a=fun(t,y); b=fun(t+h/2,[v+h*d/2 for v,d in zip(y,a)])
    c=fun(t+h/2,[v+h*d/2 for v,d in zip(y,b)])
    d=fun(t+h,[v+h*w for v,w in zip(y,c)])
    return [v+h*(aa+2*bb+2*cc+dd)/6 for v,aa,bb,cc,dd in zip(y,a,b,c,d)]

def simulate(writing='dissipation', damping='coefficient', dt=0.01):
    y=[0.0]*6; max_res=0.0; probe=0.0; entry=None
    n=round(80/dt); h=80/n
    assert n%8==0
    for i in range(n):
        t=i*h; phase='train' if i<n*3//8 else 'rest' if i<n*5//8 else 'probe'
        old=y
        y=rk4(lambda tt,yy:rhs(tt,yy,writing,damping,phase),t,y,h)
        assert all(math.isfinite(x) for x in y) and -1e-12<=y[2]<=1+1e-12
        k=(1+PARAMS['beta']*y[2])/PARAMS['L']**2
        energy=0.5*y[1]**2+0.5*k*y[0]**2
        max_res=max(max_res,abs(energy-y[3]+y[4]-y[5]))
        if i+1==n*5//8: entry=y[:3]
        if i>=n*5//8:probe+=h*(old[0]**2+y[0]**2)/2
    return dict(writing=writing,damping=damping,dt=h,probe_entry_q_p_m=entry,probe_rms=math.sqrt(probe/30),max_abs_energy_residual=max_res,final_state=y[:3])

def exact(t, near_boundary=False):
    m=0.9+0.05*math.sin(math.pi*t/2) if near_boundary else 0.4+0.1*math.sin(0.7*t)
    return [math.sin(t),math.cos(t),m]

def exact_dot(t, near_boundary=False):
    dm=0.025*math.pi*math.cos(math.pi*t/2) if near_boundary else 0.07*math.cos(0.7*t)
    return [math.cos(t),-math.sin(t),dm]

def manufactured(writing,damping,near_boundary=False):
    # Add a known source vector to all three equations. This verifies the
    # integrator on a forced exact solution, not the physical constitutive law.
    def forcing(t,y):
        f=rhs(t,y,writing,damping,'train')[:3]
        fstar=rhs(t,exact(t,near_boundary),writing,damping,'train')[:3]
        return [a+b-c for a,b,c in zip(f,exact_dot(t,near_boundary),fstar)]
    rows=[]
    for h in [0.2,0.1,0.05,0.025]:
        y=exact(0,near_boundary); err=0.0
        for i in range(round(2/h)):
            y=rk4(forcing,i*h,y,h)
            err=max(err,max(abs(a-b) for a,b in zip(y,exact((i+1)*h,near_boundary))))
        order=None if not rows else math.log(rows[-1]['max_state_error']/err,2)
        rows.append(dict(dt=h,max_state_error=err,observed_order=order))
    assert rows[-1]['observed_order'] > 3.7
    return dict(writing=writing,damping=damping,near_boundary=near_boundary,rows=rows)

def relaxation_check():
    alpha,tau,A=0.1,15.0,1.7
    rate=alpha*A*A+1/tau; target=alpha*A*A/rate
    # An exact embedding into generic dx/dt=(xeq(u)-x)/tau_k(u).
    delta=0.0
    for m in [0,0.1,0.5,1]:
        delta=max(delta,abs(alpha*A*A*(1-m)-m/tau-(target-m)*rate))
    # The original law imposes xeq=1-tau_k/tau and is even in drive.
    return dict(max_abs_rhs_difference=delta,rate=rate,target=target,
        caveat='Embedding in a generic relaxation family; not equality to all voltage-dependent CALM circuits or experimental fits.')

def matched_history():
    a,tau,target=0.5,30.0,0.3
    records=[]
    for A in [1.0,2.0]:
        rate=a*A*A+1/tau; plateau=a*A*A/rate
        T=-math.log(1-target/plateau)/rate
        m=plateau*(1-math.exp(-rate*T))
        records.append(dict(amplitude=A,duration=T,terminal_m=m))
    assert abs(records[0]['terminal_m']-records[1]['terminal_m'])<1e-14
    return dict(histories=records,scope='Equal terminal scalar states. Equal q and p and identical future input must also be imposed to obtain equal oscillator readouts. Any mechanical reset is an intervention with an unaccounted implementation cost.')

def main():
    variants=[simulate(w,d) for w in ['drive','strain','dissipation'] for d in ['ratio','coefficient']]
    convergence=[manufactured(w,d) for w in ['drive','strain','dissipation'] for d in ['ratio','coefficient']]
    for v in variants:
        fine=simulate(v['writing'],v['damping'],0.005)
        assert fine['max_abs_energy_residual'] < v['max_abs_energy_residual']/5
        v['fine_energy_residual']=fine['max_abs_energy_residual']
        v['fine_probe_rms']=fine['probe_rms']
        v['probe_rms_refinement_difference']=abs(v['probe_rms']-fine['probe_rms'])
    return dict(model_version='0.3-candidate-revision-1',status='Computational verification; no material fit, measured validation or novelty claim.',parameters=PARAMS,
        writing_laws={'drive':'alpha*u^2*(1-m)-m/tau','strain':'alpha*q^2*(1-m)-m/tau','dissipation':'alpha*c*p^2*(1-m)-m/tau'},
        damping_laws={'ratio':'c=2*zeta*sqrt(k)','coefficient':'c=2*zeta/L'},
        rest_caveat='State-written variants can keep writing during mechanical ring-down even when u=0. Their rest law is not generally pure exponential decay.',
        manufactured_solution=convergence,comparisons=variants,relaxation_family=relaxation_check(),matched_history=matched_history())



def saturation_check():
    # Unmodified scalar ODE under constant drive, compared with its exact solution.
    alpha,A,tau=2.0,4.0,20.0
    rate=alpha*A*A+1/tau; plateau=alpha*A*A/rate
    out=[]
    for initial in [0.0,1.0]:
        runs=[]
        for h in [0.02,0.01,0.005,0.0025]:
            y=[initial];err=0.0;lo=initial;hi=initial
            for i in range(round(1/h)):
                y=rk4(lambda t,y:[alpha*A*A*(1-y[0])-y[0]/tau],i*h,y,h)
                expected=plateau+(initial-plateau)*math.exp(-rate*(i+1)*h)
                err=max(err,abs(y[0]-expected));lo=min(lo,y[0]);hi=max(hi,y[0])
            assert 0<=lo<=hi<=1
            runs.append(dict(dt=h,max_abs_error=err,minimum_m=lo,maximum_m=hi,final_m=y[0]))
        out.append(dict(initial_m=initial,runs=runs))
    return dict(alpha=alpha,drive=A,tau=tau,equilibrium_m=plateau,runs=out,
                boundary_derivatives=dict(at_m0=alpha*A*A,at_m1=-1/tau),
                scope='Scalar constant-drive law only; no clipping, no added forcing, not a full coupled-material saturation validation.')

def revision_results():
    result=main()
    result['primary_candidate']={'id':'M03-D','writing':'dissipation','damping':'coefficient',
        'reason':'Fixed research choice for prospective comparisons, not selected by largest RMS or a measured fit.'}
    result['near_boundary_manufactured_solution']=[manufactured(w,d,True) for w in ['drive','strain','dissipation'] for d in ['ratio','coefficient']]
    result['scalar_saturation_check']=saturation_check()
    alpha=PARAMS['alpha'];PARAMS['alpha']=0.0
    baseline=simulate('drive','coefficient',0.005)
    PARAMS['alpha']=alpha
    result['no_memory_baseline']=baseline
    result['no_memory_baseline']['definition']='alpha=0, initially m=0; otherwise same parameters, drive and initial mechanics.'
    for item in result['comparisons']:
        item['rms_difference_from_no_memory']=item['fine_probe_rms']-baseline['probe_rms']
    return result

if __name__=='__main__': print(json.dumps(revision_results(),indent=2,allow_nan=False))
