"""Composable finite-type branching model; floating-point exploration, not certification."""
from dataclasses import dataclass
from typing import Protocol
import math
import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import solve
from model import source, pair


class Hazard(Protocol):
    def values(self, states, sites): ...


@dataclass(frozen=True)
class StepHazard:
    protected: float = .01
    exposed: float = .3

    def values(self, states, sites):
        return np.array([self.protected if a > r else self.exposed for a, r in states])


@dataclass(frozen=True)
class SmoothHazard:
    slope: float = 8.
    protected: float = .01
    exposed: float = .3

    def values(self, states, sites):
        from scipy.special import expit
        if not np.isfinite(self.slope) or self.slope < 0:
            raise ValueError('slope must be finite and nonnegative')
        return np.array([self.protected + (self.exposed-self.protected)*expit(-self.slope*(a-r)/sites) for a,r in states])


@dataclass
class MolecularSource:
    sites: int = 2
    writing: float = .01
    recruitment: float = 1.
    erasure: float = .01
    antagonism: float = .5

    def __post_init__(self):
        if type(self.sites) is not int or not 1 <= self.sites <= 32:
            raise ValueError('dense example supports integer site counts 1..32')
        if any(not np.isfinite(x) or x < 0 for x in (self.writing,self.recruitment,self.erasure,self.antagonism)):
            raise ValueError('molecular rates must be finite and nonnegative')
        self.states,self.index,self.Q,self.D,self.pairs = source(self.sites,e=self.erasure,w=self.writing,k=self.recruitment,h=self.antagonism,dense=True)


class SisterLaw(Protocol):
    def product(self, source, f, g): ...


class Complementary:
    def product(self, source, f, g):
        return pair(source.pairs,f,g,len(source.states))


class Independent:
    def product(self, source, f, g):
        return (source.D@f)*(source.D@g)


@dataclass
class BranchingPopulation:
    source: MolecularSource
    hazard: Hazard
    sisters: SisterLaw
    division_rate: float = .1
    phases: int = 1

    def __post_init__(self):
        if not np.isfinite(self.division_rate) or self.division_rate <= 0:
            raise ValueError('division rate must be positive')
        if type(self.phases) is not int or not 1 <= self.phases <= 16:
            raise ValueError('Erlang phases must be an integer in 1..16')
        self.d=np.asarray(self.hazard.values(self.source.states,self.source.sites),float)
        if self.d.shape != (len(self.source.states),) or np.any(~np.isfinite(self.d)) or np.any(self.d<0):
            raise ValueError('hazard must provide one finite nonnegative death rate per state')

    def extinction(self, tolerance=2e-13, max_iterations=100000):
        """Numerical least fixed point; stopping residual is not a rigorous error bound."""
        n=len(self.d); rate=self.phases*self.division_rate
        V=solve(np.diag(self.d+rate)-self.source.Q,rate*np.eye(n))
        W=np.linalg.matrix_power(V,self.phases)
        q=np.zeros(n)
        for it in range(max_iterations):
            nxt=1-W@np.ones(n)+W@self.sisters.product(self.source,q,q)
            if np.max(abs(nxt-q))<tolerance:
                return nxt, it+1
            q=nxt
        raise RuntimeError('extinction iteration budget exhausted; no result certified')

    def pgf(self, times, terminal=0.):
        times=self._times(times); n=len(self.d); m=self.phases; rate=m*self.division_rate
        v=np.broadcast_to(np.asarray(terminal,float),(n,)).copy()
        if np.any(~np.isfinite(v)) or np.any(v<0) or np.any(v>1):
            raise ValueError('terminal PGF arguments must lie in [0,1]')
        def rhs(t,y):
            z=y.reshape(m,n)
            next_phase=np.vstack((z[1:],self.sisters.product(self.source,z[0],z[0])))
            return (z@self.source.Q.T+self.d*(1-z)+rate*(next_phase-z)).ravel()
        return self._integrate(rhs,np.tile(v,m),times)[:n].T

    def moments(self,times):
        """First and second factorial moments, including molecular evolution in every phase."""
        times=self._times(times); n=len(self.d); m=self.phases; rate=m*self.division_rate
        def rhs(t,y):
            mean,fact=y.reshape(2,m,n)
            next_m=np.vstack((mean[1:],2*self.source.D@mean[0]))
            next_f=np.vstack((fact[1:],2*self.source.D@fact[0]+2*self.sisters.product(self.source,mean[0],mean[0])))
            return np.stack((mean@self.source.Q.T-self.d*mean+rate*(next_m-mean),fact@self.source.Q.T-self.d*fact+rate*(next_f-fact))).ravel()
        z=self._integrate(rhs,np.r_[np.ones(m*n),np.zeros(m*n)],times)
        mean=z[:n].T; fact=z[m*n:m*n+n].T
        return mean,fact+mean-mean**2

    @staticmethod
    def _times(times):
        times=np.asarray(times,float)
        if times.ndim!=1 or not len(times) or np.any(~np.isfinite(times)) or times[0]<0 or np.any(np.diff(times)<=0):
            raise ValueError('times must be nonnegative, finite and strictly increasing')
        return times

    @staticmethod
    def _integrate(rhs,initial,times):
        if times[-1]==0:return initial[:,None]
        sol=solve_ivp(rhs,(0,times[-1]),initial,t_eval=times,method='DOP853',rtol=2e-11,atol=2e-13)
        if not sol.success:raise RuntimeError(sol.message)
        return sol.y


@dataclass(frozen=True)
class IndependentCapture:
    probability: float

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

    def nondetection(self,population,times):
        return population.pgf(times,1-self.probability)

    def moments(self,mean,variance):
        p=self.probability
        return p*mean,p*p*variance+p*(1-p)*mean


@dataclass(frozen=True)
class PairedCloneObservation:
    mother_mean_range: float = 0.
    misclassification: float = 0.
    missing_fraction: float = 0.

    def interval(self,covariance,pairs=None,alpha=.05):
        if any(not 0<=x<=1 for x in (self.mother_mean_range,self.misclassification,self.missing_fraction)):
            raise ValueError('sensitivity inputs must lie in [0,1]')
        sampling=0.
        if pairs is not None:
            if type(pairs) is not int or pairs<1 or not 0<alpha<1:raise ValueError('invalid independent-pair sample size or alpha')
            sampling=3*math.sqrt(math.log(6/alpha)/(2*pairs))
        error=4*self.misclassification+3*self.missing_fraction+sampling
        return covariance-self.mother_mean_range**2/4-error,covariance+error


def balanced_survival_bound(log_sites,time,slope=8.,contrast_scaled=0.):
    """Theorem 12 numerical evaluation for the paper's symmetric rates, smooth hazard only.

    contrast_scaled = sqrt(N)*abs((a-r)/N); log_sites avoids astronomic integers.
    Values greater than one are correctly reported as uninformative bounds.
    """
    if min(log_sites,time,slope,contrast_scaled)<0 or not all(math.isfinite(x) for x in (log_sites,time,slope,contrast_scaled)):
        raise ValueError('arguments must be finite and nonnegative')
    theta=.055; psi=.145; growth=.915; noise=1.07; lip=.29*slope/4
    log_second=(-theta+psi+growth)*time-log_sites/2
    second=0. if lip*time==0 else math.exp(log_second+math.log(lip*time)+.5*math.log(contrast_scaled**2+noise*time))
    return math.exp(-theta*time)+second