"""Full resident chemistry, complementary division and arbitrary-cycle physical source.
Adapted from the companion two-cycle example. Analytical marks are separate.
"""
from dataclasses import dataclass,replace,asdict
from fractions import Fraction as Q
from itertools import combinations
import math
import numpy as np
import sympy as sp
from mpmath import mp
from scipy.integrate import solve_ivp
from resident import Cell,ResidentChemistry,ChemicalRegions,Interval,MATRICES,A,reconstruction,residual
ALPHA=Q(1,10**12);MIN_SIZE=14*10**19
SPECIES=('A','B','z','H');RESERVOIRS=('F','G','RA','RB','WH')
# Gross exchange for each directed resident label, including supplied and collected molecules.
EXCHANGE=(None,None,0,0,None,None,2,2,3,3,1,1,4)
def number(ctx,x):
    x=Q(str(x));return ctx.mpf(x.numerator)/x.denominator
def dot(a,b):return sum(x*y for x,y in zip(a,b))

class ChemicalModel(ResidentChemistry):
    def __init__(self):
        super().__init__();self.jumps=np.array([np.subtract(c.outputs,c.inputs) for c in self.channels]);self.coefficients=np.array([float(c.coefficient) for c in self.channels])
    def rates(self,cell):
        a,b,z,h=cell.counts;m=cell.size
        return np.array([a,b*z/m,16*z,h,h,2*z*(z-1)/m,6*m,a,27*m,b,b/100000,a*(a-1)/(100000*m),h/10000],float)
    def density(self,x):
        a,b,z,h=x
        return np.array([6-2*a+b*z+2e-5*(b-a*a),27+a-(1+z)*b-1e-5*(b-a*a),a-b*z-16*z-4*z*z+3*h,16*z+2*z*z-2.0001*h])
    def certificates(self):
        # Full species A,B,z,H,F,G,RA,RB,WH, with maintained reservoirs projected out.
        left=((0,),(2,4),(3,),(6,),(7,),(1,5));right=((1,2),(3,),(2,2),(0,),(1,),(0,0))
        element=(2,1,1,2,1,3,2,1,2);boltzmann=(Q(1),Q(1),Q(1),Q(2),Q(1,8),Q(1),Q(1,6),Q(1,27),Q(1))
        for j,(l,r) in enumerate(zip(left,right)):
            if sum(element[i] for i in l)!=sum(element[i] for i in r):raise ArithmeticError('Full material balance failed.')
            ratio=math.prod(boltzmann[i] for i in r)/math.prod(boltzmann[i] for i in l)
            if ratio!=self.channels[2*j].coefficient/self.channels[2*j+1].coefficient:raise ArithmeticError('Common chemical potentials failed.')
            for channel,reactants,products in [(self.channels[2*j],l,r),(self.channels[2*j+1],r,l)]:
                if channel.inputs!=tuple(reactants.count(i) for i in range(4)) or channel.outputs!=tuple(products.count(i) for i in range(4)):raise ArithmeticError('Reservoir projection failed.')
        if element[3]!=element[8]:raise ArithmeticError('Sink material imbalance.')
        matrices={}
        for tag,P in MATRICES.items():
            P=sp.Matrix(P)/10**6
            if any((P-sp.eye(4)/200)[:k,:k].det()<=0 or (42*sp.eye(4)-P)[:k,:k].det()<=0 for k in range(1,5)):raise ArithmeticError('Energy sandwich failed.')
            matrices[tag]=P.tolist()
        z=sp.Symbol('z');coords=reconstruction(z);a,b,zz,h=coords;eps=sp.Rational(1,100000)
        f=(6-2*a+b*zz+2*eps*(b-a*a),27+a-(1+zz)*b-eps*(b-a*a),a-b*zz-16*zz-4*zz*zz+3*h,16*zz+2*zz*zz-sp.Rational(20001,10000)*h)
        r=residual(z)
        if any(sp.simplify(x)!=0 for x in (f[0]+2*r,f[1]-r,f[2],f[3])):raise ArithmeticError('Stationary reconstruction failed.')
        return dict(element=element,boltzmann_weights=list(map(str,boltzmann)),energy_matrices=matrices,stationary_curve_identity=True,scope='Exact chemical, source projection and energy-matrix checks. The local stochastic exponential-generator inequality is imported from the paper.')
    def recovery_density(self,cell,duration=5376,method='Radau'):
        initial=np.array(cell.counts,float)/cell.size
        sol=solve_ivp(lambda t,x:self.density(x),(0,duration),initial,method=method,rtol=2e-12,atol=1e-14,dense_output=True)
        if not sol.success:raise ArithmeticError(sol.message)
        return sol

@dataclass(frozen=True)
class Population:
    cells:tuple
    precursor:int
    omega:int
    endpoint:int
    divisions:int=0
    @property
    def size(self):return sum(c.size for c in self.cells)

class ComplementaryDivision:
    def split(self,counts,rng):
        if max(counts)>=2**53:raise ValueError('Numerical division requires molecule counts below 2^53.')
        first=tuple(int(rng.binomial(n,.5)) for n in counts)
        return first,tuple(n-a for n,a in zip(counts,first))

class PopulationSource:
    """Physical source only: no type-dependent rates, reset, or analytical failure censoring."""
    def __init__(self,N,gamma,chemistry=None,division=None):
        self.N=N;self.gamma=Q(str(gamma));self.chemistry=chemistry or ChemicalModel();self.division=division or ComplementaryDivision()
        if type(N) is not int or N<1 or self.gamma<=0:raise ValueError('Positive integer N and growth coefficient required.')
    def refill(self,cells):
        cells=tuple(cells)
        if not cells or any(not self.N<=c.size<2*self.N for c in cells):raise ValueError('Cells must lie in the live-size interval.')
        W=sum(c.size for c in cells);return Population(cells,4*W,4*W,W)
    def growth_rate(self,pop,cell):return self.gamma*Q(pop.precursor,pop.omega)*cell.counts[2] if pop.precursor else Q(0)
    def step(self,pop,index,label,rng,allocation=None):
        if type(index) is not int or not 0<=index<len(pop.cells) or type(label) is not int or not 0<=label<=13:raise ValueError('Unknown cell or event label.')
        cell=pop.cells[index];Qnext=pop.precursor;divisions=pop.divisions
        if label<13:children=(self.chemistry.channels[label].apply(cell),)
        elif label==13:
            if self.growth_rate(pop,cell)<=0:raise ValueError('Disabled growth event.')
            n=list(cell.counts);n[2]-=1;grown=Cell(cell.tag,cell.size+1,tuple(n));Qnext-=1
            if grown.size==2*self.N:
                if allocation is None:first,second=self.division.split(grown.counts,rng)
                else:
                    first=tuple(allocation);second=tuple(n-a for n,a in zip(grown.counts,first))
                    if len(first)!=4 or min(first+second)<0:raise ValueError('Invalid complementary allocation.')
                children=(Cell(cell.tag,self.N,first),Cell(cell.tag,self.N,second));divisions+=1
            elif grown.size<2*self.N:children=(grown,)
            else:raise ValueError('Division threshold exceeded.')
        else:raise ValueError('Unknown event label.')
        result=replace(pop,cells=pop.cells[:index]+children+pop.cells[index+1:],precursor=Qnext,divisions=divisions)
        if result.size+result.precursor!=pop.size+pop.precursor:raise ArithmeticError('Precursor/size conservation failed.')
        return result
    def run(self,pop,duration,rng,event_budget,collect_endpoint=True,quota=None,hard_limit=None):
        if duration<0 or type(event_budget) is not int or event_budget<0:raise ValueError('Nonnegative duration/event budget required.')
        if quota is not None and (type(quota) is not int or quota<1):raise ValueError('Positive integer service quota required.')
        if hard_limit is not None and (type(hard_limit) is not int or hard_limit<0):raise ValueError('Nonnegative integer hard budget required.')
        time=0.;gross=[0]*5;steps=0;status='event_budget';counter=0
        if duration==0:return pop,dict(status='duration',time=0.,events=0,gross_exchange=gross,service_counter=0,saturated=False,complete=True)
        for _ in range(event_budget):
            if collect_endpoint and pop.precursor==pop.endpoint:status='endpoint';break
            arrays=[np.r_[self.chemistry.rates(c),float(self.growth_rate(pop,c))] for c in pop.cells];rates=np.concatenate(arrays);total=float(rates.sum())
            proposed=time+float(rng.exponential(1/total))
            if proposed>duration:time=float(duration);status='duration';break
            if proposed<=time:raise ArithmeticError('Time resolution exhausted; use a smaller diagnostic.')
            chosen=min(int(np.searchsorted(np.cumsum(rates),rng.random()*total,side='right')),len(rates)-1);index,label=divmod(chosen,14)
            reservoir=EXCHANGE[label] if label<13 else None
            if hard_limit is not None and reservoir is not None and gross[reservoir]>=hard_limit:status='hard_service_cutoff';break
            pop=self.step(pop,index,label,rng);steps+=1;time=proposed
            if reservoir is not None:gross[reservoir]+=1
            counter=counter+1 if quota is None else min(quota,counter+1)
            if collect_endpoint and pop.precursor==pop.endpoint:status='endpoint';break
        # Saturating the counter never changes a physical transition or terminates the source.
        return pop,dict(status=status,time=time,events=steps,gross_exchange=gross,service_counter=counter,saturated=quota is not None and counter==quota,complete=status in ('endpoint','duration'))

class UniformTransfer:
    def sample(self,cells,M,rng):
        cells=tuple(cells)
        if type(M) is not int or not 1<=M<=len(cells):raise ValueError('Sample size must fit the endpoint population.')
        chosen=set(map(int,rng.choice(len(cells),M,replace=False)))
        return tuple(c for i,c in enumerate(cells) if i in chosen),tuple(c for i,c in enumerate(cells) if i not in chosen)
    def enumerate(self,cells,M,tolerance):
        cells=tuple(cells);eps=Q(tolerance);n=len(cells)
        if not 0<eps<1 or not 1<=M<=n or math.comb(n,M)>100000:raise ValueError('Invalid transfer or exact enumeration exceeds 100,000 subsets.')
        means={tag:Q(M,n)*sum(c.size for c in cells if c.tag==tag) for tag in ('H','L')};values=[];bad=0
        for indices in combinations(range(n),M):
            sizes={tag:sum(cells[i].size for i in indices if cells[i].tag==tag) for tag in ('H','L')}
            failure=any(abs(sizes[tag]-means[tag])>eps*means[tag] for tag in means);bad+=failure;values.append(sizes)
        moments={}
        for tag in means:
            mean=sum(Q(row[tag]) for row in values)/len(values);variance=sum((row[tag]-mean)**2 for row in values)/len(values)
            moments[tag]=dict(mean=str(mean),variance=str(variance))
        return dict(subsets=len(values),bad=bad,failure=str(Q(bad,len(values))),means={k:str(v) for k,v in means.items()},moments=moments)
    @staticmethod
    def weighted_moments(weights,M):
        weights=tuple(map(Q,weights));n=len(weights)
        if not 1<=M<=n or any(not 0<=w<=1 for w in weights):raise ValueError('Weights in [0,1] and valid sample size required.')
        p=Q(M,n);joint=Q(M*(M-1),n*(n-1)) if n>1 else Q(0)
        mean=p*sum(weights);variance=p*(1-p)*sum(w*w for w in weights)+2*(joint-p*p)*sum(weights[i]*weights[j] for i in range(n) for j in range(i+1,n))
        if not 0<=variance<=mean:raise ArithmeticError('Without-replacement variance bound failed.')
        return dict(mean=str(mean),variance=str(variance),pair_covariance=str(joint-p*p) if n>1 else None)

class SerialProtocol:
    def __init__(self,source,M,recovery_duration=5376):
        if type(M) is not int or M<1 or recovery_duration<0:raise ValueError('Positive retained count and nonnegative recovery duration required.')
        self.source=source;self.M=M;self.recovery_duration=recovery_duration;self.transfer=UniformTransfer()
    def run(self,cells,cycles,seed,event_budget):
        if type(cycles) is not int or cycles<0 or type(event_budget) is not int or event_budget<0:raise ValueError('Nonnegative integer cycle count and event budget required.')
        rng=np.random.default_rng(seed);pop=self.source.refill(cells);history=[];used=0
        for cycle in range(cycles):
            W0=pop.size;pop,batch=self.source.run(pop,float(8/self.source.gamma),rng,event_budget-used);used+=batch['events']
            if batch['status']!='endpoint':return dict(complete=False,stage='batch',history=history,partial=batch,population=asdict(pop),events=used)
            if len(pop.cells)<self.M:return dict(complete=False,stage='insufficient_cells',history=history,events=used)
            selected,discarded=self.transfer.sample(pop.cells,self.M,rng);before=selected;residual=pop.precursor
            recovering=replace(pop,cells=selected,precursor=0)
            recovered,recovery=self.source.run(recovering,self.recovery_duration,rng,event_budget-used,collect_endpoint=False);used+=recovery['events']
            record=dict(cycle=cycle,batch=batch,recovery=recovery,start_size=W0,consumed_precursor=3*W0,residual_discarded=residual,discarded_size=sum(c.size for c in discarded),selected_size=sum(c.size for c in selected),selected=[asdict(c) for c in before],recovered=[asdict(c) for c in recovered.cells])
            history.append(record)
            if recovery['status']!='duration':return dict(complete=False,stage='recovery',history=history,population=asdict(recovered),events=used)
            pop=self.source.refill(recovered.cells)
        return dict(complete=True,history=history,population=asdict(pop),events=used,scope='Physical diagnostic completion, not a theorem-success classification; analytical energy and odds marks are not applied.')