"""Family extinction and a genuinely pruned retained-branch experiment."""
from dataclasses import dataclass
from fractions import Fraction as F
import numpy as np
from scipy.integrate import solve_ivp


@dataclass(frozen=True)
class SisterKernel:
    marginal_S:F=F(1,2)
    covariance:F=F(0)
    def __post_init__(self):
        object.__setattr__(self,'marginal_S',F(self.marginal_S));object.__setattr__(self,'covariance',F(self.covariance))
        if not 0<=self.marginal_S<=1 or min(self.entries)<0:raise ValueError('Infeasible binary sister kernel.')
    @property
    def entries(self):
        p=self.marginal_S;c=self.covariance
        return (p*p+c,p*(1-p)-c,p*(1-p)-c,(1-p)**2+c)
    def generating_function(self,s,t):
        a,b,c,d=map(float,self.entries);return a*s*s+(b+c)*s*t+d*t*t
    def retained_marginal(self):
        a,b,c,d=self.entries;return (a+(b+c)/2,d+(b+c)/2)


@dataclass(frozen=True)
class Action:
    label:str
    death_S:F
    death_T:F
    def __post_init__(self):
        for k in ['death_S','death_T']:
            object.__setattr__(self,k,F(getattr(self,k)))
            if getattr(self,k)<0:raise ValueError('Death rates must be nonnegative.')


@dataclass(frozen=True)
class Phase:
    action:Action
    duration:float
    def __post_init__(self):
        if not np.isfinite(self.duration) or self.duration<=0:raise ValueError('Finite positive duration required.')


A=Action('A',F(2),F(4));B=Action('B',F(3),F(3))


class FamilySource:
    def __init__(self,founder_kernel=SisterKernel(),descendant_division=F(0),descendant_kernels=None):
        self.founder=founder_kernel;self.epsilon=F(descendant_division)
        if self.epsilon<0:raise ValueError('Nonnegative descendant division rate required.')
        self.descendants=tuple(descendant_kernels) if descendant_kernels is not None else (founder_kernel,founder_kernel)
        if len(self.descendants)!=2:raise ValueError('Supply S and T offspring kernels.')
    def field(self,action,state):
        p,s,t=state;e=float(self.epsilon)
        return np.array([self.founder.generating_function(s,t)-p,float(action.death_S)*(1-s)+e*(self.descendants[0].generating_function(s,t)-s),float(action.death_T)*(1-t)+e*(self.descendants[1].generating_function(s,t)-t)])
    def extinction(self,phases,terminal=(0,0,0)):
        x=np.asarray(terminal,float)
        if x.shape!=(3,) or np.any(x<0) or np.any(x>1):raise ValueError('Three terminal probabilities in [0,1] required.')
        for phase in reversed(tuple(phases)):
            sol=solve_ivp(lambda u,y:self.field(phase.action,y),(0,phase.duration),x,method='DOP853',rtol=2e-12,atol=2e-14)
            if not sol.success:raise RuntimeError(sol.message)
            x=sol.y[:,-1]
        return x
    def mean_operator(self,action):
        M=[[F(0)]*3 for _ in range(3)];M[0][0]=-F(1)
        for i,kernel in enumerate((self.founder,)+self.descendants):
            birth=F(1) if i==0 else self.epsilon
            M[i][1]+=2*birth*kernel.marginal_S;M[i][2]+=2*birth*(1-kernel.marginal_S)
            if i:M[i][i]-=birth+(action.death_S if i==1 else action.death_T)
        return M
    def observation_signature(self,actions):
        return dict(founder_division=F(1),descendant_division=self.epsilon,death_rates=tuple((a.death_S,a.death_T) for a in actions),
            retained_marginals=tuple(k.retained_marginal() for k in (self.founder,)+self.descendants),founder_death=F(0))
    def founder_configuration_extinction(self,phases,counts=(1,0,0)):
        if len(counts)!=3 or any(not isinstance(n,int) or n<0 for n in counts):raise ValueError('Three nonnegative integer founder counts required.')
        return float(np.prod(self.extinction(phases)**np.array(counts)))


class RetainedBranchExperiment:
    """Records all events on one branch; hidden sister states are never sampled.

    This is exact marginal simulation under the declared pruning protocol, not
    simulation of an independently branching family. It retains self-transitions.
    """
    def __init__(self,source):self.source=source
    def sample(self,phases,seed):
        rng=np.random.default_rng(seed);state=0;time=0.;records=[]
        for phase in phases:
            end=time+phase.duration
            while time<end:
                birth=1. if state==0 else float(self.source.epsilon)
                death=0. if state==0 else float(phase.action.death_S if state==1 else phase.action.death_T)
                rate=birth+death
                wait=np.inf if rate==0 else rng.exponential(1/rate)
                if time+wait>=end:time=end;break
                time+=wait
                if rng.random()<death/rate:
                    records.append(dict(time=time,event='death',state=state,action=phase.action.label));return records
                kernel=self.source.founder if state==0 else self.source.descendants[state-1]
                before=state;state=1 if rng.random()<float(kernel.retained_marginal()[0]) else 2
                records.append(dict(time=time,event='division',mother=before,retained=state,action=phase.action.label))
        records.append(dict(time=time,event='censor',state=state));return records