"""Literal copying operations adapted from example 44; extraction remains active.

These small direct-method diagnostics are not the enormous finite stopped kernels
used by the paper theorem. No failure mass is renormalized or discarded.
"""
from dataclasses import dataclass,replace,asdict
from fractions import Fraction as Q
from itertools import combinations
import math
import numpy as np
from resident import Cell,ResidentChemistry
EXCHANGE=(None,None,0,0,None,None,2,2,3,3,1,1,4,5)

@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 ResidentChemistry();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<=14:raise ValueError('Unknown cell or event label.')
        cell=pop.cells[index];Qnext=pop.precursor;divisions=pop.divisions
        if label<14:children=(self.chemistry.channels[label].apply(cell),)
        elif label==14:
            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]*6;steps=0;status='event_budget';counter=0;extracted=formation=growth=0
        if duration==0:return pop,dict(status='duration',time=0.,events=0,gross_exchange=gross,service_counter=0,saturated=False,complete=True,extracted=0,signed_internal_formation=0,growth_consumed=0)
        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,15)
            reservoir=EXCHANGE[label] if label<14 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 label==13:extracted+=1
            elif label==14:growth+=1
            else:formation+=self.chemistry.channels[label].outputs[2]-self.chemistry.channels[label].inputs[2]
            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'),extracted=extracted,signed_internal_formation=formation,growth_consumed=growth)

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_z=sum(c.counts[2] for c in discarded),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.')