"""Founder preparation, independent thinning, and a reusable branching source."""
from dataclasses import dataclass
from fractions import Fraction as F
import math
import numpy as np
from scipy.sparse import lil_matrix
from scipy.sparse.linalg import expm_multiply
from scipy.stats import binom,nbinom


@dataclass(frozen=True)
class FounderPreparation:
    dependence: F = F(0)  # Probability of a shared class draw for the whole well.

    def __post_init__(self):
        if not 0<=self.dependence<=1:raise ValueError('Dependence must lie in [0,1]')

    def cdf_two(self,k):
        if k<2:return F(0)
        independent=1-F(k+5,4*2**k);shared=1-F(k+1,2**(k+1))
        return (1-self.dependence)*independent+self.dependence*shared

    def pgf(self,z,founders=2,detection=F(1)):
        if not isinstance(founders,int) or founders<0 or not 0<=z<=1 or not 0<=detection<=1:raise ValueError('Invalid source/observation inputs')
        w=1-detection+detection*z;g0=w;g1=w/(2-w)
        return (1-self.dependence)*((g0+g1)/2)**founders+self.dependence*(g0**founders+g1**founders)/2

    def detected_coefficients(self,k,detection=F(1),founders=2):
        """Exact coefficients through k, retaining the infinite geometric tail analytically."""
        d=F(detection)
        if k<0 or founders<0 or not 0<=d<=1:raise ValueError('Invalid coefficient request')
        slow=[1-d,d]+[F(0)]*max(0,k-1);slow=slow[:k+1]
        fast=[(1-d)/(1+d)]+[2*d**j/(1+d)**(j+1) for j in range(1,k+1)]
        def power(a,m):
            result=[F(1)]+[F(0)]*k
            for _ in range(m):result=[sum((result[j]*a[i-j] for j in range(i+1)),F(0)) for i in range(k+1)]
            return result
        independent=power([(a+b)/2 for a,b in zip(slow,fast)],founders)
        s,f=power(slow,founders),power(fast,founders)
        return [(1-self.dependence)*i+self.dependence*(a+b)/2 for i,a,b in zip(independent,s,f)]

    def many_cdf_numeric(self,k,founders):
        if k<founders:return 0.
        if founders==0:return 1.
        # Given j fast founders: m plus the number of geometric failures from j successes.
        shared=.5+.5*nbinom.cdf(k-founders,founders,.5)
        j=np.arange(1,founders+1)
        independent=2.**(-founders)+np.sum(binom.pmf(j,founders,.5)*nbinom.cdf(k-founders,j,.5))
        return float((1-float(self.dependence))*independent+float(self.dependence)*shared)

    def sample(self,wells,founders=2,detection=1.,seed=56092026):
        if wells<1 or founders<0 or not 0<=detection<=1:raise ValueError('Invalid sampling request')
        rng=np.random.default_rng(seed)
        shared=rng.random(wells)<float(self.dependence)
        classes=rng.random((wells,founders))<.5
        classes[shared,:]=(rng.random(np.count_nonzero(shared))<.5)[:,None]
        families=np.where(classes,rng.geometric(.5,size=classes.shape),1)
        latent=families.sum(axis=1)
        return latent,rng.binomial(latent,detection)


@dataclass(frozen=True)
class TwoTypeBranching:
    birth_s: float=0.
    death_s: float=.3
    switch_s: float=.001
    birth_r: float=.1
    death_r: float=0.
    switch_r: float=.00001

    def __post_init__(self):
        if any(not math.isfinite(v) or v<0 for v in vars(self).values()):raise ValueError('Finite nonnegative rates required')

    def killed_distribution(self,horizon,cap=40,resistant_probability=5/24):
        """Forward count law with an absorbing overflow. Returning paths are deliberately lost.

        Numeric retained probabilities give lower approximations; overflow bounds
        missing histories in exact mathematics, not numerical roundoff.
        """
        if horizon<0 or not isinstance(cap,int) or not 1<=cap<=150 or not 0<=resistant_probability<=1:raise ValueError('Invalid finite-state calculation')
        states=[(s,n-s) for n in range(cap+1) for s in range(n+1)];index={x:i for i,x in enumerate(states)}
        overflow=len(states);Q=lil_matrix((overflow+1,overflow+1),dtype=float)
        for i,(s,r) in enumerate(states):
            moves=[((s+1,r),s*self.birth_s),((s-1,r),s*self.death_s),((s-1,r+1),s*self.switch_s),
                   ((s,r+1),r*self.birth_r),((s,r-1),r*self.death_r),((s+1,r-1),r*self.switch_r)]
            for dest,rate in moves:
                if rate:Q[i,index.get(dest,overflow)]+=rate;Q[i,i]-=rate
        initial=np.zeros(overflow+1);initial[index[(1,0)]]=1-resistant_probability;initial[index[(0,1)]]=resistant_probability
        probabilities=expm_multiply(Q.tocsr().T*horizon,initial)
        counts=np.zeros(cap+1)
        for (s,r),probability in zip(states,probabilities):counts[s+r]+=probability
        return counts,float(probabilities[-1])


@dataclass(frozen=True)
class BirthDeathEnvelope:
    birth_cap: float
    death_floor: float
    horizon: float=7.

    def __post_init__(self):
        if any(not math.isfinite(v) or v<0 for v in vars(self).values()):raise ValueError('Finite nonnegative demographic bounds required')

    def cdf(self,k):
        if k<0:return 0.
        l,m,t=self.birth_cap,self.death_floor,self.horizon
        if t==0:return float(k>=1)
        if l==m:return 1-(l*t)**k/(1+l*t)**(k+1)
        if l>m:
            one_minus=-math.expm1(-(l-m)*t);den=l-m+m*one_minus
            a=m*one_minus/den;q=l*one_minus/den
        else:
            decay=math.exp(-(m-l)*t);den=m-l*decay
            a=m*(1-decay)/den;q=l*(1-decay)/den
        return 1-(1-a)*q**k

    def exact_supercritical_lower(self,k,degree=7):
        l,m,t=map(lambda v:F(str(v)),[self.birth_cap,self.death_floor,self.horizon])
        if k<1 or l<=m or k*l<(k+1)*m or degree<1 or degree%2!=1:
            raise ValueError('This Taylor/monotonicity certificate does not cover those inputs')
        x=(l-m)*t;E=sum(((-x)**j/F(math.factorial(j)) for j in range(degree+1)),F(0))
        if not 0<E<1:raise ValueError('Taylor lower bound uninformative; increase odd degree')
        a=m*(1-E)/(l-m*E);q=l*(1-E)/(l-m*E)
        return 1-(1-a)*q**k