"""Literal finite-food chemistry, complementary division, and restart kernels.

Exact rational certificates and numerical finite-chain evaluations are separate.
"""
# EDITABLE INPUTS: normalized-volume hypothetical source rates, not calibration.
MODULE_CAPACITY = 9
FORWARD_RATE = '1000'
REVERSE_RATE = '1'
CROSS_CATALYSIS = '1/10'
BATCH_DEADLINE = '20'
ENCODED_WORD = (0, 1)
INITIAL_SELECTED_COUNTS = (1, 1)
RANDOM_SEED = 22092026
MAX_SSA_EVENTS = 100000
LINEAGE_CYCLES = 200
FAMILY_DEPTH = 5

import os
os.environ['OPENBLAS_NUM_THREADS']='1'
os.environ['OMP_NUM_THREADS']='1'
from dataclasses import dataclass
from fractions import Fraction as F
from functools import cached_property
from itertools import product
from math import comb
import argparse
import csv
import hashlib
import json
from pathlib import Path
import platform
import numpy as np
from scipy.linalg import expm

MANUSCRIPT_SHA256='03705999826adddef5723e0bdc3c684f4d98a55cfcc85eac11fef4c1e6b4eb00'
SPECIES=('X0','Y0','X1','Y1','F0','F1')


@dataclass(frozen=True)
class Parameters:
    capacity:int=MODULE_CAPACITY
    forward:F=F(FORWARD_RATE)
    reverse:F=F(REVERSE_RATE)
    coupling:F=F(CROSS_CATALYSIS)

    def __post_init__(self):
        if not isinstance(self.capacity,int) or not 2<=self.capacity<=40:raise ValueError('This dense reference implementation supports 2 <= K <= 40.')
        for key in ('forward','reverse','coupling'):object.__setattr__(self,key,F(getattr(self,key)))
        if self.forward<=0 or self.reverse<=0 or self.coupling<0:raise ValueError('Positive forward/reverse rates and nonnegative coupling required.')


@dataclass(frozen=True)
class MolecularState:
    counts:tuple[int,...]

    def __post_init__(self):
        if len(self.counts)!=6 or any(not isinstance(v,(int,np.integer)) or v<0 for v in self.counts):raise ValueError('Six nonnegative integer counts required.')
    @property
    def totals(self):return tuple(sum(self.counts[2*i:2*i+2])+self.counts[4+i] for i in range(2))
    @property
    def word(self):
        bits=[]
        for i in range(2):
            x,y=self.counts[2*i:2*i+2]
            if (x>0)==(y>0):return None  # empty or contaminated module
            bits.append(int(y>0))
        return tuple(bits)
    @property
    def composition(self):
        total=sum(self.counts[:4])
        if total==0:return None
        return tuple(F(v,total) for v in self.counts[:4])
    def admitted(self,K,word=None):
        actual=self.word
        return actual is not None and (word is None or actual==tuple(word)) and self.totals==(K,K) and all(1<=sum(self.counts[2*i:2*i+2])<K for i in range(2))
    @classmethod
    def pure(cls,K,word,selected):
        if len(word)!=2 or any(v not in (0,1) for v in word) or len(selected)!=2 or any(not isinstance(n,(int,np.integer)) or not 1<=n<=K for n in selected):raise ValueError('Invalid word or pure counts.')
        values=[0]*6
        for i,(bit,n) in enumerate(zip(word,selected)):values[2*i+bit]=int(n);values[4+i]=K-int(n)
        return cls(tuple(values))


def falling(n,power):
    value=1
    for i in range(power):value*=max(n-i,0)
    return value


@dataclass(frozen=True)
class Reaction:
    name:str
    inputs:tuple[int,...]
    outputs:tuple[int,...]
    rate:F

    def propensity(self,state):
        value=self.rate
        for n,p in zip(state.counts,self.inputs):value*=falling(n,p)
        return value
    def fire(self,state):
        if any(n<p for n,p in zip(state.counts,self.inputs)):raise ValueError('Insufficient reactants.')
        return MolecularState(tuple(n+b-a for n,a,b in zip(state.counts,self.inputs,self.outputs)))


class FiniteFoodChemistry:
    """All 24 directional channels; catalysts keep their input multiplicities."""
    def __init__(self,parameters=Parameters()):
        self.parameters=parameters;self.reactions=self._build()
    def _build(self):
        reactions=[];p=self.parameters
        for module in range(2):
            for bit in (0,1):
                target=2*module+bit
                for catalyst in (None,2*(1-module),2*(1-module)+1):
                    left=[0]*6;right=[0]*6;left[target]=1;left[4+module]=1;right[target]=2
                    factor=1 if catalyst is None else p.coupling
                    if catalyst is not None:left[catalyst]+=1;right[catalyst]+=1
                    name=f'{SPECIES[target]}_cat_{"none" if catalyst is None else SPECIES[catalyst]}'
                    reactions += [Reaction(name+'_birth',tuple(left),tuple(right),p.forward*factor),Reaction(name+'_reverse',tuple(right),tuple(left),p.reverse*factor)]
        return tuple(reactions)

    def simulate(self,state,deadline,rng,max_events=MAX_SSA_EVENTS):
        if state.totals!=(self.parameters.capacity,)*2 or deadline<0 or max_events<0:raise ValueError('Invalid batch totals, deadline or event budget.')
        t=0.;history=[(t,*state.counts)];events=0
        # An event cap is computational only. An unfinished run has no cycle outcome.
        while t<deadline:
            rates=np.array([float(r.propensity(state)) for r in self.reactions]);total=rates.sum()
            if total==0:
                t=float(deadline);break
            if events>=max_events:return {'completed':False,'time':t,'state':state,'events':events,'history':history}
            next_time=t+rng.exponential(1/total)
            if next_time>deadline:t=float(deadline);break
            r=self.reactions[int(rng.choice(len(rates),p=rates/total))];state=r.fire(state);t=next_time;events+=1;history.append((t,*state.counts))
        history.append((t,*state.counts));return {'completed':True,'time':t,'state':state,'events':events,'history':history}


class FairDivision:
    """Every resident AND food molecule goes to exactly one daughter."""
    def __init__(self,capacity):self.capacity=capacity
    @staticmethod
    def empty_probability(n):
        if n<1:raise ValueError('Observable defined on nonempty modules.')
        return F(1,2**(n-1))
    @classmethod
    def joint_success(cls,n):return (1-cls.empty_probability(n[0]))*(1-cls.empty_probability(n[1]))
    def refill(self,state):
        if any(t>self.capacity for t in state.totals):raise ValueError('Cannot refill an overfull module.')
        additions=tuple(self.capacity-t for t in state.totals);c=list(state.counts)
        for i,value in enumerate(additions):c[4+i]+=value
        return MolecularState(tuple(c)),additions
    def allocate(self,parent,left_counts):
        if parent.totals!=(self.capacity,)*2:raise ValueError('Parent must have conserved batch totals.')
        left=MolecularState(tuple(left_counts))
        if any(a>b for a,b in zip(left.counts,parent.counts)):raise ValueError('Allocation exceeds parent.')
        right=MolecularState(tuple(b-a for a,b in zip(left.counts,parent.counts)));a,fa=self.refill(left);b,fb=self.refill(right)
        success=parent.word is not None and a.admitted(self.capacity,parent.word) and b.admitted(self.capacity,parent.word)
        return {'success':success,'left_before':left,'right_before':right,'left':a,'right':b,'food_added_left':fa,'food_added_right':fb,'total_food_added':sum(fa)+sum(fb)}
    def sample(self,parent,rng):return self.allocate(parent,[int(rng.binomial(n,.5)) for n in parent.counts])


class ScheduledCycle:
    def __init__(self,chemistry,deadline=F(BATCH_DEADLINE)):
        if deadline<0:raise ValueError('Nonnegative deadline required.')
        self.chemistry=chemistry;self.deadline=float(deadline);self.division=FairDivision(chemistry.parameters.capacity)
    def run(self,initial,rng,max_events=MAX_SSA_EVENTS):
        if not initial.admitted(self.chemistry.parameters.capacity):raise ValueError('Cycle must start in an admitted region.')
        batch=self.chemistry.simulate(initial,self.deadline,rng,max_events)
        if not batch['completed']:return {'batch':batch,'division':None,'success':None}
        return self._finish(batch,rng)
    def _finish(self,batch,rng):
        division=self.division.sample(batch['state'],rng)
        return {'batch':batch,'division':division,'success':division['success']}


class CountChain:
    """Exact rates on one invariant pure face; shared by all four words."""
    def __init__(self,parameters=Parameters()):
        self.parameters=parameters;K=parameters.capacity
        self.states=tuple(product(range(1,K+1),repeat=2));self.index={n:i for i,n in enumerate(self.states)}
        self.admitted=tuple(n for n in self.states if max(n)<K);self.admitted_indices=np.array([self.index[n] for n in self.admitted])
    def transitions(self,state):
        K=self.parameters.capacity;p=self.parameters
        if state not in self.index:raise ValueError('Outside pure-face state space.')
        for i in (0,1):
            n,m=state[i],state[1-i];factor=1+p.coupling*m
            for change,rate in ((1,p.forward*n*(K-n)*factor),(-1,p.reverse*n*(n-1)*factor)):
                if rate:
                    target=list(state);target[i]+=change;yield tuple(target),rate
    def generator_on(self,function,state):return sum(rate*(function(target)-function(state)) for target,rate in self.transitions(state))
    @cached_property
    def matrix(self):
        L=np.zeros((len(self.states),)*2)
        for i,state in enumerate(self.states):
            for target,rate in self.transitions(state):j=self.index[target];L[i,j]+=float(rate);L[i,i]-=float(rate)
        return L
    def transition_matrix(self,time):
        if time<0:raise ValueError('Nonnegative time required.')
        P=expm(float(time)*self.matrix)
        error=float(np.max(abs(P.sum(axis=1)-1)))
        if P.min()<-1e-11 or error>2e-8:raise ArithmeticError(f'Numerical transition matrix failed diagnostic: {error:g}')
        return P # No clipping or renormalization; preserve evaluation error visibly.
    def stationary(self):
        p=self.parameters;K=p.capacity;R=p.forward/p.reverse
        one={n:F(comb(K,n))*R**n/((1+R)**K-1) for n in range(1,K+1)}
        return {state:one[state[0]]*one[state[1]] for state in self.states}
    def verify_detailed_balance(self):
        pi=self.stationary()
        if sum(pi.values())!=1:raise ArithmeticError('Stationary normalization failed.')
        for state in self.states:
            for target,rate in self.transitions(state):
                reverse=dict(self.transitions(target))[state]
                if pi[state]*rate!=pi[target]*reverse:raise ArithmeticError('Detailed balance failed.')
        return True


@dataclass(frozen=True)
class DriftCertificate:
    decay:F
    residual:F
    maximum_violation:F

    @classmethod
    def compute(cls,chain,decay):
        decay=F(decay)
        if decay<=0:raise ValueError('Positive decay rate required.')
        g=lambda n:FairDivision.empty_probability(n[0])
        residual=max(F(0),max(chain.generator_on(g,n)+decay*g(n) for n in chain.states))
        violation=max(chain.generator_on(g,n)+decay*g(n)-residual for n in chain.states)
        return cls(decay,residual,violation)
    def lower(self,time):
        if time<0 or self.maximum_violation>0:raise ValueError('Invalid time or certificate.')
        raw=1-2/(1+self.decay*F(time))-2*self.residual/self.decay
        return max(F(0),raw)
    def record(self,time):return {'decay':str(self.decay),'residual':str(self.residual),'maximum_violation':str(self.maximum_violation),'joint_return_lower':str(self.lower(time)),
        'scope':'Exact finite-state generator inequality and rational exp bound; no simulation used, Lean not rerun.'}
    @classmethod
    def operating_box(cls):
        # Lg is affine in each of a,b,gamma separately: its maximum on a box
        # occurs at a vertex, even though a*gamma and b*gamma appear.
        corners=list(product((F(500),F(1500)),(F(1,2),F(2)),(F(1,20),F(2,9))))
        certificates=[cls.compute(CountChain(Parameters(9,*p)),1800) for p in corners]
        return cls(F(1800),max(c.residual for c in certificates),F(0))


class RestartKernel:
    """Substochastic selected-daughter kernel; missing mass is joint failure.

    Family recursion retains complementary daughter counts, never samples them
    independently. Independent compartments are used only after their split.
    """
    def __init__(self,chain,deadline):
        self.chain=chain;self.P=chain.transition_matrix(deadline);self.admitted_P=self.P[chain.admitted_indices]
        B=np.zeros((len(chain.states),len(chain.admitted)));self.allocations=[]
        for i,n in enumerate(chain.states):
            choices=[]
            for y in product(range(1,n[0]),range(1,n[1])):
                weight=F(comb(n[0],y[0])*comb(n[1],y[1]),2**sum(n));j=(y[0]-1)*(chain.parameters.capacity-1)+y[1]-1
                other=(n[0]-y[0],n[1]-y[1]);k=(other[0]-1)*(chain.parameters.capacity-1)+other[1]-1
                B[i,j]=float(weight);choices.append((j,k,float(weight)))
            self.allocations.append(choices)
        self.matrix=self.admitted_P@B
        expected=self.admitted_P@np.array([float(FairDivision.joint_success(n)) for n in chain.states])
        self.row_error=float(np.max(abs(self.matrix.sum(axis=1)-expected)))
    def lineage(self,cycles):
        if not isinstance(cycles,int) or not 0<=cycles<=10000:raise ValueError('Cycles must be an integer in 0..10000.')
        values=np.ones(len(self.chain.admitted));records=[values.copy()]
        for _ in range(cycles):values=self.matrix@values;records.append(values.copy())
        return np.array(records)
    def family(self,depth):
        if not isinstance(depth,int) or not 0<=depth<=20:raise ValueError('Depth must be an integer in 0..20.')
        values=np.ones(len(self.chain.admitted));records=[values.copy()]
        for _ in range(depth):
            terminal=np.array([sum(w*values[a]*values[b] for a,b,w in choices) for choices in self.allocations])
            values=self.admitted_P@terminal;records.append(values.copy())
        return np.array(records)


class CompositionRegions:
    def __init__(self,K):
        if not isinstance(K,int) or not 2<=K<=20:raise ValueError('Explicit composition regions support 2 <= K <= 20.')
        self.K=K
        self.regions={word:tuple(MolecularState.pure(K,word,n).composition for n in product(range(1,K),repeat=2)) for word in product((0,1),repeat=2)}
    def nearest(self,measurement):
        if len(measurement)!=4 or any(not np.isfinite(float(v)) for v in measurement):raise ValueError('Four finite measurements required.')
        # Do not normalize/clamp a noisy reading: that would change its L1 error.
        p=tuple(F(str(v)) for v in measurement)
        distance={w:min(sum(abs(a-b) for a,b in zip(p,q)) for q in region) for w,region in self.regions.items()}
        minimum=min(distance.values());best=[w for w in distance if distance[w]==minimum]
        return {'word':best[0] if len(best)==1 else None,'distance':minimum,'tie':len(best)>1}
    def margin(self):return F(2,self.K),F(1,self.K)


def budget_ceiling(total):
    if total<2:return F(0)
    return FairDivision.joint_success((total//2,total-total//2))


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'));args=parser.parse_args();out=args.output;out.mkdir(parents=True,exist_ok=True)
    parameters=Parameters();chain=CountChain(parameters);chemistry=FiniteFoodChemistry(parameters);rng=np.random.default_rng(RANDOM_SEED)
    initial=MolecularState.pure(parameters.capacity,ENCODED_WORD,INITIAL_SELECTED_COUNTS)
    cycle=ScheduledCycle(chemistry,F(BATCH_DEADLINE)).run(initial,rng);batch=cycle['batch']
    np.savetxt(out/'literal_batch.csv',np.array(batch['history']),delimiter=',',header=','.join(('time',*SPECIES)),comments='')
    def serialize(value):
        if isinstance(value,MolecularState):return dict(zip(SPECIES,map(int,value.counts)))
        if isinstance(value,F):return str(value)
        raise TypeError(type(value).__name__)
    cycle_record={'completed':batch['completed'],'events':batch['events'],'terminal_time':batch['time'],'parent':batch['state'],'division':cycle['division'],'success':cycle['success']}
    (out/'sample_cycle.json').write_text(json.dumps(cycle_record,default=serialize,indent=2)+'\n')
    # Fixed manuscript reproductions stay separate from editable study parameters.
    reference=[]
    for K,s in ((9,3900),(20,9000)):
        c=CountChain(Parameters(K,F(1000),F(1),F(1,10)));cert=DriftCertificate.compute(c,s);c.verify_detailed_balance()
        P=c.transition_matrix(20);return_values=P@np.array([float(FairDivision.joint_success(n)) for n in c.states]);admitted=return_values[c.admitted_indices]
        stationary=sum(prob*FairDivision.joint_success(n) for n,prob in c.stationary().items())
        reference.append({'K':K,**cert.record(F(20)),'numerical_return_min':float(min(admitted)),'numerical_return_max':float(max(admitted)),
            'stationary_return_exact':str(stationary),'stationary_return_float':float(stationary),'partition_ceiling':str(FairDivision.joint_success((K,K))),
            'transition_row_sum_error':float(np.max(abs(P.sum(axis=1)-1)))})
    decay=3900 if parameters.capacity==9 else 9000 if parameters.capacity==20 else max(1,int(parameters.forward*(parameters.capacity-1)/3))
    cert=DriftCertificate.compute(chain,decay);kernel=RestartKernel(chain,F(BATCH_DEADLINE));lineage=kernel.lineage(LINEAGE_CYCLES);family=kernel.family(FAMILY_DEPTH)
    q=cert.lower(F(BATCH_DEADLINE));u=FairDivision.joint_success((parameters.capacity,)*2)
    def write_csv(name,header,rows):
        with (out/name).open('w',newline='') as f:w=csv.writer(f);w.writerow(header);w.writerows(rows)
    write_csv('lineage.csv',['cycles','rational_bound_float','numerical_min','numerical_max','partition_upper_float'],[[n,float(q**n),float(min(v)),float(max(v)),float(u**n)] for n,v in enumerate(lineage)])
    write_csv('family.csv',['depth','internal_cycles','union_lower_float','numerical_min','numerical_max','partition_upper_float'],[[n,2**n-1,float(max(F(0),1-(2**n-1)*(1-q))),float(min(v)),float(max(v)),float(u**(2**n-1))] for n,v in enumerate(family)])
    np.savetxt(out/'restart_kernel.csv',kernel.matrix,delimiter=',',header=','.join(f'A_{a}_{b}' for a,b in chain.admitted),comments='')
    times=np.r_[0,np.geomspace(1e-6,.02,70),.1,1.,20.];H=np.array([float(FairDivision.empty_probability(n[0])+FairDivision.empty_probability(n[1])) for n in chain.states]);observable=[]
    for time in times:
        value=chain.transition_matrix(time)@H
        envelope=2*np.exp(-float(cert.decay)*time)+float(2*cert.residual/cert.decay)*(1-np.exp(-float(cert.decay)*time))
        observable.append([time,value[chain.index[(1,1)]],value[chain.index[(parameters.capacity,)*2]],envelope])
    write_csv('empty_daughter.csv',['time','E_H_from_1_1','E_H_from_K_K','certified_envelope_numerically_evaluated'],observable)
    result={'reference_reproductions':reference,'operating_box':DriftCertificate.operating_box().record(F(20)),
        'configured_model':{'capacity':parameters.capacity,'forward':str(parameters.forward),'reverse':str(parameters.reverse),'coupling':str(parameters.coupling),**cert.record(F(BATCH_DEADLINE))},
        'best_17_resident_partition':str(budget_ceiling(17)),'sharp_composition_gap':str(F(2,parameters.capacity)),'measurement_tolerance_strict':str(F(1,parameters.capacity)),
        'kernel_row_identity_error':kernel.row_error,'sample_cycle_completed':batch['completed'],'sample_cycle_success':cycle['success'],
        'lineage_final_range':[float(min(lineage[-1])),float(max(lineage[-1]))],'family_final_range':[float(min(family[-1])),float(max(family[-1]))],
        'scope':'Finite-chain matrix exponentials and restart recursions are numerical evaluations, not certified probabilities. Exact rational generator checks are separate. The four labels use disjoint supports, not bistability or contamination correction.'}
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    lines=[f'K=9 exact lower bound: {reference[0]["joint_return_lower"]}; numerical one-cycle minimum: {reference[0]["numerical_return_min"]:.10g}.',
        f'K=20 exact lower bound: {reference[1]["joint_return_lower"]}; numerical one-cycle minimum: {reference[1]["numerical_return_min"]:.10g}.',
        f'17-resident partition ceiling: {budget_ceiling(17)} = {float(budget_ceiling(17)):.10g}, below 99%.',
        f'Configured literal cycle: completed={batch["completed"]}, success={cycle["success"]}, events={batch["events"]}.',
        f'Selected lineage after {LINEAGE_CYCLES} cycles: numerical min {min(lineage[-1]):.8g}; both daughters inspected at every split.',
        'Failed splits remain missing kernel mass. Siblings are complementary; no independent daughter sampling.']
    (out/'console.txt').write_text('\n'.join(lines)+'\n');print('\n'.join(lines))
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4.2),layout='constrained')
    short=chemistry.simulate(initial,.003,np.random.default_rng(RANDOM_SEED+1));ar=np.array(short['history'])
    for i in (0,1):axs[0].step(ar[:,0]*1000,ar[:,1+2*i]+ar[:,2+2*i],where='post',label=f'Module {i}')
    axs[0].set(xlabel='Time (thousandths of source unit)',ylabel='Selected resident molecules',title='A literal 24-channel sample path');axs[0].legend(fontsize=8)
    ar=np.array(observable)
    for column,label in [(1,'Numerical E H, start (1,1)'),(2,'Numerical E H, start (K,K)'),(3,'Proved envelope (evaluated)')]:axs[1].loglog(ar[1:,0],ar[1:,column],label=label)
    axs[1].set(xlabel='Time (source units)',ylabel='Sum of module empty-daughter probabilities',title='Amplification reduces partition risk');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'amplification.png',dpi=180);fig.savefig(out/'amplification.svg');plt.close(fig)
    fig,axs=plt.subplots(1,2,figsize=(10,4.2),layout='constrained');budgets=range(4,53)
    axs[0].semilogy(budgets,[float(1-budget_ceiling(k)) for k in budgets],label='Fair-partition failure floor')
    for row in reference:
        K=row['K'];axs[0].scatter([2*K],[float(1-F(row['joint_return_lower']))],label=f'Proved failure bound, K={K}')
    axs[0].axhline(.01,color='gray',ls=':');axs[0].set(xlabel='Peak resident capacity',ylabel='One-cycle failure probability',title='18 residents suffice in this architecture');axs[0].legend(fontsize=8)
    ns=np.arange(len(lineage));axs[1].plot(ns,[float(q**int(n)) for n in ns],label='Uniform lower bound');axs[1].plot(ns,lineage.min(axis=1),label='Numerical restart-kernel minimum');axs[1].plot(ns,[float(u**int(n)) for n in ns],ls='--',label='Partition upper bound')
    axs[1].set(xlabel='Selected-lineage cycles',ylabel='All-cycle success probability',title='Probability of success through repeated\ncopying cycles');axs[1].legend(fontsize=8)
    for ax in axs:ax.grid(alpha=.2)
    fig.savefig(out/'budget_lineage.png',dpi=180);fig.savefig(out/'budget_lineage.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()