"""Internal twenty-channel chemistry and molecular pulse, adapted from example 43.
Finite-bath coefficients are supplied by finite_bath.py, never frozen here.
"""
from dataclasses import dataclass,asdict
from fractions import Fraction as Q
from itertools import product
import math
import numpy as np
import sympy as sp
SURVIVAL=('49/50',)*6
REFILL_ERRORS=('-1/200','-1/200')
SPECIES=('U','W','X','C1','C2','Z')
A=(1,0,1,2,2,2);B=(0,1,1,1,2,2);I=(0,0,1,1,1,2);Y=(Q(0),Q(0),Q(1),Q(9,8),Q(7,5),Q(9,5))
LEDGER=('collected_I','collected_X','food_U','food_W','gross_drive','all_wash_I','net_synthesis')
V0=200000000000
def dot(a,b):return sum(x*y for x,y in zip(a,b))
def num(ctx,x):
    x=Q(str(x));return ctx.mpf(x.numerator)/x.denominator
def ceildiv(n,d):return -((-n)//d)
def counts(values):
    values=tuple(values)
    if len(values)!=6 or any(type(n) is not int or n<0 for n in values):raise ValueError('Six nonnegative Python integer counts required.')
    return values
def restart(N,V):
    N=counts(N)
    return 159*V<=160*dot(A,N)<=161*V and 159*V<=160*dot(B,N)<=161*V and 40*dot(Y,N)>=2*V

@dataclass(frozen=True)
class Channel:
    label:int
    inputs:tuple
    outputs:tuple
    coefficient:Q
    @property
    def jump(self):return tuple(self.outputs.count(i)-self.inputs.count(i) for i in range(6))
    def propensity(self,N,V):
        value=self.coefficient*Q(V)**(1-len(self.inputs));used={}
        for i in self.inputs:value*=N[i]-used.get(i,0);used[i]=used.get(i,0)+1
        return value
    def marks(self,collect):
        wash=I[self.label-12] if 12<=self.label<=17 else 0
        return (wash if collect else 0,int(collect and self.label==14),int(self.label==10),int(self.label==11),int(self.label in (18,19)),wash,dot(I,self.jump) if self.label<10 or self.label>=18 else 0)

class Reactor:
    """Literal twenty-label source; count and density propensities share stoichiometry."""
    def __init__(self,release='20',cleavage='1/50'):
        self.r=Q(release);self.d=Q(cleavage)
        if self.r<=0 or self.d<0:raise ValueError('Positive release and nonnegative drive required.')
        eps=Q(1,500000000);eta=Q(1,8000000000)
        pairs=[((0,1),(2,),eps,eps/10),((2,0),(3,),Q(20),Q(20)),((3,1),(4,),Q(20),Q(20)),((4,),(5,),Q(20),Q(2)),((5,),(2,2),self.r,self.r)]
        channels=[]
        for j,(left,right,kf,kr) in enumerate(pairs):channels.extend([Channel(2*j,left,right,kf),Channel(2*j+1,right,left,kr)])
        channels.extend(Channel(10+i,(),(i,),Q(1)) for i in (0,1))
        channels.extend(Channel(12+i,(i,),(),Q(1)) for i in range(6))
        channels.extend([Channel(18,(2,),(0,1),self.d),Channel(19,(0,1),(2,),self.d*eta)])
        self.channels=tuple(channels);self.jumps=np.array([c.jump for c in channels]);self.k=np.array([float(c.coefficient) for c in channels]);self.orders=np.array([len(c.inputs) for c in channels])
        self.inputs=np.full((20,2),6,int);self.offsets=np.zeros((20,2),int)
        for j,c in enumerate(channels):
            for k,i in enumerate(c.inputs):self.inputs[j,k]=i;self.offsets[j,k]=c.inputs[:k].count(i)
    def scope(self):return 19<=self.r<=21 and Q(1,50)<=self.d<=Q(1,25)
    def rates(self,state,V=None):
        factors=np.r_[state,1][self.inputs]
        if V is not None:factors=factors-self.offsets
        rates=self.k*np.prod(factors,axis=1)
        return rates if V is None else rates*np.power(float(V),1-self.orders)
    def symbolic_identities(self):
        N=sp.symbols('u w x c1 c2 z');V=sp.Symbol('V',positive=True)
        # Use the literal channels; independent concentration drift uses ordinary powers.
        rate=[c.coefficient*V**(1-len(c.inputs))*sp.prod(N[i]-c.inputs[:k].count(i) for k,i in enumerate(c.inputs)) for c in self.channels]
        density=[c.coefficient*sp.prod(N[i]/V for i in c.inputs) for c in self.channels]
        L=lambda weights:sp.expand(sum(a*dot(weights,c.jump) for a,c in zip(rate,self.channels)))
        corrections=[sp.expand(L(A)-(V-dot(A,N))),sp.expand(L(B)-(V-dot(B,N))),sp.expand(L(Y)-V*sum(a*dot(Y,c.jump) for a,c in zip(density,self.channels))-self.r*N[2]/(5*V))]
        if any(c!=0 for c in corrections):raise ArithmeticError('Generator identity failed.')
        return dict(material_A='V-A',material_B='V-B',stock_correction=str(self.r)+'*N_X/(5V)',chemical_inventory_marks=[c.marks(False)[6] for c in self.channels],identities_exact=True)

@dataclass(frozen=True)
class Intervention:
    q:object='1/4'
    survival:tuple=SURVIVAL
    refill:tuple=REFILL_ERRORS
    def __post_init__(self):
        object.__setattr__(self,'q',Q(self.q));object.__setattr__(self,'survival',tuple(map(Q,self.survival)));object.__setattr__(self,'refill',tuple(map(Q,self.refill)))
        if not Q(1,4)<=self.q<=Q(3,4) or len(self.survival)!=6 or any(not Q(49,50)<=l<=1 for l in self.survival) or len(self.refill)!=2 or any(abs(e)>Q(1,200) for e in self.refill):raise ValueError('Intervention outside the admitted box.')
    def probabilities(self,i):return self.q*self.survival[i],1-self.q,self.q*(1-self.survival[i])
    def doses(self,V):return tuple(math.floor(V*(1-self.q+e)) for e in self.refill)

@dataclass(frozen=True)
class PulseOutcome:
    retained:tuple
    withdrawn:tuple
    lost:tuple
    doses:tuple
    @property
    def start(self):return tuple(n+(self.doses[i] if i<2 else 0) for i,n in enumerate(self.retained))

class MolecularPulse:
    def sample(self,N,V,intervention,rng):
        N=counts(N)
        if type(V) is not int or V<=0:raise ValueError('Positive integer copy scale required.')
        if max(N)>=2**53:raise ValueError('Numerical multinomial sampler limited to counts below 2^53.')
        draws=[tuple(map(int,rng.multinomial(n,list(map(float,intervention.probabilities(i)))))) for i,n in enumerate(N)]
        return PulseOutcome(*(tuple(row[k] for row in draws) for k in range(3)),intervention.doses(V))
    def enumerate(self,N,V,intervention):
        """Exact rational product law for small pedagogical populations; no renormalization."""
        N=counts(N)
        if type(V) is not int or V<=0:raise ValueError('Positive integer copy scale required.')
        if sum(N)>12:raise ValueError('Exact enumeration is restricted to at most twelve molecules.')
        species=[]
        for i,n in enumerate(N):
            pr,pw,pl=intervention.probabilities(i);terms=[]
            for nr in range(n+1):
                for nw in range(n-nr+1):
                    nl=n-nr-nw;mass=math.comb(n,nr)*math.comb(n-nr,nw)*pr**nr*pw**nw*pl**nl
                    if mass:terms.append(((nr,nw,nl),mass))
            species.append(terms)
        for terms in product(*species):
            triples=[t[0] for t in terms];yield PulseOutcome(*(tuple(t[k] for t in triples) for k in range(3)),intervention.doses(V)),math.prod(t[1] for t in terms)