"""Ordered polymer source and uninterrupted fed/diluted reaction trajectories.

The channel/ledger layout adapts the separately authored manuscript-55 example;
this version uses the manuscript-59 mark law and has no weak-food assumption.
"""
from dataclasses import dataclass
from collections import Counter
from itertools import product
import math
import numpy as np
from scipy.integrate import solve_ivp

FOOD=('0','1','00','01','10','11')
SELECTED=('00','11','0011')
EPS=1/500000000
MARKS=(1.,1.5,2.)


@dataclass(frozen=True)
class Channel:
    inputs: tuple
    outputs: tuple
    coefficient: float
    kind: str = 'internal'
    label: str = ''

    @property
    def change(self):
        a,b=Counter(self.inputs),Counter(self.outputs)
        return {z:b[z]-a[z] for z in a.keys()|b.keys() if b[z]!=a[z]}

    def propensity(self,counts,V):
        rate=self.coefficient*V
        for z,m in Counter(self.inputs).items():
            for j in range(m):rate*=max(0,counts[z]-j)/V
        return rate


class Catalogue:
    def __init__(self,n=4):
        if type(n)!=int or not 4<=n<=8:raise ValueError('Explicit catalogue budget: n=4..8; analytic source calculations have a separate budget.')
        self.n=n;self.words=tuple(''.join(x) for j in range(1,n+1) for x in product('01',repeat=j))
        self.splits=tuple((z[:j],z[j:],z) for z in self.words for j in range(1,len(z)))
        self.index={z:i for i,z in enumerate(self.words)}

    def productive_labels(self):
        return tuple((z,r) for r in self.splits if r[0] in FOOD and r[1] in FOOD and len(r[2])>2 for z in FOOD+(r[2],))


@dataclass(frozen=True)
class Environment:
    """Fixed assignments and marks; incidence (word, split index, H mark)."""
    S: tuple
    B: tuple
    incidences: tuple

    def validate(self,catalogue):
        R=len(catalogue.splits)
        if len(self.S)!=R or len(self.B)!=R or any(x not in MARKS for x in self.S+self.B):raise ValueError('One allowed S and B mark per split required.')
        seen=set();degrees=Counter()
        for z,r,H in self.incidences:
            if z not in catalogue.words or type(r)!=int or not 0<=r<R or H not in MARKS or (z,r) in seen:raise ValueError('Invalid or duplicate incidence.')
            seen.add((z,r));degrees[z]+=1
        if any(d>=R for d in degrees.values()):raise ValueError('The shifted capped law excludes degree R.')
        if 2*(R+len(self.incidences))+len(catalogue.words)+6>200000:raise ValueError('Explicit channel budget exceeded.')

    @classmethod
    def matched(cls,catalogue,kind='self'):
        if kind not in ['self','food','deleted']:raise ValueError('Unknown matched intervention.')
        R=len(catalogue.splits);r=catalogue.splits.index(SELECTED)
        incidence=() if kind=='deleted' else ((('0011' if kind=='self' else '0'),r,1.),)
        return cls((1.,)*R,(1.,)*R,incidence)

    def witness_class(self,catalogue):
        return not any(z in FOOD for z,r,H in self.incidences) and any(z=='0011' and catalogue.splits[r]==SELECTED for z,r,H in self.incidences)

    def deleted(self):return Environment(self.S,self.B,())


class PolymerReactor:
    def __init__(self,catalogue,environment):
        environment.validate(catalogue);self.catalogue=catalogue;self.environment=environment
        groups={r:[] for r in range(len(catalogue.splits))}
        for z,r,H in environment.incidences:groups[r].append((z,4*environment.S[r]*H))
        channels=[]
        for j,(u,v,w) in enumerate(catalogue.splits):
            for z,k in [(None,EPS*environment.S[j]*environment.B[j])]+groups[j]:
                left=(u,v)+((z,) if z else ());right=(w,)+((z,) if z else ())
                channels.extend([Channel(left,right,k,label=f'{j}:{z}:forward'),Channel(right,left,k,label=f'{j}:{z}:reverse')])
        channels.extend(Channel((),(z,),1.,'feed',z) for z in FOOD)
        channels.extend(Channel((z,),(),1.,'wash',z) for z in catalogue.words)
        self.channels=tuple(channels);X=len(catalogue.words)
        self.stoich=np.zeros((X,len(channels)),dtype=np.int64)
        self.roles=np.full((len(channels),3),X,dtype=int);self.offsets=np.zeros_like(self.roles)
        for j,ch in enumerate(channels):
            for z,d in ch.change.items():self.stoich[catalogue.index[z],j]=d
            used=Counter()
            for i,z in enumerate(ch.inputs):self.roles[j,i]=catalogue.index[z];self.offsets[j,i]=used[z];used[z]+=1
        self.coefficient=np.array([ch.coefficient for ch in channels]);self.length=np.array(list(map(len,catalogue.words)))
        self.nonfood=np.where(self.length>2,self.length,0);self.internal=np.array([ch.kind=='internal' for ch in channels])
        self.export=np.array([len(ch.inputs[0]) if ch.kind=='wash' and len(ch.inputs[0])>2 else 0 for ch in channels])
        self.arrival=np.array([ch.kind=='feed' for ch in channels],int)
        self.supply=np.array([len(ch.outputs[0]) if ch.kind=='feed' else 0 for ch in channels])
        self.synthesis=(self.nonfood@self.stoich)*self.internal

    def initial(self,V=1):return np.array([V if z in FOOD else 0 for z in self.catalogue.words])

    def density_flux(self,x):
        return self.coefficient*np.prod(np.maximum(np.r_[x,1.][self.roles],0),axis=1)

    def count_rates(self,counts,V):
        # Integer subtraction precedes normalization; repeated catalyst/reactant
        # roles use falling factorials with no additional symmetry divisor.
        reactants=np.r_[counts,V][self.roles]-self.offsets
        return self.coefficient*V*np.prod(np.maximum(reactants,0)/V,axis=1)

    def deterministic(self,times):
        times=np.asarray(times,float)
        if len(times)<2 or times[0]!=0 or np.any(np.diff(times)<=0):raise ValueError('Increasing times starting at zero required.')
        X=len(self.catalogue.words)
        def rhs(t,y):
            f=self.density_flux(y[:X]);return np.r_[self.stoich@f,self.export@f,self.arrival@f,self.supply@f,self.synthesis@f]
        solution=solve_ivp(rhs,(0,times[-1]),np.r_[self.initial(),np.zeros(4)],method='BDF',rtol=2e-10,atol=1e-23,t_eval=times)
        if not solution.success:raise ArithmeticError(solution.message)
        return solution.y.T

    def simulate(self,V=100,horizon=1.,seed=59092026,event_budget=100000):
        if type(V)!=int or not 1<=V<=10**9 or not math.isfinite(horizon) or horizon<=0 or type(event_budget)!=int or event_budget<1:raise ValueError('Invalid count-scale/horizon/event budget; direct SSA supports V<=1e9.')
        rng=np.random.default_rng(seed);counts=self.initial(V).astype(np.int64);t=0.;events=0;Q=A=B=J=0
        first_nonfood=first_selected=first_exit=None;maxmass=10*V;windows=[0,0];observations={}
        due=[x for x in [1.,100.,199.] if x<=horizon];tetramers=np.array([len(z)==4 for z in self.catalogue.words]);selected=self.catalogue.index['0011']
        while t<horizon and events<event_budget:
            rates=self.count_rates(counts,V);total=rates.sum();tn=t+rng.exponential(1/total)
            for obs in due:
                if obs not in observations and t<=obs<tn:observations[obs]=dict(mass=int(self.length@counts),largest_tetramer=int(max(counts[tetramers])),selected=int(counts[selected]))
            if tn>horizon:t=horizon;break
            mark=min(int(np.searchsorted(np.cumsum(rates),rng.random()*total)),len(rates)-1)
            counts+=self.stoich[:,mark];t=tn;events+=1
            if min(counts)<0:raise ArithmeticError('Impossible negative population.')
            Q+=int(self.export[mark]);A+=int(self.arrival[mark]);B+=int(self.supply[mark]);J+=int(self.synthesis[mark])
            if 1<t<=100:windows[0]+=int(self.export[mark])
            if 100<t<=199:windows[1]+=int(self.export[mark])
            mass=int(self.length@counts);maxmass=max(maxmass,mass)
            if first_nonfood is None and self.nonfood@counts>0:first_nonfood=t
            if first_selected is None and counts[selected]>0:first_selected=t
            if first_exit is None and mass>11*V:first_exit=t
        residual=int(self.nonfood@counts)+Q-J
        if residual:raise ArithmeticError('Signed nonfood synthesis ledger mismatch.')
        complete=t>=horizon
        mission=(complete and horizon==199 and len(observations)==3 and maxmass<=11*V and A<=1195*V and all(q*10>V for q in windows)
                 and all(o['mass']*2<=21*V and o['largest_tetramer']*3*10**18>=V for o in observations.values()))
        return dict(completed=complete,time=t,V=V,seed=seed,events=events,first_nonfood=first_nonfood,first_selected=first_selected,first_mass_exit=first_exit,
                    observations=observations,export=Q,window_exports=windows,food_arrivals=A,food_monomers=B,signed_synthesis=J,
                    final_nonfood=int(self.nonfood@counts),ledger_residual=residual,maximum_mass=maxmass,final_counts=counts.tolist(),
                    mission=mission if complete and horizon==199 else None,
                    evidence='Direct marked SSA diagnostic; incomplete runs stay incomplete; not a theorem-scale simulation')