"""Six-species reversible binding reactor: establish, retain, and export.

Count propensities, phase monitoring, deterministic illustration and exact
algebraic budget checks remain separate. No stochastic theorem is inferred
from the deterministic curves or from a limited number of count trajectories.
"""
from __future__ import annotations

from dataclasses import dataclass
from fractions import Fraction as Q
from pathlib import Path
import argparse
import csv
import hashlib
import json
import math
import platform
import numpy as np
from scipy.integrate import solve_ivp

# EDITABLE REACTOR INPUTS. Nondimensional time uses unit washout; V is the
# bimolecular divisor and each food's feed rate, not a physical vessel volume.
COPY_SCALE = 100_000_000
EQUILIBRIUM_BIAS = Q(10)
DUPLEX_RELEASE = Q(20)
BASAL_RATE = Q(1,500_000_000)
PARAMETER_SWEEP = ((8,18),(10,20),(12,22))
ENTRY_THRESHOLD = 40_000
RESIDENCE_FLOOR = 20_000
ENTRY_DEADLINE = 500.0
END_TIME = 1000.0
EXPORT_THRESHOLD = 10_000_000
RESOURCE_LOWER,RESOURCE_UPPER = Q(9,10),Q(11,10)
ODE_RTOL,ODE_ATOL = 1e-10,1e-16
SSA_EVENT_LIMIT = 20_000       # computation limit; NEVER count unfinished as failure
SSA_SEED = 18092026
SPECIES = ('U','W','X','C1','C2','Z')
RESOURCE_A = (1,0,1,2,2,2)
RESOURCE_B = (0,1,1,1,2,2)
CATALYTIC_WEIGHTS = (Q(0),Q(0),Q(1),Q(9,8),Q(7,5),Q(9,5))
COVALENT_STOCK = (0,0,4,4,4,8)
MANUSCRIPT_SHA256 = '5f10655592fdba876cc0e8fbf37bbd3862c4a8506a4f365128ef6756c48950a2'


@dataclass(frozen=True)
class Reaction:
    name: str
    inputs: tuple[int,...]
    outputs: tuple[int,...]
    coefficient: Q
    export_mark: int=0
    catalytic_ligation: bool=False

    @property
    def jump(self):
        return tuple(self.outputs.count(i)-self.inputs.count(i) for i in range(6))

    def count_rate(self,counts,volume):
        rate=self.coefficient*Q(volume)**(1-len(self.inputs))
        seen={}
        for i in self.inputs:
            rate*=max(0,int(counts[i])-seen.get(i,0));seen[i]=seen.get(i,0)+1
        return rate


@dataclass(frozen=True)
class ReactorParameters:
    K: Q=EQUILIBRIUM_BIAS
    release: Q=DUPLEX_RELEASE
    epsilon: Q=BASAL_RATE

    def __post_init__(self):
        if min(self.K,self.release,self.epsilon) <= 0:
            raise ValueError('All kinetic parameters must be positive.')


class BindingReactor:
    def __init__(self,parameters=ReactorParameters(),enabled=True):
        self.parameters,self.enabled=parameters,enabled
        K,r,e=parameters.K,parameters.release,parameters.epsilon
        pairs=(('basal',(0,1),(2,),e,e/K,False),
               ('first_binding',(2,0),(3,),Q(20),Q(20),False),
               ('second_binding',(3,1),(4,),Q(20),Q(20),False),
               ('templated_ligation',(4,),(5,),Q(20),Q(20)/K,True),
               ('duplex_release',(5,),(2,2),r,r,False))
        channels=[]
        for name,left,right,forward,reverse,catalytic in pairs:
            if enabled or not catalytic:
                channels.extend((Reaction(name+'_forward',left,right,forward,0,catalytic),
                                 Reaction(name+'_reverse',right,left,reverse,0,catalytic)))
        channels += [Reaction('feed_'+SPECIES[i],(),(i,),Q(1)) for i in (0,1)]
        channels += [Reaction('washout_'+SPECIES[i],(i,),(),Q(1),COVALENT_STOCK[i]) for i in range(6)]
        self.channels=tuple(channels)
        self.jumps=np.array([c.jump for c in channels],dtype=np.int64)
        self.marks=np.array([c.export_mark for c in channels],dtype=np.int64)
        self.orders=np.array([len(c.inputs) for c in channels])
        self.coefficients=np.array([float(c.coefficient) for c in channels])
        self.inputs=np.full((len(channels),2),6,dtype=int)
        self.offsets=np.zeros((len(channels),2),dtype=int)
        for j,c in enumerate(channels):
            seen={}
            for k,i in enumerate(c.inputs):
                self.inputs[j,k]=i;self.offsets[j,k]=seen.get(i,0);seen[i]=seen.get(i,0)+1

    def rates(self,values,volume=None):
        factors=np.r_[values,1][self.inputs]
        if volume is not None:
            factors=np.maximum(0,factors-self.offsets)
        answer=self.coefficients*np.prod(factors,axis=1)
        return answer if volume is None else answer*np.power(float(volume),1-self.orders)

    def generator(self,counts,volume,weights):
        rates=[c.count_rate(counts,volume) for c in self.channels]
        deltas=[sum(w*d for w,d in zip(weights,c.jump)) for c in self.channels]
        return sum(a*d for a,d in zip(rates,deltas)),sum(a*d*d for a,d in zip(rates,deltas))

    def rhs(self,time,state):
        rates=self.rates(state[:6])
        return np.r_[rates@self.jumps,rates@self.marks]

    def deterministic(self,policy=None,rtol=ODE_RTOL,method='LSODA'):
        policy=policy or OperatingPolicy()
        def entry(time,state):
            return float(np.array(CATALYTIC_WEIGHTS,dtype=float)@state[:6])-policy.entry/policy.volume
        entry.direction=1
        solution=solve_ivp(self.rhs,(0,policy.end),[1,1,0,0,0,0,0],method=method,
                           rtol=rtol,atol=ODE_ATOL,dense_output=True,max_step=1,events=entry)
        if not solution.success:raise RuntimeError(solution.message)
        return solution

    def stochastic(self,policy=None,seed=SSA_SEED,limit=SSA_EVENT_LIMIT):
        policy=policy or OperatingPolicy()
        if limit < 1 or policy.volume >= 10**12:
            raise ValueError('Positive event budget and V below 10^12 required for this SSA.')
        rng=np.random.default_rng(seed)
        monitor=OperatingMonitor(policy,self.enabled)
        time,events=0.,0
        trace=[(time,*monitor.counts.tolist(),monitor.phase,monitor.counter)]
        while events < limit and monitor.status == 'active' and time < policy.end:
            rates=self.rates(monitor.counts,policy.volume)
            total=rates.sum()
            next_time=time+rng.exponential(1/total)
            boundary=policy.deadline if time < policy.deadline else policy.end
            if next_time > boundary:
                time=boundary
                monitor.advance(time)
                continue  # exponential memorylessness; no state or counter reset
            event=int(np.searchsorted(np.cumsum(rates),rng.random()*total,side='right'))
            time=next_time;events+=1
            monitor.jump(time,monitor.counts+self.jumps[event],int(self.marks[event]))
            if events % 1000 == 0:
                trace.append((time,*monitor.counts.tolist(),monitor.phase,monitor.counter))
        if time >= policy.end:monitor.advance(time)
        trace.append((time,*monitor.counts.tolist(),monitor.phase,monitor.counter))
        result=monitor.report()
        result.update(time=time,events=events,trace=trace)
        if monitor.status == 'active':
            result.update(status='unfinished',event_verdict=None)
        return result


@dataclass(frozen=True)
class OperatingPolicy:
    volume: int=COPY_SCALE
    entry: int=ENTRY_THRESHOLD
    floor: int=RESIDENCE_FLOOR
    deadline: float=ENTRY_DEADLINE
    end: float=END_TIME
    export: int=EXPORT_THRESHOLD
    lower: Q=RESOURCE_LOWER
    upper: Q=RESOURCE_UPPER

    def __post_init__(self):
        if (not isinstance(self.volume,int) or self.volume < 1 or not 0 < self.floor < self.entry
                or not 0 < self.deadline < self.end or self.export < 1 or not 0 < self.lower < 1 < self.upper):
            raise ValueError('Invalid operating thresholds or resource corridor.')


class OperatingMonitor:
    """Stateful monitor for the exact stopped law, including the deadline gate."""
    def __init__(self,policy,enabled=True):
        self.policy,self.enabled=policy,enabled
        self.counts=np.array((policy.volume,policy.volume,0,0,0,0),dtype=np.int64)
        self.phase,self.counter,self.status=0,0,'active'
        self.time,self.entry_time=0.,None

    def corridor(self,counts):
        return all(self.policy.lower*self.policy.volume <= sum(w*int(n) for w,n in zip(weights,counts))
                   <= self.policy.upper*self.policy.volume for weights in (RESOURCE_A,RESOURCE_B))

    def advance(self,time):
        if time < self.time:raise ValueError('Monitor time cannot go backwards.')
        if self.status != 'active':return
        self.time=time
        if self.enabled and time >= self.policy.deadline and self.phase == 0:
            self.status='missed_deadline'
        elif time >= self.policy.end:
            self.status='complete'

    def jump(self,time,new_counts,mark):
        if self.status != 'active':raise ValueError('Stopped law cannot jump.')
        if time < self.time or time > self.policy.end or mark < 0:
            raise ValueError('Invalid event time or export mark.')
        if time > self.policy.deadline and self.time <= self.policy.deadline:
            self.advance(self.policy.deadline)
            if self.status != 'active':return
        self.time=time
        self.counts=np.asarray(new_counts,dtype=np.int64)
        if len(self.counts) != 6 or np.min(self.counts) < 0:raise ValueError('Invalid molecule counts.')
        if time > self.policy.deadline:
            self.counter=min(self.policy.export,self.counter+mark)
        weighted40=sum(w*int(n) for w,n in zip((0,0,40,45,56,72),self.counts))
        if self.enabled:
            if self.phase == 0 and weighted40 >= 40*self.policy.entry:
                self.phase=1;self.entry_time=time
            elif self.phase == 1 and weighted40 < 40*self.policy.floor:
                self.phase=2;self.status='return_failure'
        if not self.corridor(self.counts):self.status='resource_exit'
        # Saturated export does not stop the residence/resource requirements.

    def report(self):
        verdict=None
        if self.status != 'active':
            verdict=(self.status == 'complete' and self.phase == 1 and self.counter == self.policy.export) if self.enabled else (
                self.status == 'resource_exit' or self.counter == self.policy.export)
        return {'status':self.status,'phase':self.phase,'counter':self.counter,'entry_time':self.entry_time,
                'counts':self.counts.tolist(),'enabled':self.enabled,'event_verdict':bool(verdict) if verdict is not None else None}


class ProbabilityBudget:
    """Exact rational envelopes at the fixed theorem scale and thresholds."""
    @staticmethod
    def evaluate():
        inverse_exp=lambda a,m:Q(math.factorial(m))/Q(a)**m
        # exp(1/50) <= 1/(1-1/50); a rational bound replaces floating evaluation.
        resource=4*(inverse_exp(100000,5)+3*10**14*Q(50,49)*inverse_exp(50000,5))
        entry=Q(1,12)+Q(1,600_000_000)
        residence=inverse_exp(1600,2)+3*10**14*inverse_exp(Q(799982,125),6)
        export=Q(27,16_250_000)
        disabled=Q(968,10**7)+resource
        a=Q(8047,3125)
        assert sum(a**j/Q(math.factorial(j)) for j in range(7)) > 12
        assert Q(5,2)**9 > 2000 and inverse_exp(600000,2) < Q(1,600_000_000)
        assert max(resource,residence,export) < Q(1,10000)
        assert disabled < Q(1,5000)
        return {'resource_failure':resource,'missed_entry':entry,'return_failure':residence,
                'insufficient_export':export,'enabled_success_lower':1-resource-entry-residence-export,
                'paper_coarse_success_lower':1-entry-Q(3,10000),
                'disabled_conservative_event_upper':disabled,
                'paper_coarse_disabled_upper':Q(123,625000)}

    @staticmethod
    def applies(parameters,policy):
        return (8 <= parameters.K <= 12 and 18 <= parameters.release <= 22
                and parameters.epsilon == Q(1,500_000_000) and policy == OperatingPolicy(
                    100_000_000,40000,20000,500,1000,10000000,Q(9,10),Q(11,10)))


def symbolic_bookkeeping():
    """Independent symbolic identities and rational coefficient inequalities."""
    import sympy as sp
    u,w,x,c1,c2,z,V,e,k,r=sp.symbols('u w x c1 c2 z V e k r',positive=True)
    rates=(e*u*w/V,e*k*x,20*x*u/V,20*c1,20*c1*w/V,20*c2,20*c2,20*k*z,
           r*z,r*x*(x-1)/V,V,V,u,w,x,c1,c2,z)
    model=BindingReactor()
    def gen(weights,quadratic=False,disabled=False):
        return sp.expand(sum(rate*sum(sp.Rational(weight)*jump for weight,jump in zip(weights,channel.jump))**(2 if quadratic else 1)
                             for rate,channel in zip(rates,model.channels) if not (disabled and channel.catalytic_ligation)))
    A=u+x+2*c1+2*c2+2*z;B=w+x+c1+2*c2+2*z
    assert sp.expand(gen(RESOURCE_A)-(V-A)) == 0
    assert sp.expand(gen(RESOURCE_B)-(V-B)) == 0
    for weights,stock in ((RESOURCE_A,A),(RESOURCE_B,B)):
        residual=sp.Poly(V+2*stock-gen(weights,True),u,w,x,c1,c2,z,V)
        assert all(coefficient >= 0 for coefficient in residual.coeffs())
    expected=(e*u*w/V-e*k*x+(5*u/(2*V)-1)*x+(-sp.Rational(29,8)+11*w/(2*V))*c1
              +sp.Rational(11,10)*c2+(r/5-sp.Rational(9,5)-8*k)*z-r*(x*x-x)/(5*V))
    assert sp.expand(gen(CATALYTIC_WEIGHTS)-expected) == 0
    assert sp.expand(gen(COVALENT_STOCK,disabled=True)+4*x+4*c1+4*c2+8*z-4*(e*u*w/V-e*k*x)) == 0
    # Uniform lower drift: free food >= 4V/5 and x/V <= 1/1000.
    lower=(Q(1)-Q(1,8*10**6)-Q(22,5000),Q(31,40),Q(11,10),Q(4,5))
    assert all(c >= Q(39,100)*weight for c,weight in zip(lower,CATALYTIC_WEIGHTS[2:]))
    # Exact upper QY coefficients after x(x-1)<=x^2, k<=1/8,r<=22,
    # e<=1e-6,u/V,w/V<=5/2 and x/V<=5/4. Basal source kept separate.
    qy=gen(CATALYTIC_WEIGHTS,True)+r*x/(25*V)
    qy=sp.expand(qy.subs({k:sp.Rational(1,8),r:22,e:sp.Rational(1,10**6),u:sp.Rational(5,2)*V,w:sp.Rational(5,2)*V}))
    polynomial=sp.Poly(qy,x,c1,c2,z)
    coefficients=[polynomial.coeff_monomial(x)+sp.Rational(5,4)*V*polynomial.coeff_monomial(x*x)]
    coefficients += [polynomial.coeff_monomial(s) for s in (c1,c2,z)]
    upper=(Q(481,160),Q(343,64),Q(267,40),Q(113,25))
    assert all(c <= sp.Rational(b) <= 5*sp.Rational(weight) for c,b,weight in zip(coefficients,upper,CATALYTIC_WEIGHTS[2:]))
    return {'resource_drift_identities':True,'resource_quadratic_bounds':True,'weighted_drift_identity':True,
            'weighted_drift_lower_coefficients':[str(v) for v in lower],
            'weighted_quadratic_upper_coefficients':[str(v) for v in upper],
            'disabled_covalent_balance':True,'scope':'algebraic checks; not a rerun of the Lean stochastic proof'}


def write_csv(path,headers,rows):
    with path.open('w',newline='') as stream:
        writer=csv.writer(stream);writer.writerow(headers);writer.writerows(rows)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=Path('outputs'))
    parser.add_argument('--ssa',action='store_true',help='Also attempt a budget-limited exact count path at the configured scale.')
    args=parser.parse_args();out=args.output;out.mkdir(parents=True,exist_ok=True)
    policy=OperatingPolicy()
    diagnostics=[];curves=[]
    times=np.unique(np.r_[np.linspace(0,60,601),np.linspace(60,policy.end,941),policy.deadline])
    settings=[(K,r,True) for K,r in PARAMETER_SWEEP]+[(EQUILIBRIUM_BIAS,DUPLEX_RELEASE,False)]
    for K,r,enabled in settings:
        model=BindingReactor(ReactorParameters(Q(K),Q(r)),enabled)
        solution=model.deterministic(policy)
        check=model.deterministic(policy,rtol=2e-12,method='Radau')
        values=solution.sol(times).T;fine=check.sol(times).T
        window=values[:,6]-solution.sol(policy.deadline)[6]
        name=f'K{K}_r{r}_{"enabled" if enabled else "disabled"}'
        write_csv(out/(name+'.csv'),['time',*SPECIES,'total_export_per_V'],zip(times,*values.T))
        diagnostics.append({'K':str(K),'r':str(r),'enabled':enabled,
            'numerical_entry_time':float(solution.t_events[0][0]) if len(solution.t_events[0]) else None,
            'window_export_per_V':float(values[-1,6]-solution.sol(policy.deadline)[6]),
            'weighted_count_at_end':float(values[-1,:6]@np.array(CATALYTIC_WEIGHTS,dtype=float)*policy.volume),
            'maximum_resource_error_per_V':float(max(np.max(np.abs(values[:,:6]@RESOURCE_A-1)),np.max(np.abs(values[:,:6]@RESOURCE_B-1)))),
            'independent_solver_max_difference':float(np.max(np.abs(values-fine))),
            'theorem_parameters_applicable':ProbabilityBudget.applies(model.parameters,policy)})
        curves.append((name,values,window))
    result={'deterministic':diagnostics,'exact_budgets':{k:str(v) for k,v in ProbabilityBudget.evaluate().items()},
            'symbolic_checks':symbolic_bookkeeping(),'count_trajectory':None}
    if args.ssa:
        path=BindingReactor().stochastic(policy)
        write_csv(out/'count_path.csv',['time',*SPECIES,'phase','counter'],path.pop('trace'))
        result['count_trajectory']=path
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    lines=[f'K={d["K"]}, r={d["r"]}, enabled={d["enabled"]}: numerical entry={d["numerical_entry_time"]}; window export/V={d["window_export_per_V"]:.9g}.' for d in diagnostics]
    budget=ProbabilityBudget.evaluate()
    lines += [f'Exact rational enabled lower bound: {float(budget["enabled_success_lower"]):.9f}.',
              f'Exact rational disabled conservative-event upper bound: {float(budget["disabled_conservative_event_upper"]):.9g}.',
              'The deterministic curves illustrate the mechanism; probability bounds use the manuscript generator argument.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axes=plt.subplots(1,2,figsize=(11,4.4),layout='constrained')
    for (name,values,window),diagnostic in zip(curves,diagnostics):
        label=f'K={diagnostic["K"]}, r={diagnostic["r"]}' if diagnostic['enabled'] else 'Catalytic ligation disabled'
        ys=values[:,:6]@np.array(CATALYTIC_WEIGHTS,dtype=float)*policy.volume
        axes[0].plot(times[ys>0],ys[ys>0],label=label)
        keep=(times>policy.deadline)&(window>0)
        axes[1].plot(times[keep],window[keep],label=label)
    axes[0].axhline(policy.entry,color='black',ls='--',lw=1,label='Entry threshold')
    axes[0].axhline(policy.floor,color='gray',ls=':',lw=1,label='Residence floor')
    axes[0].set(xlim=(0,60),yscale='log',ylim=(.01,1e8),xlabel='Time (washout units)',
                ylabel='Weighted catalytic count',title='Deterministic startup: free and bound forms')
    axes[1].axhline(policy.export/policy.volume,color='black',ls='--',lw=1,label='Export threshold / V')
    axes[1].set(xlim=(policy.deadline,policy.end),yscale='log',xlabel='Time (washout units)',
                ylabel='Covalent export over (500,t] / V',title='Same output window and disabled control')
    for ax in axes:ax.legend(fontsize=7);ax.grid(alpha=.2)
    fig.savefig(out/'operation.png',dpi=180);fig.savefig(out/'operation.svg');plt.close(fig)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    (out/'run_metadata.json').write_text(json.dumps({'manuscript_sha256':MANUSCRIPT_SHA256,
        'source_sha256':digest(Path(__file__)),'python':platform.python_version(),
        'output_sha256':{p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name != 'run_metadata.json'}},indent=2)+'\n')


if __name__ == '__main__':main()
