"""Capacity is stationary; recovery and duration require the full retained state."""
from fractions import Fraction as F

# EDITABLE SCENARIO. Concentrations in micromolar, time in seconds.
SUPPLY='linear'                 # maintained, linear, saturating
SOURCE_STRENGTH=.12
INITIAL_DONOR=120.
DONOR_CONVERSION=1000.          # concentration-to-source scale, NOT a volume
DONOR_AFFINITY=.01
REPAIR_SCALE=1.
QUOTAS=(10.,4.)
DEADLINE=.004
HORIZON=1.004
PREPARATION_OFFSET=(0.,)*8      # x,z,e1,e2,zT,h,w,v, added to the supplied p
SOURCE_AMPLITUDE=0.             # maintained continuous sinusoidal perturbation
SOURCE_FREQUENCY=20.
MANUSCRIPT_SHA256='87f19d4283d6ab714b707ac855fee97713c9002e269eab4d6cb764edb6f88e34'

import argparse,csv,hashlib,json,platform
from pathlib import Path
from dataclasses import replace
from math import log
import numpy as np
from branches import nominal
from dynamics import OperatingModel,MaintainedSupply,LinearDonor,SaturatingDonor
from certificates import ModalCertificate,TubeCertificate,fixed_band_stock,R,S
import tube_kernel as kernel


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--output',type=Path,default=Path('outputs'));args=parser.parse_args();out=args.output;out.mkdir(parents=True,exist_ok=True)
    def default(v):
        if isinstance(v,F):return str(v)
        if isinstance(v,np.ndarray):return v.tolist()
        if isinstance(v,np.generic):return v.item()
        raise TypeError(type(v).__name__)
    def dump(name,value):(out/name).write_text(json.dumps(value,default=default,indent=2)+'\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)
    modal=ModalCertificate();nominal_checks=modal.check();dump('modal_recovery_checks.json',nominal_checks)
    tubes={};slices={}
    for name in ['tube_nominal','tube_Q120','tube_Q60']:
        certificate=TubeCertificate(name);tubes[name]=certificate.verify();slices[name]=certificate.slices()
        table(name+'_slices.csv',['time','HT_lower','HT_upper','source_lower','source_upper'],slices[name])
    assert tubes['tube_nominal']['first_grid_time_inside_S']=='6191921/25000000'
    dump('fresh_tube_checks.json',tubes)
    # Stationary benchmark reconstructed from exact branch responses.
    exact=nominal(True);_,L=exact.trx.response.invert(F(4));g=exact.gpx.response.enclose(L)
    threshold=[(exact.gpx.response.current_from_carrier(v)+4)/exact.source.unit_current(L) for v in g]
    assert F(11371266,10**8)<threshold[0]<threshold[1]<F(11371268,10**8)
    same=[]
    for sigma in [F(1,100),F(1,10),F(1),F(10)]:
        branch=replace(exact.trx,b=exact.trx.b*sigma,c=exact.trx.c*sigma);assert branch.response==exact.trx.response
        same.append(dict(sigma=sigma,response_unchanged=True))
    dump('stationary_benchmark.json',dict(nadph_floor=L,source_threshold=threshold,repair_family=same,scope='Companion stationary criterion, freshly evaluated. Both matched damage and repair rates scale with sigma; it is not a time rescaling of the eight-state dynamics.'))
    stocks=[fixed_band_stock(F(251,250)),fixed_band_stock(F(251,250),200,F(1,100)),fixed_band_stock(F(251,250),1500,1)]
    assert stocks[0]['linear_sufficient_initial_stock']==1927680 and all(d.get('saturating_passes',True) for d in stocks)
    dump('supply_law_certificates.json',dict(cases=stocks,storage_identity='d(Q+W)/dt = -HG-HT',scope='Fixed-band theorems apply to all of R. The 120/60 moving tubes apply to the different independent 0.1% preparation box. A saturating law is a different model, not an improved bound for the same linear law.'))
    p=np.array(modal.preparation,float);times=np.unique(np.r_[np.linspace(0,.012,121),np.linspace(.012,1.004,501)])
    model=OperatingModel();references={}
    for name,supply in [('nominal',MaintainedSupply()),('Q120',LinearDonor()),('Q60',LinearDonor(60,500)),('saturating200',SaturatingDonor())]:
        result=model.solve(p,supply,times);references[name]=result
        write_trajectory(table,name+'.csv',result)
    if len(PREPARATION_OFFSET)!=8:raise ValueError('Eight preparation offsets required.')
    initial=p+np.array(PREPARATION_OFFSET,float)
    if not 0<=DEADLINE<=HORIZON or HORIZON<=0:raise ValueError('A positive horizon after the deadline is required.')
    laws={'maintained':lambda:MaintainedSupply(SOURCE_STRENGTH,SOURCE_AMPLITUDE,SOURCE_FREQUENCY),'linear':lambda:LinearDonor(INITIAL_DONOR,DONOR_CONVERSION),'saturating':lambda:SaturatingDonor(INITIAL_DONOR,DONOR_AFFINITY,SOURCE_STRENGTH)}
    if SUPPLY not in laws:raise ValueError('Choose maintained, linear or saturating supply.')
    configured_model=OperatingModel(repair_scale=REPAIR_SCALE,quotas=QUOTAS);grid=np.unique(np.r_[np.linspace(0,HORIZON,501),DEADLINE])
    configured=configured_model.solve(initial,laws[SUPPLY](),grid);write_trajectory(table,'configured_trajectory.csv',configured)
    smallbox=np.all(np.abs(initial-p)<=.001*p+1e-15)
    applicable=REPAIR_SCALE==1 and QUOTAS==(10.,4.) and DEADLINE>=.004 and HORIZON<=1.004 and smallbox
    applicable=bool(applicable and ((SUPPLY=='maintained' and SOURCE_STRENGTH==.12 and SOURCE_AMPLITUDE==0) or (SUPPLY=='linear' and DONOR_CONVERSION==1000 and abs(INITIAL_DONOR-120)<=.001)))
    mask=grid>=DEADLINE
    dump('configured_mission.json',dict(supply=SUPPLY,initial=initial,repair_scale=REPAIR_SCALE,quotas=QUOTAS,deadline=DEADLINE,horizon=HORIZON,
        fixed_success_tube_applicable=applicable,sampled_service_minima=configured['services'][mask].min(axis=0),quota_crossings=configured['crossings'],integrated_shortfall=configured['shortfalls'][-1],
        donor_spent=configured['regeneration'][-1],storage_drawdown=configured['storage'][0]-configured['storage'][-1],storage_residual=configured['storage_identity_residual'],donor_residual=configured['donor_identity_residual'],
        scope='Numerical edited scenario. Tube applicability is checked only for the listed fixed success witnesses, not inferred from a plot. Other scenarios may satisfy separate fixed-band checks or require a new certificate.'))
    dump('numerical_reference_accounts.json',{name:dict(final_services=r['services'][-1],final_source=r['scales'][-1],regeneration=r['regeneration'][-1],integrated_service=r['integrated_services'][-1],storage_drawdown=r['storage'][0]-r['storage'][-1],shortfall=r['shortfalls'][-1],storage_residual=r['storage_identity_residual'],donor_residual=r['donor_identity_residual'],crossings=r['crossings'],evidence=r['evidence']) for name,r in references.items()})
    # Deliberately damaged preparation, separate from the small certified box.
    base=nominal();eq=base.state(base.design.equilibrium(.12));damaged=eq.copy();damaged[5]=0;damaged[6]=.95*base.trx.E;damaged[7]=0;damaged[4]=base.trx.z0
    clocks=[]
    for sigma in [1.,.1,.01]:
        m=OperatingModel(repair_scale=sigma);necessary=m.necessary_repair_delay(damaged[6]);result=m.solve(damaged,MaintainedSupply(),np.linspace(0,900/sigma,301))
        crossing=float(result['crossings'][0]) if len(result['crossings']) else None
        assert crossing is None or crossing>=necessary
        clocks.append(dict(sigma=sigma,necessary_delay=necessary,numerical_first_crossing=crossing,final_service=result['services'][-1],same_stationary_state_residual=float(max(abs(m.kinetics.rhs(0,eq,.12))))))
    dump('hidden_repair_clock.json',dict(preparation=damaged,results=clocks,scope='Analytic delay is necessary at every source strength; numerical crossings do not certify a recovery deadline. This heavily damaged state is outside the positive preparation theorem.'))
    curve=[]
    for x in np.geomspace(.01,29,100):curve.append((x,base.gpx.response.current(x),base.trx.response.current(x)))
    table('stationary_responses.csv',['NADPH','GPx_service','Trx_service'],curve)
    plot(out,slices,references,clocks,curve,threshold)
    print('All three 1206-step rational tubes replayed; maintained and 120 uM missions pass, 60 uM fails at the mission end.',flush=True)
    print(f'120 uM certificate: Trx floor {tubes["tube_Q120"]["min_LT_after_tau"]:.9f}; stored-carrier drawdown {tubes["tube_Q120"]["storage_drawdown_W0_minus_WH"]}.',flush=True)
    print(f'Configured fixed success tube applicable: {applicable}. Numerical trajectories are separate from uniform certificates.',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']},input_sha256={p.name:digest(p) for p in sorted(here.glob('*.json')) if p.name not in ['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 write_trajectory(table,name,result):
    table(name,['time','x','z','e1','e2','zT','h','w','v','donor','HG','HT','source_scale','integrated_HG','integrated_HT','regeneration','storage','shortfall_G','shortfall_T'],
        [tuple([t])+tuple(result['states'][i])+tuple([result['donor'][i]])+tuple(result['services'][i])+tuple([result['scales'][i]])+tuple(result['integrated_services'][i])+tuple([result['regeneration'][i],result['storage'][i]])+tuple(result['shortfalls'][i]) for i,t in enumerate(result['times'])])


def plot(out,slices,references,clocks,curve,threshold):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for i,label in enumerate(['GPx','Trx']):axs[0].semilogx([r[0] for r in curve],[r[i+1] for r in curve],label=label)
    axs[0].axhline(10,color='gray',ls=':');axs[0].axhline(4,color='gray',ls=':');axs[0].set(xlabel='Stationary NADPH (µM)',ylabel='Service (µM/s)',title='Steady branch turnover under matched repair\nscaling');axs[0].legend()
    axs[1].loglog([r['sigma'] for r in clocks],[r['necessary_delay'] for r in clocks],'o-',label='Necessary delay');axs[1].loglog([r['sigma'] for r in clocks],[r['numerical_first_crossing'] for r in clocks],'o--',label='Numerical first quota crossing');axs[1].set(xlabel='Matched repair multiplier σ',ylabel='Time (s)',title='Recovery time versus repair-rate multiplier');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'capacity_recovery.png',dpi=180);fig.savefig(out/'capacity_recovery.svg');plt.close(fig)
    fig,axs=plt.subplots(1,3,figsize=(12,4),layout='constrained')
    nominal=np.array(slices['tube_nominal']);axs[0].fill_between(1000*nominal[:,0],nominal[:,1],nominal[:,2],alpha=.3,label='Exact tube slices');axs[0].plot(1000*references['nominal']['times'],references['nominal']['services'][:,1],label='Numerical center');axs[0].axvline(4,color='gray',ls=':');axs[0].axhline(4,color='black',lw=.7);axs[0].set(xlim=(0,10),xlabel='Time (ms)',ylabel='Trx service (µM/s)',title='Recovery from the 0.1% box');axs[0].legend(fontsize=7)
    for key,label in [('Q120','120 µM'),('Q60','60 µM')]:
        a=np.array(slices['tube_'+key]);axs[1].fill_between(a[:,0],a[:,1],a[:,2],alpha=.22);axs[1].plot(references[key]['times'],references[key]['services'][:,1],label=label);axs[2].fill_between(a[:,0],a[:,3],a[:,4],alpha=.22);axs[2].plot(references[key]['times'],references[key]['scales'],label=label)
    axs[1].axhline(4,color='black',ls=':');axs[1].set(xlabel='Time (s)',ylabel='Trx service (µM/s)',title='Thioredoxin service at two donor\ninventories');axs[1].legend(fontsize=8)
    axs[2].axhline(float(threshold[0]),color='black',ls=':',label='Stationary threshold');axs[2].set(xlabel='Time (s)',ylabel='Source scale Q/V',title='Source depletion relative to the\nsteady-service threshold');axs[2].legend(fontsize=7)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'finite_operation.png',dpi=180);fig.savefig(out/'finite_operation.svg');plt.close(fig)


if __name__=='__main__':main()