"""Design prescribed phosphorylation equilibria, then test their physical dynamics."""
# EDIT THESE INPUTS. x>1 gives the prescribed free-enzyme ratio u=(x*x-1)/8.
AUXILIARY_ROOTS=('2','3','4','5','6')
ENZYME_TOTAL_RATIO='5'
CONCENTRATION_UNIT_MICROMOLAR='0.1'   # illustration, not a fitted biochemical system
TIME_UNIT_SECONDS='10'
SUBSTRATE_TOTAL_OFFSETS=('-1/500','0','1/500')
KINETIC_SPEED_FACTORS=('1/10','1','10')
TRAJECTORY_DURATION=100000.0
FREE_SUBSTRATE_PERTURBATION=0.01
FULL_SOURCE_CENSUS=False             # True extends exact Routh census from n<=6 to n<=10; slow

import argparse,csv,hashlib,json,platform
from pathlib import Path
from fractions import Fraction as Q
import numpy as np
import scipy
import sympy as sp
from scipy.optimize import brentq
import phos_sharp as ps
from construction import PrescribedStates,KineticFreedom,EquilibriumChart,stability,substrate_window
from reactor import Reactor
import paper_checks
MANUSCRIPT_SHA256='84259d8ba8da08bb2db4d3b1950eb2cb032902be1e42970ff64d6f7c1411d805'

def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'));out=parser.parse_args().output;out.mkdir(parents=True,exist_ok=True)
    def dump(name,x):(out/name).write_text(json.dumps(x,indent=2,default=lambda v:v.tolist() if isinstance(v,np.ndarray) else str(v))+'\n')
    def table(name,rows):
        with (out/name).open('w',newline='') as f:w=csv.DictWriter(f,fieldnames=list(rows[0]));w.writeheader();w.writerows(rows)
    print('Replaying source construction, interlacing, coefficient positivity, determinant and Routh checks.',flush=True)
    replay=paper_checks.replay(FULL_SOURCE_CENSUS);dump('source_arithmetic_replay.json',replay)
    rec=PrescribedStates(AUXILIARY_ROOTS,Q(ENZYME_TOTAL_RATIO)).construct();rates=rec['rates'];n=len(rates)
    dump('configured_design.json',rec);audit=EquilibriumChart(rates,rec['totals']).audit();dump('physical_root_audit.json',audit)
    if audit['positive_count']!=2*n-1:raise ArithmeticError('Independent physical root audit disagrees with construction.')
    rows=[]
    for j,s in enumerate(rec['states']):
        z=s['z'];H,dH,*_=ps.chart(rec['A'],rec['B'],rec['D'],rec['r'],Q(2),s['u']);exact=stability(rates,z)
        if any(ps.vector_field(rates,z)) or ps.totals(z)!=rec['totals']:raise ArithmeticError('Exact equilibrium validation failed.')
        J=ps.reduced_jacobian(rates,z);lead=float(max(np.linalg.eigvals(np.array(J,float)).real));flux=sum(row[2]*z[n+3+i] for i,row in enumerate(rates))
        rows.append(dict(index=j,x=str(s['x']),enzyme_ratio=str(s['u']),free_readout=float(z[n]),total_readout=float(z[n]+z[-1]),phosphorylation_flux=str(flux),relative_substrate_sensitivity=float(rec['totals'][2]/(s['u']*abs(dH))),slope=str(dH),leading_real_numerical=lead,**exact))
    dump('equilibrium_classification.json',rows)
    # Keep the paper's certified benchmark separate from configurable geometry.
    benchmark=PrescribedStates(('2','3','4','5','6'),Q(5)).construct();probes=list(map(Q,['1/20','69/128','4327/3200','7917/3200','6373/1600','399/80']))
    window=substrate_window(benchmark,probes,Q(1,500));dump('certified_substrate_window.json',window)
    wide=substrate_window(benchmark,probes,Q(1,100));dump('wider_window_attempt.json',wide)
    continuation=[]
    for offset in SUBSTRATE_TOTAL_OFFSETS:
        st=Q(12)+Q(offset);chart=EquilibriumChart(benchmark['rates'],(Q(10),Q(2),st));aa=chart.audit();dump('substrate_offset_'+str(Q(offset)).replace('/','_')+'.json',aa)
        for k,(lo,hi) in enumerate(zip(probes,probes[1:])):
            try:u=brentq(lambda v:float(ps.chart(benchmark['A'],benchmark['B'],benchmark['D'],Q(5),Q(2),v)[0])-float(st),float(lo),float(hi),xtol=1e-13)
            except ValueError:continue
            _,_,s,f,*_=ps.chart(benchmark['A'],benchmark['B'],benchmark['D'],Q(5),Q(2),u);z=ps.state_from(benchmark['A'],benchmark['B'],benchmark['D'],u,s,f)
            continuation.append(dict(substrate_total=str(st),branch=k,ratio_numerical=u,readout_numerical=float(z[3]+z[-1]),physical_root_count_exact=aa['positive_count']))
    table('substrate_total_continuation.csv',continuation)
    # A symmetric one-site system has u=r and would be lost by the regular chart.
    exceptional=EquilibriumChart([(2,1,1,2,1,1)],(1,1,2)).audit();dump('exceptional_ratio_example.json',exceptional)
    bad=list(map(Q,['1059/1000','1098/1000','1003/1000','1029/1000','1038/1000']));threshold=[]
    for r in map(Q,['64/100','66/100','675/1000']):
        b=ps.build(bad,r);threshold.append(dict(r=str(r),r_exceeds_every_requested_ratio=r>max(map(ps.ratio,bad)),all_coefficients_positive=b['positive'],leading_B_raw=str(b['Braw'][-1])))
    dump('coefficient_positivity_counterexample.json',threshold)
    sweep=[]
    for fac in map(Q,KINETIC_SPEED_FACTORS):
        freedom=KineticFreedom(tuple(row[1]*fac for row in rates),tuple(row[4]*fac for row in rates),tuple(row[5]*fac for row in rates));rr=freedom.realize(rec)
        for j,s in enumerate(rec['states']):
            if any(ps.vector_field(rr,s['z'])):raise ArithmeticError('Kinetic freedom moved a prescribed equilibrium.')
            lead=float(max(np.linalg.eigvals(np.array(ps.reduced_jacobian(rr,s['z']),float)).real));sweep.append(dict(speed_factor=str(fac),state=j,unstable=stability(rr,s['z']).get('unstable'),leading_real_numerical=lead))
    table('same_equilibria_different_speeds.csv',sweep)
    reactor=Reactor(rates,rec['totals']);curves=[];diagnostics=[]
    if not 0<FREE_SUBSTRATE_PERTURBATION<.1:raise ValueError('Use a small positive fraction below 0.1.')
    for row,state in zip(rows,rec['states']):
        if row.get('unstable')!=0:continue
        z=np.array(state['z'],float);q=z[reactor.indices].copy();q[n-1]*=1-FREE_SUBSTRATE_PERTURBATION;initial=reactor.species(q);run=reactor.integrate(initial,TRAJECTORY_DURATION);other=reactor.integrate(initial,TRAJECTORY_DURATION,samples=101,method='Radau');curves.append((row['index'],run,float(z[n]+z[-1])))
        diagnostics.append(dict(state=row['index'],initial_distance=float(np.linalg.norm(initial-z)),final_distance=float(np.linalg.norm(run['endpoint']-z)),solver_difference=float(max(abs(run['endpoint']-other['endpoint']))),conservation_drift=run['conservation_drift'],minimum_concentration=run['minimum_concentration']))
        table(f'full_reactor_sink_{row["index"]}.csv',[dict(time_model=t,**{label:float(run['species'][l,k]) for l,label in enumerate(reactor.labels)}) for k,t in enumerate(run['time'])])
    dump('numerical_recovery_checks.json',diagnostics)
    table('physical_rate_constants.csv',[dict(site=i+1,**{label:float(value/(Q(TIME_UNIT_SECONDS)*(Q(CONCENTRATION_UNIT_MICROMOLAR) if j in [0,3] else 1))) for j,(label,value) in enumerate(zip(['a_per_uM_s','b_per_s','c_per_s','alpha_per_uM_s','beta_per_s','gamma_per_s'],row))}) for i,row in enumerate(rates)])
    plot(out,rec,rows,curves,continuation)
    print('Positive physical equilibria:',audit['positive_count'],'; Routh unstable counts:',[r.get('unstable') for r in rows],flush=True)
    print('Paper substrate window:',window['status'],'; wider attempted window:',wide['status'],flush=True)
    print('Source checks:',replay['checks_passed'],'; exact Routh census n <=',replay['routh_census_nmax'],flush=True)
    print('The all-n equilibrium count is a paper theorem. Even-index stability beyond the finite census is not proved here. Numerical trajectories are not basin or robustness certificates; Lean not rerun.',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(),numpy=np.__version__,scipy=scipy.__version__,sympy=sp.__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,rec,rows,curves,continuation):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for row in rows:
        stable=row.get('unstable')==0;axs[0].plot(float(Q(row['enzyme_ratio'])),row['total_readout']*float(Q(CONCENTRATION_UNIT_MICROMOLAR)),'o',mfc='C0' if stable else 'white',mec='C0',ms=8)
    axs[0].set(xlabel='Free enzyme ratio E/F',ylabel='Fully phosphorylated pool / micromolar',title='Prescribed equilibrium ratios and their\nstability')
    for j,run,base in curves:axs[1].semilogx(run['time'][1:]*float(Q(TIME_UNIT_SECONDS))/3600,((run['species'][len(rec['rates'])]+run['species'][-1])/base-1)[1:],label=f'State {j}')
    axs[1].set(xlabel='Time / hours (declared units)',ylabel='Relative readout displacement',title='Full mass-action recovery (numerical)');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'prescribed_states_recovery.png',dpi=180);fig.savefig(out/'prescribed_states_recovery.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for k in range(5):
        points=[r for r in continuation if r['branch']==k];axs[0].plot([float(Q(r['substrate_total'])) for r in points],[r['ratio_numerical'] for r in points],'o-',label=f'Branch {k}')
    axs[0].set(xlabel='Total substrate (model units)',ylabel='Equilibrium E/F (numerical)',title='Five roots persist in the certified window');axs[0].legend(fontsize=7)
    xs=np.arange(1,13);axs[1].plot(xs,2*xs-1,'o-',label='Maximum positive equilibria (theorem)');axs[1].plot(np.arange(1,7),np.arange(1,7),'s',label='Stable states freshly checked, n <= 6');axs[1].set(xlabel='Sites n',ylabel='Count',title='Equilibrium theorem and finite stability\ncounts');axs[1].legend(fontsize=7)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'window_and_count.png',dpi=180);fig.savefig(out/'window_and_count.svg');plt.close(fig)

if __name__=='__main__':main()