"""Construct equilibrium geometry, choose kinetics, and verify stable-state capacity."""
# EDITABLE INPUTS: an exact designed family and the separate published benchmark.
DESIGN_SITES = 3
DESIGN_STAGE_BUDGET = 24
CUSTOM_ROOTS = ('2','7/2','5','13/2','8')
ENZYME_TOTAL_RATIO = '9'
CATALYTIC_SCALES = ('1','1','1')
KINASE_DISSOCIATION = ('1','1','1')
PHOSPHATASE_DISSOCIATION = ('1','1','1')
BINDING_SCALE = '1'
INDEPENDENT_PARAMETER_RADIUS = '1/671088640000'
READOUT_ERROR_MICROMOLAR = '1/1000'
CONCENTRATION_UNIT_MICROMOLAR = '1/10'
TIME_UNIT_SECONDS = '10'
TRAJECTORY_DURATION = 50000.
TRAJECTORY_RELATIVE_PERTURBATION = .01 # Exploration outside the tiny certified initial balls.
FULL_PAPER_REPLAY = True             # Includes all 48 rational states for n=2,...,7.

from pathlib import Path
from fractions import Fraction as Q
from dataclasses import asdict
import argparse,csv,hashlib,json,math,platform
import numpy as np
import scipy
from mpmath import mp
import phos_sharp as ps
import phos_capacity as pc
import paper_checks
from model import Geometry,Kinetics,Reactor,stability,ordered_design
from recovery import LocalRecovery

MANUSCRIPT_SHA256='2cda97c29c4bf2773242ede84a0c57667a8aba2e1e90466844a6b5cdce51242f'

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 asdict(v) if hasattr(v,'__dataclass_fields__') 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 exact construction, all rational witnesses, and local interval bounds.',flush=True)
    replay=paper_checks.replay(FULL_PAPER_REPLAY);dump('source_arithmetic_replay.json',replay)
    design=ordered_design(DESIGN_SITES,DESIGN_STAGE_BUDGET);dump('ordered_design.json',design)
    geometry=Geometry(CUSTOM_ROOTS,Q(ENZYME_TOTAL_RATIO));rec=geometry.construct();kinetics=Kinetics(CATALYTIC_SCALES,KINASE_DISSOCIATION,PHOSPHATASE_DISSOCIATION,Q(BINDING_SCALE));rates=kinetics.rates(geometry)
    dump('configured_geometry.json',dict(polynomials={k:list(map(str,rec[k])) for k in ['N','J','A','B','D']},roots=rec['xs'],rates=rates,totals=rec['totals'],states=[dict(state=s['z'],**stability(rates,s['z'])) for s in rec['states']],scope='Distinct roots give distinct equilibria. Geometry is preserved by admitted kinetic retuning; their stability is checked separately, not inferred from count.'))
    here=Path(__file__).parent;benchmark=json.loads((here/'worked_example.json').read_text());certificate=json.loads((here/'operational_certificate.json').read_text());rates=paper_checks.rate_tuple_list(benchmark['record']['rates']);states=[list(map(Q,p['state'])) for p in benchmark['record']['profiles']];totals=tuple(map(Q,certificate['totals']));n=3
    geometry=Geometry(tuple(benchmark['selection']['xs']),Q(benchmark['selection']['r']));rec=geometry.construct();reactor=Reactor(rates,totals)
    profiles=[]
    for j,z in enumerate(states):
        exact=stability(rates,z);matrix=ps.reduced_jacobian(rates,z);eig=np.linalg.eigvals(np.array(matrix,float));u=z[4]/z[5];phi,slope,*_=ps.chart(rec['A'],rec['B'],rec['D'],Q(9),Q(2),u)
        if any(ps.vector_field(rates,z)) or ps.totals(z)!=totals:raise ArithmeticError('Exact benchmark equilibrium failed.')
        profiles.append(dict(index=j,x=str(rec['xs'][j]),u=str(u),state=list(map(str,z)),readout_exact=str(z[3]+z[11]),readout_micromolar=float((z[3]+z[11])*Q(CONCENTRATION_UNIT_MICROMOLAR)),slope=str(slope),leading_real_numerical=float(max(eig.real)),**exact))
    dump('benchmark_equilibria.json',profiles)
    recovery=LocalRecovery(certificate,states).evaluate(INDEPENDENT_PARAMETER_RADIUS,READOUT_ERROR_MICROMOLAR,CONCENTRATION_UNIT_MICROMOLAR,TIME_UNIT_SECONDS);dump('local_recovery_certificate.json',recovery)
    wider=LocalRecovery(certificate,states).evaluate('1/1000',READOUT_ERROR_MICROMOLAR,CONCENTRATION_UNIT_MICROMOLAR,TIME_UNIT_SECONDS);dump('larger_tolerance_attempt.json',wider)
    # The unbound-only mass matrix is intentionally compared with the correct
    # loaded matrix. It is not silently used as the dynamics of this reactor.
    co=pc.coalesced(3,Q(4));state=co['states'][0]['z'];M,h,K=pc.loaded(3,state,Q(8),Q(2));R=pc.incidence(3);wrong=pc.matmul(pc.matmul([[-R[i][j]+h[j] for j in range(4)] for i in range(3)],[[Q(int(i==j))/state[i] for j in range(4)] for i in range(4)]),pc.transpose(R))
    dump('loaded_inventory_comparison.json',dict(mass_matrix=M,loaded_K=K,unbound_only_K=wrong,entrywise_difference=max(abs(K[i][j]-wrong[i][j]) for i in range(3) for j in range(3)),scope='The loaded mass matrix retains enzyme-bound substrate. Replacing it by free-substrate diagonal changes the slow model.'))
    limiting=[]
    for nn in range(2,21):
        v=(3*nn-2)*Q(2,9)**(nn-1);gap=v*(50*nn-45)/(9*(14-8*v));limiting.append(dict(sites=nn,feedback_deficit_exact=str(gap),feedback_deficit=float(gap),asymptote=float(Q(25,21)*nn*nn*Q(2,9)**(nn-1))))
    table('limiting_feedback_margin.csv',limiting)
    kinetic=[]
    for scale in [Q(1),Q(1,10),Q(1,100)]:
        recg=Geometry(tuple(Q(3)+Q(1,2)*(j-2) for j in range(5)),Q(4));kin=Kinetics.from_coalesced_currents(3,Q(4),binding_scale=scale);rr=kin.rates(recg)
        for j,s in enumerate(recg.construct()['states']):
            assert all(v==0 for v in ps.vector_field(rr,s['z']));e=np.linalg.eigvals(np.array(ps.reduced_jacobian(rr,s['z']),float));kinetic.append(dict(binding_scale=str(scale),state=j,unstable=stability(rr,s['z']).get('unstable'),leading_real=float(max(e.real))))
    table('same_geometry_kinetic_sweep.csv',kinetic)
    splitting=split_law();dump('splitting_law.json',splitting)
    curves=[];summaries=[]
    if not 0<TRAJECTORY_RELATIVE_PERTURBATION<.05:raise ValueError('Use a small positive exploratory perturbation below five percent.')
    for j in [0,2,4]:
        z=np.array(states[j],float);q=z[reactor.indices].copy();q[2]*=1-TRAJECTORY_RELATIVE_PERTURBATION;initial=reactor.species(q)
        run=reactor.integrate(initial,TRAJECTORY_DURATION);other=reactor.integrate(initial,TRAJECTORY_DURATION,method='Radau',samples=101);rd=run['species'][3]+run['species'][11];curves.append((j,run))
        summaries.append(dict(sink=j,relative_perturbation=-TRAJECTORY_RELATIVE_PERTURBATION,final_distance=float(np.linalg.norm(run['endpoint']-z)),solver_difference=float(max(abs(run['endpoint']-other['endpoint']))),conservation_drift=run['conservation_drift'],minimum=run['minimum_concentration'],scope='Free S3 is decreased and free S0 increased by the same amount. Numerical recovery outside the very small certified initial neighbourhood.'))
        table(f'recovery_sink_{j}.csv',[dict(time_model=t,readout_micromolar=rd[k]*float(Q(CONCENTRATION_UNIT_MICROMOLAR)),**{label:run['species'][l,k] for l,label in enumerate(reactor.labels)}) for k,t in enumerate(run['time'])])
    dump('nonlinear_recovery_diagnostics.json',summaries)
    table('benchmark_rates_physical.csv',[dict(site=i+1,association_kinase=float(row[0]/(Q(CONCENTRATION_UNIT_MICROMOLAR)*Q(TIME_UNIT_SECONDS))),dissociation_kinase=float(row[1]/Q(TIME_UNIT_SECONDS)),catalysis_kinase=float(row[2]/Q(TIME_UNIT_SECONDS)),association_phosphatase=float(row[3]/(Q(CONCENTRATION_UNIT_MICROMOLAR)*Q(TIME_UNIT_SECONDS))),dissociation_phosphatase=float(row[4]/Q(TIME_UNIT_SECONDS)),catalysis_phosphatase=float(row[5]/Q(TIME_UNIT_SECONDS))) for i,row in enumerate(rates)])
    plot(out,profiles,recovery,curves,limiting)
    print('Exact benchmark unstable counts:',[p['unstable'] for p in profiles],flush=True)
    print('Fresh independent-parameter certificate:',recovery['all_certified'],'; tenfold envelope seconds:',[s.get('tenfold_seconds_upper') for s in recovery['sinks']],flush=True)
    print('Source checks:',replay['checks'],'; equilibrium capacity is 2n-1, stable-equilibrium capacity n. Neither theorem bounds periodic attractors or stochastic memory retention. 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 split_law():
    mp.dps=60;n=3;r=Q(4);co=pc.coalesced(n,r);kin=Kinetics.from_coalesced_currents(n,r);rates=kin.rates(Geometry((Q(3),)*5,r));cp=ps.charpoly(ps.reduced_jacobian(rates,co['states'][0]['z']));theta=abs(cp[-2]);M=ps.val(co['B'],Q(1))-ps.val(co['D'],Q(1));pref=math.prod((b+c)*al*ga for a,b,c,al,be,ga in rates)
    a=4*(2*r*ps.val(co['N'],r)+ps.val(co['J'],r))*M/((r-1)*ps.val(co['Draw'],Q(1))*theta)*pref;rows=[]
    for delta in [Q(1,2),Q(1,4),Q(1,8),Q(1,16)]:
        geo=Geometry(tuple(Q(3)+delta*(j-2) for j in range(5)),r);rr=kin.rates(geo)
        for j,st in enumerate(geo.construct()['states']):
            matrix=ps.reduced_jacobian(rr,st['z']);J=mp.matrix([[mp.mpf(v.numerator)/v.denominator for v in row] for row in matrix]);e=mp.eig(J,left=False,right=False);slow=min(e,key=abs);leading=-a*delta**4*math.prod(j-k for k in range(5) if k!=j);prediction=mp.mpf(leading.numerator)/leading.denominator
            rows.append(dict(delta=str(delta),state=j,slow_eigenvalue=mp.nstr(slow,30),leading_prediction=str(leading),ratio=mp.nstr(slow/prediction,25)))
    return dict(a_exact=str(a),rows=rows,scope='Exact positive splitting coefficient and high-precision numerical eigenvalue comparison. The delta^-4 recovery tradeoff belongs to this three-site coalescing family, not a universal speed limit.')

def plot(out,profiles,recovery,curves,limiting):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    for p in profiles:axs[0].plot(float(Q(p['x'])),p['readout_micromolar'],'o',mfc='C0' if p['unstable']==0 else 'white',mec='C0',ms=8)
    if recovery['all_certified']:
        for s in recovery['sinks']:
            a,b=map(lambda x:float(Q(x)),s['readout_interval_micromolar']);x=float(Q(profiles[s['index']]['x']));axs[0].vlines(x,a,b,color='C1',lw=5)
    axs[0].set(xlabel='Auxiliary equilibrium coordinate x',ylabel='Fully phosphorylated pool / micromolar',title='Five equilibria, three stable labels');axs[0].plot([],[],'o',mfc='C0',mec='C0',label='Sink');axs[0].plot([],[],'o',mfc='white',mec='C0',label='One-direction saddle');axs[0].legend(fontsize=8)
    for j,run in curves:
        nominal=profiles[j]['readout_micromolar'];rd=(run['species'][3]+run['species'][11])*float(Q(CONCENTRATION_UNIT_MICROMOLAR));axs[1].semilogx(run['time'][1:]*float(Q(TIME_UNIT_SECONDS))/3600,(rd[1:]-nominal)/nominal,label=f'Sink {j}')
    axs[1].set(xlabel='Time / hours (declared units)',ylabel='Relative readout deviation',title='Full nonlinear recovery (numerical)');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'stable_labels_recovery.png',dpi=180);fig.savefig(out/'stable_labels_recovery.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(11,4),layout='constrained')
    ns=np.arange(1,13);axs[0].plot(ns,2*ns-1,'o-',label='Positive equilibria: 2n - 1');axs[0].plot(ns,ns,'s-',label='Stable equilibria: n');axs[0].set(xlabel='Number of sites n',ylabel='Maximum count in one class',title='Maximum equilibrium and stable-state counts');axs[0].legend(fontsize=8)
    axs[1].semilogy([r['sites'] for r in limiting],[r['feedback_deficit'] for r in limiting],'o-',label='Exact feedback deficit');axs[1].semilogy([r['sites'] for r in limiting],[r['asymptote'] for r in limiting],'--',label='Asymptotic formula');axs[1].set(xlabel='Number of sites n',ylabel='1 - feedback gain',title='Feedback stability margin versus site count');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'capacity_and_margin.png',dpi=180);fig.savefig(out/'capacity_and_margin.svg');plt.close(fig)

if __name__=='__main__':main()