"""Rest and rhythm in one phosphorylation reactor, with single-rate actuation."""
# EDITABLE INPUTS. Exact source concentrations/currents are in the named JSON.
SOURCE_FILE = 'finite_source_exact.json'
ORBIT_SEEDS_FILE = 'orbit_seeds.json' # Supply new guesses/section for substantially changed sources.
RATE_MULTIPLIERS = {}           # e.g. {'alpha1': 1.001}; exploratory, no inherited certificate.
ACTUATED_RATE = 'alpha1'        # Any one of eighteen names; published pulse uses alpha1.
PULSE_AMPLITUDE = .1
PULSE_PERIODS = 3
ON_PHASE_PI = 1.
OFF_PHASE_PI = 1/3
WAIT_AFTER_ON_PERIODS = 14
ON_SETTLING_PERIODS = 400
OFF_SETTLING_PERIODS = 600
CONCENTRATION_UNIT_MICROMOLAR = .1
TIME_UNIT_SECONDS = 100.
OBSERVATION_GAP = .5            # Model units; readout diagnostic, not a noise certificate.
OBSERVATION_ERROR = .01         # Absolute readout units per sample.

from pathlib import Path
import argparse,csv,hashlib,json,platform
import numpy as np
import scipy
import sympy as sp
from mpmath import mp
from clock import finite_model,patch_model,rate_list,RATE_NAMES,GH_CENTER,mp_JB,critical_pair,normal_form
from phos import Model,certify_hopf,bounds
from reactor import MassActionReactor,SwitchingExperiment
import local_certificate,input_rank

MANUSCRIPT_SHA256='efe1d3ed4fc2b556aed9c0eb6b3b62957ffeb32c9cfbf1edf67d93e11728a266'

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,value):(out/name).write_text(json.dumps(value,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)
    here=Path(__file__).parent;data=json.loads((here/SOURCE_FILE).read_text());x=list(map(sp.Rational,data['xstar']));q=list(map(sp.Rational,data['currents']));r=sp.Rational(data['r'])
    if len(x)!=12 or len(q)!=3 or min(*x,*q,r)<=0:raise ValueError('The three-site source needs twelve positive concentrations, three positive currents and a positive reverse ratio.')
    model=Model(3,x,q,[sp.Rational(1,100)]*3,[r,sp.Rational(1,100),sp.Rational(1,100)])
    rates=list(map(float,rate_list(model)))
    for name,multiplier in RATE_MULTIPLIERS.items():
        if name not in RATE_NAMES or not np.isfinite(multiplier) or multiplier<=0:raise ValueError('Each rate multiplier must be named, finite and positive.')
        rates[RATE_NAMES.index(name)]*=multiplier
    reactor=MassActionReactor(model,rates,ACTUATED_RATE)
    dump('exact_source.json',dict(concentrations=list(map(str,x)),currents=list(map(str,q)),rates=dict(zip(RATE_NAMES,map(str,rate_list(model)))),totals=list(map(str,model.totals())),configured_rates=dict(zip(RATE_NAMES,rates)),configured_equilibrium_residual=float(max(abs(reactor.f(0,np.zeros(9))))),scope='The exact reconstruction fixes a baseline equilibrium. Independent rate multipliers generally move it; reference certificates below apply to the published source only.'))
    rest_state,rest_residual=reactor.equilibrium()
    assert all(v==0 for v in model.field()) and model.Rm*model.P==sp.eye(9)
    eig=np.linalg.eigvals(reactor.jac(0,rest_state));dump('rest_stability.json',dict(equilibrium_chart=rest_state,equilibrium_residual=rest_residual,eigenvalues=[str(v) for v in eig],largest_real=float(max(eig.real)),e_folding_model_time=float(-1/max(eig.real)),scope='Numerical restricted Jacobian at the configured rates and totals; modified rates require a newly located rest state.'))
    print('Computing fresh generalized-Hopf interval certificate and all input ranks.',flush=True)
    gh=local_certificate.certify();dump('generalized_hopf_certificate.json',gh);dump('all_input_rank_certificate.json',input_rank.certify(gh))
    mp.dps=45;J,B=mp_JB(patch_model(sp.Rational(GH_CENTER[0]),sp.Rational(GH_CENTER[1])));lam,qq,ll,_=critical_pair(J,mp.mpf(GH_CENTER[2]));coeff,_=normal_form(J,B,lam,qq,ll)
    dump('independent_normal_form.json',dict(eigenvalue=str(lam),cubic=str(coeff[1]),quintic=str(coeff[2]),l2=str(mp.re(coeff[2])/mp.im(lam)),scope='Independent high-precision numerical recurrence; fresh directed certificate is stored separately.'))
    hc=certify_hopf(lambda rr:finite_model(rr));dump('finite_family_hopf.json',{k:bounds(hc[k]) for k in ('r','omega','l1','crossing')})
    imported=json.loads((here/'imported_orbit_certificate.json').read_text());dump('orbit_proof_status.json',dict(status='imported_not_rerun',source_sha256=hashlib.sha256((here/'imported_orbit_certificate.json').read_bytes()).hexdigest(),period=imported['T_interval'],readout=imported['readout_p2p'],gershgorin=imported['gershgorin'],scope='The manuscript independently validates the finite orbit. This run freshly solves numerical shooting and variational equations; it does not rerun the multiple-shooting interval proof or the Fourier proof.'))
    seeds=json.loads((here/ORBIT_SEEDS_FILE).read_text());orbits=[];curves=[]
    for name in ['stable','unstable']:
        orbit=reactor.shoot(seeds[name+'_y0'],seeds[name+'_T']);run=reactor.integrate(orbit['anchor'],orbit['period'],ledger=True,samples=2001)
        readout=run['species'][3]+run['species'][11];orbit.update(case=name,readout_range=float(np.ptp(readout)),min_species=run['minimum_concentration'],conservation_drift=run['conservation_drift'],kinase_turnover=float(sum(run['ledger'][[2,5,8]])),phosphatase_turnover=float(sum(run['ledger'][[11,14,17]])))
        orbits.append(orbit);curves.append((name,run));table(name+'_orbit.csv',[dict(time=t,**{f'x{i}':run['species'][i,j] for i in range(12)},readout=readout[j]) for j,t in enumerate(run['time'])])
        print(name,'period',orbit['period'],'range',orbit['readout_range'],'leading nontrivial multiplier',orbit['largest_nontrivial'],flush=True)
    if max(eig.real)>=0 or orbits[0]['largest_nontrivial']>=1:raise RuntimeError('Configured numerical rest/rhythm pair is not attracting; do not classify this as bistability.')
    dump('numerical_orbits.json',[{**o,'multipliers':[str(v) for v in o['multipliers']]} for o in orbits]);T=orbits[0]['period']
    experiment=SwitchingExperiment(reactor,T,seeds['XI'],seeds['ZETA'],PULSE_AMPLITUDE,PULSE_PERIODS)
    segments,on,off=experiment.run(np.pi*ON_PHASE_PI,np.pi*OFF_PHASE_PI,WAIT_AFTER_ON_PERIODS,rest_state)
    # Independent stiff solver cross-checks only the finite pulse endpoints.
    cross_on=experiment.pulse(on['chart'][:,0],np.pi*ON_PHASE_PI,method='Radau');cross_off=experiment.pulse(off['chart'][:,0],np.pi*OFF_PHASE_PI,method='Radau')
    on_fate=experiment.terminal_window(on['endpoint'],ON_SETTLING_PERIODS,rest_state);off_fate=experiment.terminal_window(off['endpoint'],OFF_SETTLING_PERIODS,rest_state)
    dump('switching_diagnostics.json',dict(actuated_rate=ACTUATED_RATE,amplitude=PULSE_AMPLITUDE,pulse_periods=PULSE_PERIODS,on=on_fate,off=off_fate,LSODA_Radau_on_difference=float(max(abs(on['endpoint']-cross_on['endpoint']))),LSODA_Radau_off_difference=float(max(abs(off['endpoint']-cross_off['endpoint']))),minimum_concentration=min(s['minimum_concentration'] for s in segments),conservation_drift=max(s['conservation_drift'] for s in segments),scope='Numerical finite-witness switching with a phase-triggered OFF pulse. Neither arbitrary-phase OFF nor finite error tolerance is certified.'))
    records=[];offset=0
    for i,s in enumerate(segments):
        for j,t in enumerate(s['time']):records.append(dict(segment=i,time=offset+t,periods=(offset+t)/T,readout=s['species'][3,j]+s['species'][11,j],relative_modulation=s['control'][j]))
        offset+=s['elapsed']
    table('switching_trace.csv',records)
    units=physical_reading(model,reactor,orbits[0],eig,T,rest_state);dump('physical_interpretation.json',units)
    # The normal-form scaling is a leading-order model, not a parameter sweep
    # of the full reactor or a guaranteed control-amplitude formula.
    d=.0444275025;b0=.012;a0=b0*b0/(8*d);scaling=[]
    for eps in [1.,.5,.25,.125,.0625]:
        a=-eps**2*a0;b=eps*b0;disc=np.sqrt(b*b+4*d*a);rp=np.sqrt((b+disc)/(2*d));rm=np.sqrt((b-disc)/(2*d))
        scaling.append(dict(epsilon=eps,inner_radius=rm,outer_radius=rp,coexistence_width=b*b/(4*d),sink_recovery_time=-1/a,cycle_recovery_time=1/(2*rp*rp*disc),relative_control_amplitude_scale=np.sqrt(eps)))
    table('leading_order_tradeoff.csv',scaling);plot(out,curves,records,scaling,reactor,T,rest_state)
    print('Switching final readout ranges:',on_fate['range'],off_fate['range'],flush=True)
    print('All rates and totals remain fixed between pulses; one rate alone is modulated during each pulse. Generalized-Hopf and rank certificates freshly replayed. Finite-orbit proof imported; switching numerical; Lean not rerun.',flush=True)
    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__,module_sha256={p.name:digest(p) for p in sorted(here.glob('*.py')) if p.name not in ['example.py','test_example.py']},input_sha256={p.name:digest(p) for p in sorted(here.glob('*.json')) if p.name!='release.json'},output_sha256={p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name!='run_metadata.json'}))

def physical_reading(model,reactor,orbit,eig,T,rest_state):
    c=CONCENTRATION_UNIT_MICROMOLAR;tau=TIME_UNIT_SECONDS
    if min(c,tau)<=0 or OBSERVATION_GAP<0 or OBSERVATION_ERROR<0:raise ValueError('Positive physical units and nonnegative observation errors required.')
    sampled=reactor.integrate(orbit['anchor'],T,samples=2001);slope=max(abs(reactor.f(t,sampled['chart'][:,i])[2]) for i,t in enumerate(sampled['time']))
    return dict(concentration_unit_micromolar=c,time_unit_seconds=tau,totals_micromolar=[float(v)*c for v in model.totals()],period_minutes=T*tau/60,readout_range_micromolar=orbit['readout_range']*c,cycle_recovery_periods=-1/np.log(orbit['largest_nontrivial']),sink_recovery_periods=-1/max(eig.real)/T,bound_phosphatase_fraction_at_rest=float(1-reactor.species(rest_state)[5]/float(model.totals()[1])),cycle_kinase_events=orbit['kinase_turnover'],rest_kinase_events=float(sum(reactor.flux(reactor.species(rest_state))[2:9:3]))*T,numerical_readout_slope_max=float(slope),sampling_range_diagnostic=orbit['readout_range']-slope*OBSERVATION_GAP-2*OBSERVATION_ERROR,scope='Uncalibrated dimensional interpretation. Slope maximum is sampled numerically, so the sampled-range calculation is diagnostic. Recovery e-folding times are not settling guarantees; no stochastic retention claim.')

def plot(out,curves,records,scaling,reactor,T,rest_state):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for name,run in curves:
        rd=run['species'][3]+run['species'][11];axs[0].plot(run['time'],rd,label=name+' cycle (numerical)');axs[1].plot(run['species'][5],rd,label=name+' cycle')
    rest=reactor.species(rest_state)[3]+reactor.species(rest_state)[11];axs[0].axhline(rest,color='black',ls=':',label='Stable rest');axs[1].plot(reactor.species(rest_state)[5],rest,'ko',label='Stable rest')
    axs[0].set(xlabel='Model time',ylabel='Fully phosphorylated pool S3 + D3',title='Steady and oscillating states at fixed\nrates and totals');axs[1].set(xlabel='Free phosphatase F',ylabel='Fully phosphorylated pool S3 + D3',title='Coexisting steady state and periodic\ntrajectories')
    for a in axs:a.legend(fontsize=8);a.grid(alpha=.2)
    fig.savefig(out/'coexisting_rest_rhythm.png',dpi=180);fig.savefig(out/'coexisting_rest_rhythm.svg');plt.close(fig)
    fig,axs=plt.subplots(2,1,figsize=(11,5),sharex=True,layout='constrained',gridspec_kw={'height_ratios':[3,1]})
    axs[0].plot([r['periods'] for r in records],[r['readout'] for r in records]);axs[0].axhline(rest,color='gray',ls=':');axs[0].set(ylabel='S3 + D3',title='Single-rate ON and phase-triggered OFF (numerical)')
    axs[1].plot([r['periods'] for r in records],[r['relative_modulation'] for r in records]);axs[1].set(xlabel='Time / autonomous cycle period',ylabel='Relative input')
    for a in axs:a.grid(alpha=.2)
    fig.savefig(out/'single_rate_switching.png',dpi=180);fig.savefig(out/'single_rate_switching.svg');plt.close(fig)

if __name__=='__main__':main()