"""Coupled refined reactors: useful operation, exact ledgers, and honest certificates."""
# EDITABLE INPUTS: normalized paper example, not calibrated chemistry.
COUNT_SCALE=224_000_000_000_000
MISSION_CYCLES=100
FAILURE_ALLOWANCE='1/100'
EXCHANGE=((0.,1.),(1.,0.))       # same symmetric graph for every species
RELEASE=('20','20')
CLEAVAGE=('3/100','3/100')
THETA='1/100'                    # stochastic theorem fixes this value
INITIAL=(('19/20','19/20','1/20','0','0','0','0'),
         ('13/14','13/14','0','0','1/28','0','0'))
DENSITY_CYCLES=10
FEEDBACK_POLICY=False           # True retains more after a small actual prior collection
SMALL_COUNT_SCALE=1400          # optional SSA, far below useful theorem scale
SMALL_CYCLES=2
EVENT_BUDGET=200000
RANDOM_SEED=51092026
CONCENTRATION_MOLAR='1/1000'     # illustrative unit conversion
TIME_UNIT_SECONDS=60.

import argparse
from fractions import Fraction as F
import csv
import hashlib
import json
import math
from pathlib import Path
import platform
import numpy as np
from mpmath import iv
import sympy as sp
from reactor import (Chemistry,CommonExchange,Intervention,HistoryPolicy,ReactorNetwork,
    DeterministicMission,MolecularMission,ready,A,B,I,J,Y,EPS,ETA,ACCOUNTS,SPECIES)

iv.dps=60
MANUSCRIPT_SHA256='6939f07afd2c656577771d525062a8375f3ec59a0a20e453f06c736e638d0820'


def interval(q):
    q=F(q);return iv.mpf(q.numerator)/q.denominator


def lower(x):return math.nextafter(float(x.a),-math.inf)
def upper(x):return math.nextafter(float(x.b),math.inf)


class MissionCertificate:
    """Imported theorem interface; numerical tail evaluated with outward bounds."""
    def __init__(self,network):self.network=network
    def bound(self,V,cycles,initial):
        if type(V) is not int or V<10000 or type(cycles) is not int or cycles<0:raise ValueError('Theorem requires integer V>=10000 and cycles>=0')
        net=self.network
        if not all(c.scope(True) for c in net.chemistries):raise ValueError('Rates or theta outside stochastic theorem')
        if not ready(initial,net.chemistries[0].theta,V):raise ValueError('Every initial node must be ready')
        refined=net.size==7;pref=29 if refined else 27;C=2*10**12 if refined else 2*10**14
        degree=net.exchange.exact_degree;exponent=F(V)/(C*(1+degree)**3)
        if cycles==0:prob=1.;tail=0.
        else:
            tail=upper(interval(pref*net.n*cycles)*iv.exp(-interval(exponent)))
            prob=max(0.,math.nextafter(1-tail,-math.inf))
        synthesis=net.n*cycles*(-(-V//56))+net.n*F(V,28)-sum(int(row@J[:net.size]) for row in initial)
        worst=net.n*V*(F(cycles,56)+F(1,28)-F(161,160)) if cycles else None
        return dict(lower_probability=prob,failure_union_upper=tail,exponent=exponent,
            node_count=net.n,degree_bound=degree,cycles=cycles,V=V,refined=refined,
            required_QI_per_node_cycle=-(-V//56),required_QX_per_node_cycle=-(-V//1080),
            max_each_food=5*V,max_gross_service=V//5,
            net_synthesis_lower_using_initial=synthesis if cycles else None,
            net_synthesis_lower_worst_ready=worst,
            scope='Theorem applied to ready initial counts and admitted literal history policy; not estimated from trajectories')

    def sufficient_scale(self,cycles,delta):
        if type(cycles) is not int or cycles<1 or not 0<F(delta)<1:raise ValueError('Positive mission and failure allowance in (0,1)')
        net=self.network;pref=29 if net.size==7 else 27;C=2*10**12 if net.size==7 else 2*10**14
        q=F(1)+net.exchange.exact_degree
        val=interval(C*q**3)*iv.ln(interval(F(pref*net.n*cycles)/F(delta)))
        return max(10000,math.ceil(upper(val)))


def exact_checks():
    chemistry=Chemistry();s=sp.symbols('u w x c1 c2 z h');V=sp.Symbol('V',positive=True)
    rates=[sp.Rational(c.k.numerator,c.k.denominator)*sp.prod(s[i] for i in c.inputs) for c in chemistry.channels]
    f=[sp.expand(sum(rate*c.jump[i] for rate,c in zip(rates,chemistry.channels))) for i in range(7)]
    assert sp.expand(sum(int(a)*v for a,v in zip(A,f))-(1-sum(int(a)*x for a,x in zip(A,s)) ))==0
    assert sp.expand(sum(int(a)*v for a,v in zip(B,f))-(1-sum(int(a)*x for a,x in zip(B,s)) ))==0
    count_rates=[]
    for c in chemistry.channels:
        count_rates.append(sp.Rational(c.k.numerator,c.k.denominator)*V**(1-len(c.inputs))*sp.prod(V*s[i]-c.inputs[:k].count(i) for k,i in enumerate(c.inputs)))
    correction=[sp.simplify(sum(rate*c.jump[i]/V for rate,c in zip(count_rates,chemistry.channels))-f[i]) for i in range(7)]
    assert correction==[0,0,40*s[2]/V,0,0,-20*s[2]/V,0]
    yw=[F(0),F(0),F(1),F(9,8),F(7,5),F(9,5),F(0)]
    assert sp.simplify(sum(w*c for w,c in zip(yw,correction))-4*s[2]/V)==0
    increments=[sum(int(w)*v for w,v in zip(J,c.jump)) for c in chemistry.channels if c.direction==1]
    assert increments==[1,0,0,1,0,0,-1]
    # Finite adjoint: polynomial row, rational exponential upper bound, and noise margins.
    H=sp.Matrix([[0,20,0,38],[0,5,20,0],[0,0,7,2],[0,0,20,24]])
    row=(sum((H**j*sp.Rational(1,28)**j/math.factorial(j) for j in range(5)),sp.zeros(4)))[0,:]
    expect=[F(1),F(961355,1229312),F(1847785,1843968),F(665611,307328)]
    assert list(row)==[sp.Rational(q.numerator,q.denominator) for q in expect]
    expupper=sum(F(6,7)**j/math.factorial(j) for j in range(8))+F(6,7)**8*F(9,8)/math.factorial(8)
    assert expupper<=F(2357,1000) and F(2357,1000)**2<=F(961355,172872)
    assert min(a/b for a,b in zip(expect,yw[2:6]))==F(961355,1382976)
    freeX=F(5145,100000)/8-F(8,10000)
    assert freeX==F(901,160000)
    assert F(3575,100000)>F(1,56) and freeX-F(1,1000)>F(1,1080) and F(181,1000)<F(1,5)
    # Multiplying by (1+Delta) makes the graph-noise inequality affine; coefficients suffice.
    D=sp.Symbol('D',nonnegative=True)
    eps=1/(10000*(1+D));eta=sp.Rational(11,2)*eps
    surplus=sp.expand((sp.Rational(2,3)-sp.Rational(11,20))*sp.Rational(59,5000)*(1+D)-(2*D+sp.Rational(2,3))*sp.Rational(11,20000))
    assert all(c>0 for c in sp.Poly(surplus,D).all_coeffs())
    assert F(361,5120000*10**8)>=F(1,2*10**12)
    taylor=sum(F(14)**j/math.factorial(j) for j in range(16))
    assert taylor>600000 and 1-F(5800,600000)>F(99,100)
    assert F(54,56)+F(1,28)-F(161,160)<0<F(55,56)+F(1,28)-F(161,160)
    # Same six coordinates, distinct h: exact projected derivatives and projected event rates.
    h=F(1,200);d=F(3,100);beta=theta=F(1,100)
    projection_difference=[d*h/theta,d*h/theta,d*beta*h/theta,F(0),F(0),F(0)]
    assert projection_difference[:3]==[F(3,200),F(3,200),F(3,20000)]
    return dict(material_identities=True,falling_factorial_correction=[str(v) for v in correction],
        inventory_forward_increments=increments,phase_polynomial_row=expect,exp_6_over_7_upper=expupper,
        phase_free_X_floor=freeX,noise_margin_polynomial=str(surplus),witness_exp14_lower=taylor,
        witness_probability_rational_lower=1-F(5800,600000),net_synthesis_positive_from_cycle=55,
        nonclosure_projected_derivative_difference=projection_difference,
        evidence='Fresh symbolic/rational identities; probability theorem imported, Lean not rerun')


def literal_initial(V,refined=True):
    vals=[[F(x)*V for x in row] for row in INITIAL]
    if any(v.denominator!=1 for row in vals for v in row):raise ValueError('Initial density is not integral at this V')
    return np.array([[int(v) for v in row[:7 if refined else 6]] for row in vals],dtype=np.int64)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',default='outputs');parser.add_argument('--simulate',action='store_true');args=parser.parse_args()
    out=Path(args.output);out.mkdir(parents=True,exist_ok=True)
    graph=CommonExchange(EXCHANGE);net=ReactorNetwork([Chemistry(r,d,THETA) for r,d in zip(RELEASE,CLEAVAGE)],graph)
    policy=HistoryPolicy(7,FEEDBACK_POLICY);initial=np.array([[float(F(v)) for v in row] for row in INITIAL])
    certificate=MissionCertificate(net);bound=certificate.bound(COUNT_SCALE,MISSION_CYCLES,literal_initial(COUNT_SCALE))
    run=DeterministicMission(net,policy).run(initial,DENSITY_CYCLES)
    donor=ReactorNetwork([Chemistry(r,d,refined=False) for r,d in zip(RELEASE,CLEAVAGE)],graph)
    post=np.array([Intervention(survival=('49/50',)*6).deterministic(row[:6])[0] for row in initial])
    full=donor.evolve(post);reduced=donor.donor_reduced(post)
    reduction_error=float(np.max(abs(full['end']-reduced)))
    refined_first=run['traces'][0]
    # Finite theta is retained: compare the illustrative endpoints, without convergence claims.
    theta_rows=[]
    for theta in ['1/100','1/1000','1/10000']:
        other=ReactorNetwork([Chemistry(r,d,theta) for r,d in zip(RELEASE,CLEAVAGE)],graph)
        after=np.array([Intervention().deterministic(row)[0] for row in initial]);flow=other.evolve(after)
        theta_rows.append(dict(theta=theta,max_endpoint_difference_to_donor=float(np.max(abs(flow['end'][:,:6]-full['end']))),
            carried_intermediate=flow['end'][:,6].tolist()))
    result=dict(inputs=dict(V=COUNT_SCALE,mission_cycles=MISSION_CYCLES,theta=THETA,release=RELEASE,cleavage=CLEAVAGE,exchange=EXCHANGE,initial=INITIAL),
        exact=exact_checks(),refined_certificate=bound,
        donor_certificate=MissionCertificate(donor).bound(COUNT_SCALE,MISSION_CYCLES,literal_initial(COUNT_SCALE,False)),
        sufficient_scale=certificate.sufficient_scale(MISSION_CYCLES,FAILURE_ALLOWANCE),
        donor_reduction_endpoint_error=reduction_error,theta_illustrations=theta_rows,
        numerical_mission={k:v for k,v in run.items() if k not in ('traces','final')},
        units=dict(concentration_molar=CONCENTRATION_MOLAR,cycle_seconds=4*TIME_UNIT_SECONDS,
            node_volume_liters=float(F(COUNT_SCALE)/(F('6.02214076e23')*F(CONCENTRATION_MOLAR)))))
    if args.simulate:
        result['small_count_path']=MolecularMission(net,policy,SMALL_COUNT_SCALE,RANDOM_SEED).run(literal_initial(SMALL_COUNT_SCALE),SMALL_CYCLES,EVENT_BUDGET)
    dump=lambda name,obj:(out/name).write_text(json.dumps(obj,indent=2,default=lambda v:v.tolist() if isinstance(v,np.ndarray) else str(v) if isinstance(v,F) else float(v))+'\n',encoding='utf-8')
    dump('results.json',result)
    def table(name,header,rows):
        with (out/name).open('w',newline='') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    table('cycles.csv',['cycle','node',*ACCOUNTS,'food_U','food_W',*SPECIES],
        [(h['cycle'],i,*h['accounts'][i],*h['food'][i],*h['endpoint'][i]) for h in run['history'] for i in range(graph.n)])
    table('first_cycle.csv',['t','node',*SPECIES,'material_A','material_B','stock_Y','inventory_I'],
        [(t,i,*row,row@A,row@B,row@Y,row@I) for t,cs in zip(refined_first['times'],refined_first['states']) for i,row in enumerate(cs)])
    sizes=[]
    for degree in [0,1,4]:
        other=ReactorNetwork(net.chemistries,CommonExchange([[0,degree],[degree,0]]));cert=MissionCertificate(other)
        for m in [1,3,10,30,100,300,1000,3000,10000]:sizes.append([degree,m,cert.sufficient_scale(m,'1/100')])
    table('sizing.csv',['maximum_weighted_degree','cycles','sufficient_V'],sizes)
    lines=[f'Refined mission lower probability: {bound["lower_probability"]:.10f}; donor bound at the same scale: {result["donor_certificate"]["lower_probability"]:.10f}.',
        f'Sufficient V for failure allowance {FAILURE_ALLOWANCE}: {result["sufficient_scale"]}.',
        f'Worst-ready net synthesis lower bound: {bound["net_synthesis_lower_worst_ready"]} inventory equivalents.',
        f'Donor full versus material-forced reduction endpoint gap: {reduction_error:.3g}.',
        f'Deterministic ledger residuals: inventory {run["inventory_telescope_residual"]:.3g}, intermediate storage {run["storage_telescope_residual"]:.3g}.',
        'Illustrations are deterministic; the probability theorem is imported. Lean is not rerun.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    plot(out,run,sizes)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    dump('run_metadata.json',dict(manuscript_sha256=MANUSCRIPT_SHA256,python=platform.python_version(),
        source_sha256=digest(Path(__file__)),module_sha256={p.name:digest(p) for p in Path(__file__).parent.glob('*.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,run,sizes):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    first=run['traces'][0];t=first['times'];c=first['states']
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    for i in range(c.shape[1]):
        axs[0].plot(t,(c@Y)[:,i],label=f'Node {i+1}')
        axs[1].plot(t,c[:,i,6],label=f'Node {i+1}')
    axs[0].plot(t,np.minimum(.0118*np.exp(.55*t),.052)-.00055,'--',color='k',label='Good-event stock bound')
    axs[1].plot(t,.009+.00201*np.exp(-3.02*t)+.0003,'--',color='k',label='Good-event intermediate bound')
    for ax,title,yl in zip(axs,['Catalytic stock during recovery','Intermediate concentration during recovery'],['Stock Y','Intermediate D']):
        ax.set(title=title,xlabel='Normalized time',ylabel=yl);ax.axvspan(3,4,alpha=.08);ax.grid(alpha=.2);ax.legend(fontsize=7)
    fig.savefig(out/'refined_flow.png',dpi=180);fig.savefig(out/'refined_flow.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4),layout='constrained')
    h=run['history'];k=[v['cycle'] for v in h]
    for i in range(c.shape[1]):
        axs[0].plot(k,[v['accounts'][i][0] for v in h],'-o',label=f'Inventory, node {i+1}')
        axs[0].plot(k,[v['accounts'][i][1] for v in h],'--',label=f'Free X, node {i+1}')
    axs[0].set(title='Collected output over successive reactor\ncycles',xlabel='Cycle',ylabel='Collected normalized amount');axs[0].legend(fontsize=7)
    for degree in [0,1,4]:
        a=np.array([r for r in sizes if r[0]==degree]);axs[1].loglog(a[:,1],a[:,2],label=f'Degree ≤ {degree}')
    axs[1].scatter([100],[224000000000000],color='k',s=20,label='Paper witness')
    axs[1].set(title='A sufficient count scale',xlabel='Mission cycles',ylabel='Sufficient count scale V');axs[1].legend(fontsize=7)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'mission.png',dpi=180);fig.savefig(out/'mission.svg');plt.close(fig)


if __name__=='__main__':main()
