"""Full multisite phosphorylation: identical static behavior, different dynamics.

Run python example.py. Inputs below are dimensionless paper witnesses, not fitted
biochemistry. Output JSON distinguishes freshly certified facts from numerics.
"""
# ------------------------- EDITABLE INPUTS -------------------------
SUBSTRATE=('2','12','1/5','2/5')
FREE_KINASE='23/10'
FREE_PHOSPHATASE='2/5'
KINASE_COMPLEXES=('23/50','3/10','17/10')
PHOSPHATASE_COMPLEXES=('23','9/50','4')
PRODUCTIVE_FLUX=('1/10','32/5','6/5')
KINASE_REVERSE=('1/100',)*3
PHOSPHATASE_REVERSE=('1','1/100','1/100')
REVERSE_PARAMETER_SWEEP=(1.,1.3,1.5,2.)
TRANSIENT_TIME=1000.
ORBIT_PARAMETER=1.
ORBIT_PERIOD_GUESS=39.59
ADDED_SITE_LOADS=('1/100',)
REVERSIBLE_DRIVE=20.
CONCENTRATION_UNIT_MICROMOLAR=.1
TIME_UNIT_SECONDS=100.
OUTPUT_DIRECTORY='outputs'
# Certificate reference cases A, H, C2 reduction and site addition are fixed
# paper reproductions; edits above change the separate exploration model.
# ------------------------------------------------------------------
import argparse,csv,hashlib,json,platform
from pathlib import Path
import numpy as np
import sympy as sp
from scipy.signal import find_peaks
import phos
from reactor import EquilibriumDesign,MassActionReactor,ConstantBinding

MANUSCRIPT_SHA256='7a0bbd174ad9a01c8765e169ed4dc1541cb7680dd3fd18d8f4529569f52c87d9'


def design():
    return EquilibriumDesign(SUBSTRATE,FREE_KINASE,FREE_PHOSPHATASE,KINASE_COMPLEXES,PHOSPHATASE_COMPLEXES,PRODUCTIVE_FLUX,KINASE_REVERSE,PHOSPHATASE_REVERSE)


def certificate_record(c):
    keys=('t','r','omega','crossing','l1','eps_derivative')
    return dict(evidence='Fresh exact polynomial isolation and outward-rounded rational arithmetic, using the attributed manuscript kernel; not Lean.',
        **{k:phos.bounds(c[k]) for k in keys if k in c},routh_first_column=[phos.bounds(v) for v in c['routh']],
        p0=[str(v) for v in c['p0']],p1=[str(v) for v in c['p1']],
        right_eigenvector=[phos.bounds(v) for v in c['right']],left_eigenvector=[phos.bounds(v) for v in c['left']],
        constant_coefficient_positive=c['const_coeff_positive'])


def certify_references():
    results={};raw={}
    for name,builder in [('A',lambda r:phos.witness('A',r)),('H',lambda r:phos.witness('H',r)),('A_kinase_clamped',lambda r:phos.witness('A',r,clampE=True)),('A_added_site',lambda r:phos.extended('A',r,list(ADDED_SITE_LOADS)))]:
        m=builder(sp.Rational(1,2));norm=2*m.n+2 if name=='A_added_site' else None
        c=phos.certify_hopf(builder,norm_index=norm);raw[name]=c;results[name]=certificate_record(c)
    reduced=phos.certify_reduced(lambda r:phos.witness('A',r),fast=4,norm_index=8)
    results['A_C2_eliminated']=certificate_record(reduced)
    results['A_C2_eliminated']['scope']='Local algebraic reduction with induced quadratic AND cubic terms; not a trajectory-error guarantee.'
    assert raw['A']['l1'].b<0 and raw['H']['l1'].a>0 and raw['A_added_site']['l1'].b<0
    assert raw['A_kinase_clamped']['l1'].a>0 and reduced['l1'].b<0
    c=raw['A'];m=phos.witness('A',sp.Rational(1,2));L,M,A=m.blocks();Jstatic=L*A.inv()*M
    cp=Jstatic.charpoly().all_coeffs()
    assert all(v>0 for v in cp) and cp[1]*cp[2]>cp[3]
    other=phos.witness('A',sp.Rational(7,3));L2,M2,A2=other.blocks();assert Jstatic==L2*A2.inv()*M2
    sensors={'S3':[0,0,1,0,0,0,0,0,-1],'S3+D3':[0,0,1,0,0,0,0,0,0]}
    gains,identity=phos.gain_constants(lambda r:phos.witness('A',r),c,sensors)
    from interval_arithmetic import I,C,dot
    a=-c['crossing'].re;scale=(a/(-c['omega']*c['l1'])).sqrt();amplitudes={}
    for name,v in sensors.items():
        z=dot([C(x) for x in v],c['right']);amplitudes[name]=phos.bounds(4*(z.re*z.re+z.im*z.im).sqrt()*scale)
    results['mechanism']=dict(static_characteristic=[str(v) for v in cp],static_independence='Exact at two distinct family values; row-factor cancellation proves the full static field independence.',
        resonance_constants={k:phos.bounds(v) for k,v in gains.items()},rank_one_identity=phos.bounds(identity*C(c['r'])),
        amplitude_prefactors=amplitudes,radial_recovery_prefactor=phos.bounds(2*a),
        all_n_scope='The positive-load site-addition theorem is imported; this run freshly certifies the configured finite extension. n<=2 means no Hopf, not no periodic orbit.',
        finite_amplitude_scope='Manuscript interval orbit proofs and validated pulse proof are not rerun here. Simulations, shooting, Floquet spectra and pulse trajectories below are numerical only.')
    return results,raw


def encoded(v):
    if isinstance(v,np.ndarray):return v.tolist()
    if isinstance(v,np.generic):return v.item()
    if isinstance(v,complex):return dict(real=v.real,imag=v.imag)
    raise TypeError(type(v).__name__)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',default=OUTPUT_DIRECTORY)
    out=Path(parser.parse_args().output);out.mkdir(exist_ok=True)
    def dump(name,data):(out/name).write_text(json.dumps(data,indent=2,default=encoded)+'\n',encoding='utf-8')
    def table(name,header,rows):
        with (out/name).open('w',newline='',encoding='utf-8') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    reference,raw=certify_references();dump('hopf_certificates.json',reference)
    print('Fresh rational Hopf certificates: A attracting, H subcritical, kinase clamp subcritical, site addition attracting.',flush=True)
    d=design();base=MassActionReactor(d.exact());n=base.n
    if n!=3:raise ValueError('The teaching plots use a three-site starting design; the reactor classes accept arbitrary site count.')
    rows=[];trajectories={};static_fields=[]
    y0=np.zeros(base.dim);y0[2]=.01
    for r in REVERSE_PARAMETER_SWEEP:
        model=MassActionReactor(d.with_reverse(r).exact());J=model.jacobian(0,np.zeros(model.dim));eig=np.linalg.eigvals(J)
        sol,X=model.integrate(y0,TRANSIENT_TIME)
        z,uf=model.static_complexes(np.array([.001,-.001,.001]));static_fields.append(uf)
        rows.append(dict(r=r,equilibrium_residual=float(max(abs(model.field(0,np.zeros(model.dim))))),leading_real=float(max(eig.real)),
            static_leading_real=float(max(np.linalg.eigvals(model.static_jacobian()).real)),static_local_field=uf,
            minimum_concentration=float(X.min()),conservation_drift=float(np.max(abs(model.conservation@X-(model.conservation@model.xstar)[:,None])))))
        trajectories[str(r)]=(sol,X)
        table(f'trajectory_r_{r}.csv',['time',*model.labels],zip(sol.t,*X))
    assert np.max(np.ptp(np.array(static_fields),axis=0))<1e-10
    dump('static_and_dynamic.json',dict(evidence='Exact reference static comparison above; editable sweep and nonlinear implicit solves are numerical.',sweep=rows))
    # Locate a repeated maximum after a long transient, then solve the return map.
    orbitmodel=MassActionReactor(d.with_reverse(ORBIT_PARAMETER).exact())
    settled,X=orbitmodel.integrate(y0,TRANSIENT_TIME,6001)
    peaks,_=find_peaks(X[3]);peaks=peaks[settled.t[peaks]>.6*TRANSIENT_TIME]
    if len(peaks)<2:raise RuntimeError('No repeated oscillation detected. Adjust the parameter, transient or disable this experiment in your own driver.')
    seed=settled.y[:,peaks[-1]];periodguess=settled.t[peaks[-1]]-settled.t[peaks[-2]]
    anchor,T,multipliers,residual=orbitmodel.periodic_orbit(seed,periodguess if periodguess>0 else ORBIT_PERIOD_GUESS)
    cyc,CX=orbitmodel.integrate(anchor,T,1601,ledger=True);led=cyc.y[orbitmodel.dim:,-1]
    trivial=int(np.argmin(abs(multipliers-1)));nontrivial=np.delete(multipliers,trivial)
    orbit=dict(evidence='Numerical shooting and variational integration, NOT a validated orbit or minimal-period proof.',r=ORBIT_PARAMETER,period=T,
        shooting_residual=residual,multipliers=multipliers,largest_nontrivial_modulus=float(max(abs(nontrivial))),
        minimum_concentration=float(CX.min()),S3_peak_to_peak=float(np.ptp(CX[3])),
        kinase_turnovers=led[:n],phosphatase_turnovers=led[n:],paired_flux_difference=led[:n]-led[n:],
        total_forward_ATP_turnover=float(sum(led[:n])),turnover_per_substrate=float(sum(led[:n])/(orbitmodel.conservation@orbitmodel.xstar)[2]))
    table('periodic_orbit.csv',['time',*orbitmodel.labels,*[f'integral_{s}' for s in ['C1','C2','C3','D1','D2','D3']]],zip(cyc.t,*CX,*cyc.y[orbitmodel.dim:]))
    dump('periodic_orbit.json',orbit)
    # Uniform slowing preserves equilibrium and static behavior; clamping changes
    # the physical conservation contract by replacing enzyme pools with buffers.
    interventions=[]
    for label,model in [('both_free_enzymes_buffered',MassActionReactor(d.with_reverse('1').exact(True,True))),
        ('complexes_faster',MassActionReactor(d.with_reverse('1').exact(),relaxation=.5)),
        ('driven_reversible',MassActionReactor(d.with_reverse('1').exact(),drive=REVERSIBLE_DRIVE))]:
        sol,X=model.integrate(y0,TRANSIENT_TIME)
        interventions.append(dict(name=label,terminal_peak_to_peak=float(np.ptp(X[3,-301:])),minimum_concentration=float(X.min()),
            equilibrium_residual_at_original_design=float(max(abs(model.field(0,np.zeros(model.dim)))))))
        trajectories[label]=(sol,X)
    # Detailed-balance construction: potentials vary with M; not fixed chemistry
    # under changing ATP concentration. Reverse currents keep the three totals.
    rm=MassActionReactor(d.with_reverse('1').exact(),drive=REVERSIBLE_DRIVE);mm=rm.model
    mu=np.zeros(len(rm.xstar));mu[:n+1]=np.arange(n+1)*REVERSIBLE_DRIVE
    mu[rm.cpx]=mu[rm.lev]-np.log(rm.kon/rm.koff)
    binding_error=np.log(rm.kon/rm.koff)-(mu[rm.lev]+mu[rm.enz]-mu[rm.cpx])
    cat_potential=mu[rm.cpx]-mu[rm.out]-mu[rm.enz];cat_potential[:n]+=2*REVERSIBLE_DRIVE
    catalysis_error=np.log(rm.kcat/rm.reverse)-cat_potential
    revsol,revX=rm.integrate(anchor,TRANSIENT_TIME,1201,ledger=True)
    # No period proof here: accumulated net fuel is reported over a finite window.
    net=float(sum(revsol.y[rm.dim:rm.dim+n,-1]))
    dump('interventions.json',dict(experiments=interventions,thermodynamics=dict(driving_M=REVERSIBLE_DRIVE,
        local_detailed_balance_residual=float(max(np.max(abs(binding_error)),np.max(abs(catalysis_error)))),
        finite_window_net_ATP=net,cycle_affinity=2*REVERSIBLE_DRIVE,
        scope='Finite-time net consumption is not an entropy-production identity away from a closed periodic orbit. The theorem gives no numerical M0. Potentials are redesigned together with M; ATP/ADP/Pi remain chemostatted.')))
    c=raw['A'];rc=phos.mid(c['r']);omega=phos.mid(c['omega']);sensor=np.array([0,0,1,0,0,0,0,0,-1]);resonance=[]
    for delta in [.01,.03,.1,.3,1.]:
        m=MassActionReactor(phos.witness('A',str(rc+delta)));J=m.jacobian(0,np.zeros(9));b=np.zeros(9);b[6]=(1+rc+delta)*.1
        gain=abs(sensor@np.linalg.solve(1j*omega*np.eye(9)-J,b))
        bounds=reference['mechanism']['resonance_constants']['S3'];K=float((sp.Rational(bounds[0])+sp.Rational(bounds[1]))/2)
        exactlaw=(1+rc+delta)*K/delta
        resonance.append(dict(delta=delta,direct_gain=gain,rank_one_law=exactlaw,relative_residual=abs(gain/exactlaw-1)))
    pulses=[]
    for r in ['13/10','3/2']:
        pm=MassActionReactor(phos.witness('A',r),ConstantBinding(1.5));ps,PX=pm.integrate(np.zeros(9),.05,51)
        pulses.append(dict(r=r,D1_terminal=float(PX[9,-1]),scope='Nominal numerical pulse only; no preparation/actuator uncertainty enclosure.'))
    dump('dynamic_readouts.json',dict(resonance=resonance,nominal_pulses=pulses,scope='The resonance law is exact for the linearized response at the certified frequency. Finite forcing, measurement error and experiment feasibility require separate analysis.'))
    dump('units.json',dict(calibrated=False,concentration_unit_micromolar=CONCENTRATION_UNIT_MICROMOLAR,time_unit_seconds=TIME_UNIT_SECONDS,
        totals_micromolar=(base.conservation@base.xstar)*CONCENTRATION_UNIT_MICROMOLAR,period_minutes=T*TIME_UNIT_SECONDS/60,
        free_phosphatase_molecules_per_femtoliter=float(base.xstar[base.model.iF]*CONCENTRATION_UNIT_MICROMOLAR*602.214076),
        caveat='Illustrative scale only; heavy enzyme loading, broad catalytic rates and low molecule numbers can invalidate a biological or deterministic interpretation.'))
    plot(out,trajectories,cyc,CX,rows,resonance)
    print(f'Numerical r={ORBIT_PARAMETER:g} orbit: period {T:.9f}; S3 swing {orbit["S3_peak_to_peak"]:.6f}; largest nontrivial multiplier {max(abs(nontrivial)):.6f}.',flush=True)
    print('Static fields agree while the full stability changes. Reverse catalysis, buffering, fuel and dynamic readouts evaluated.',flush=True)
    here=Path(__file__).parent;digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,source_sha256=digest(Path(__file__)),python=platform.python_version(),
        module_sha256={p.name:digest(p) for p in sorted(here.glob('*.py')) if p.name not in ['example.py','test_example.py']},
        output_sha256={p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name!='run_metadata.json'}))


def plot(out,trajectories,cyc,X,rows,resonance):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for key in [str(REVERSE_PARAMETER_SWEEP[1]),str(REVERSE_PARAMETER_SWEEP[2])]:
        sol,Z=trajectories[key];axs[0].plot(sol.t,Z[3],label='r = '+key)
    axs[0].set(xlabel='Time (model units)',ylabel='Free fully phosphorylated S3',title='Phosphorylation trajectories at matched\nequilibria');axs[0].legend()
    axs[1].plot([r['r'] for r in rows],[r['leading_real'] for r in rows],'o-',label='Full mass action')
    axs[1].plot([r['r'] for r in rows],[r['static_leading_real'] for r in rows],'s--',label='Static elimination')
    axs[1].axhline(0,color='black',lw=.7);axs[1].set(xlabel='D1 dissociation ratio r',ylabel='Largest real eigenvalue',title='Stability of full and complex-eliminated\nmodels');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'mechanism.png',dpi=180);fig.savefig(out/'mechanism.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for i in [3,6,8,9,11]:axs[0].semilogy(cyc.t/cyc.t[-1],X[i],label=['S0','S1','S2','S3','E','F','C1','C2','C3','D1','D2','D3'][i])
    axs[0].set(xlabel='Fraction of numerical period',ylabel='Concentration',title='An attracting numerical cycle');axs[0].legend(fontsize=8,ncol=2)
    axs[1].loglog([r['delta'] for r in resonance],[r['direct_gain'] for r in resonance],'o',label='Matrix response')
    axs[1].loglog([r['delta'] for r in resonance],[r['rank_one_law'] for r in resonance],label='Exact resonance law')
    axs[1].set(xlabel='Distance above Hopf threshold',ylabel='Linear S3 gain at Hopf frequency',title='Linear substrate response near the Hopf\nthreshold');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'cycle_and_readout.png',dpi=180);fig.savefig(out/'cycle_and_readout.svg');plt.close(fig)


if __name__=='__main__':main()
