"""Composable seven-type branching models and explicit terminal objectives."""
from dataclasses import dataclass
from fractions import Fraction as F
from typing import Protocol
import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import expm
import source


class SisterLaw(Protocol):
    def pairs(self):...


@dataclass(frozen=True)
class ComplementarySisters:
    def pairs(self):return source.PAIRS


@dataclass(frozen=True)
class IndependentSisters:
    def pairs(self):return source.INDEPENDENT


@dataclass(frozen=True)
class Action:
    erasure:F
    extra_molecular_death:F=F(0)
    def __post_init__(self):
        for name in self.__dataclass_fields__:
            object.__setattr__(self,name,F(getattr(self,name)))
            if getattr(self,name)<0:raise ValueError('Rates must be nonnegative.')


@dataclass(frozen=True)
class Phase:
    action:Action
    duration:F=F(10)
    def __post_init__(self):
        object.__setattr__(self,'duration',F(self.duration))
        if self.duration<=0:raise ValueError('Duration must be positive.')


@dataclass(frozen=True)
class Acquisition:
    scale:F=F(1,10)
    profile:tuple=tuple(source.KAPPA)
    def __post_init__(self):
        object.__setattr__(self,'scale',F(self.scale));object.__setattr__(self,'profile',tuple(map(F,self.profile)))
        if self.scale<0 or len(self.profile)!=6 or any(k<0 or not 0<=self.scale*k<=1 for k in self.profile):raise ValueError('Nonnegative scale/profile and acquisition probabilities in [0,1] required.')


@dataclass(frozen=True)
class ResistantLineage:
    birth:F=F(1,10)
    death:F=F(1,50)
    def __post_init__(self):
        object.__setattr__(self,'birth',F(self.birth));object.__setattr__(self,'death',F(self.death))
        if self.birth<0 or self.death<0:raise ValueError('Nonnegative lineage rates required.')
    @property
    def extinction(self):
        if self.birth==0:return F(1) if self.death>0 else F(0)
        return min(F(1),self.death/self.birth)


class BranchingSource:
    """Only the sister-pair law changes between the compared sources.

    Exactly one sampled daughter is replaced by M with probability mu_i.
    Custom laws must be normalized, exchangeable and share the fixed marginal.
    """
    def __init__(self,law=None,acquisition=Acquisition(),resistant=ResistantLineage()):
        self.law=law or ComplementarySisters();self.acquisition=acquisition;self.resistant=resistant
        self.pairs=self.law.pairs();self.marginal=source.L;self.mu=[acquisition.scale*k for k in acquisition.profile]
        if len(self.pairs)!=6:raise ValueError('Six conditional sister laws required.')
        for i,row in enumerate(self.pairs):
            if any(not (0<=j<6 and 0<=k<6) or p<0 for j,k,p in row) or sum(p for j,k,p in row)!=1:raise ValueError('Invalid daughter law.')
            for j in range(6):
                if sum(p for a,b,p in row if a==j)!=self.marginal[i][j] or sum(p for a,b,p in row if b==j)!=self.marginal[i][j]:raise ValueError('Daughter marginal differs; this is not the matched comparison.')
            terms={(j,k):sum(p for a,b,p in row if (a,b)==(j,k)) for j,k,p in row}
            if any(p!=terms.get((k,j),0) for (j,k),p in terms.items()):raise ValueError('The sampled pair must be exchangeable.')

    def polynomial(self,action):
        # Reuse only the literal molecular chemical generator and base deaths;
        # reconstruct the offspring polynomial for the injected law/profile.
        d,B,_=source.exact_source(action.erasure,action.extra_molecular_death,F(0),'J',self.resistant.birth,self.resistant.death)
        C=[]
        for i,row in enumerate(self.pairs):
            C.append([(j,k,F(1,10)*(1-self.mu[i])*p) for j,k,p in row]+[(j,6,F(1,10)*self.mu[i]*self.marginal[i][j]) for j in range(6) if self.marginal[i][j]])
        C.append([(6,6,self.resistant.birth)])
        return d,B,C

    def mean_matrix(self,action):
        d,B,C=self.polynomial(action);A=[row[:] for row in B]
        for i,row in enumerate(C):
            for j,k,p in row:A[i][j]+=p;A[i][k]+=p
        return A

    def numerical_field(self,action):
        d,B,C=self.polynomial(action);d=np.array(d,float);B=np.array(B,float);C=[[(j,k,float(p)) for j,k,p in row] for row in C]
        return lambda t,z:d+B@z+np.array([sum(p*z[j]*z[k] for j,k,p in row) for row in C])

    def compose(self,phases,terminal=None,killed_acquisition=False):
        z=np.array([1.]*6+[float(self.resistant.extinction)] if terminal is None else terminal,dtype=float)
        if z.shape!=(7,) or z.min()<0 or z.max()>1:raise ValueError('Terminal payoff must lie in the seven-dimensional cube.')
        for phase in reversed(phases):
            if killed_acquisition:
                # Event: no M has ever appeared. Delete success terms, retain
                # division removal and the no-acquisition daughter terms.
                d,B,C=self.polynomial(phase.action);d=np.array(d[:6],float);B=np.array(B,float)[:6,:6]
                C=[[(j,k,float(p)) for j,k,p in row if j<6 and k<6] for row in C[:6]]
                fun=lambda t,x:d+B@x+np.array([sum(p*x[j]*x[k] for j,k,p in row) for row in C])
                start=z[:6]
            else:fun=self.numerical_field(phase.action);start=z
            sol=solve_ivp(fun,(0,float(phase.duration)),start,method='DOP853',rtol=2e-12,atol=2e-14)
            if not sol.success:raise RuntimeError(sol.message)
            z=np.r_[sol.y[:,-1],0.] if killed_acquisition else sol.y[:,-1]
        return z

    def means(self,phases,initial):
        m=np.array(initial,dtype=float);records=[m.copy()]
        for phase in phases:m=m@expm(float(phase.duration)*np.array(self.mean_matrix(phase.action),float));records.append(m.copy())
        return np.array(records)

    def clearing_terminal(self):
        """Exact sufficient verification for the declared mutation-free source."""
        clear=BranchingSource(self.law,Acquisition(F(0)),self.resistant)
        A=clear.mean_matrix(Action(F(3,10)));w=list(map(F,[14,11,10,84,43,107]))
        slack=[-F(9,100)*w[i]-sum(A[i][j]*w[j] for j in range(6)) for i in range(6)]
        if min(slack)<0:raise ValueError('Molecular clearing drift not established.')
        return [F(1)]*6+[self.resistant.extinction],slack


def mechanism(phases,epsilon=F(1,10)):
    """Numerical h, second-order covariance transport and exact-response equation.

    Exact refers to the identity being integrated, not its floating solution.
    Sensitivity and response states persist across backward phase boundaries.
    """
    rho=.2;lam=np.array(source.L,float);kap=np.array(source.KAPPA,float);eps=float(epsilon);mu=eps*kap
    z=np.r_[np.ones(12),np.zeros(24)];times=[];records=[];elapsed=0
    for phase in reversed(phases):
        e,c=phase.action.erasure,phase.action.extra_molecular_death
        fj=BranchingSource(ComplementarySisters(),Acquisition(epsilon)).numerical_field(phase.action)
        fi=BranchingSource(IndependentSisters(),Acquisition(epsilon)).numerical_field(phase.action)
        zero=BranchingSource(acquisition=Acquisition(F(0)));A=np.array(zero.mean_matrix(phase.action),float)[:6,:6]
        _,Bl,_=zero.polynomial(phase.action);Bl=np.array(Bl,float)[:6,:6]
        pair=lambda P,x:np.array([sum(float(p)*x[j]*x[k] for j,k,p in row) for row in P])
        def rhs(t,z):
            xj,xi,h,delta,response,absolute=np.split(z,6)
            forcing=.1*(1-mu)*(pair(source.INDEPENDENT,xj)-pair(source.PAIRS,xj))
            B=Bl+(.1*(1-mu)*(lam@(xi+xj))+.1*mu*rho)[:,None]*lam
            return np.r_[fj(t,np.r_[xj,rho])[:6],fi(t,np.r_[xi,rho])[:6],A@h+.1*kap*(1-rho),
                A@delta+.1*(pair(source.INDEPENDENT,h)-pair(source.PAIRS,h)),B@response+forcing,B@absolute+abs(forcing)]
        sol=solve_ivp(rhs,(0,float(phase.duration)),z,method='DOP853',rtol=2e-12,atol=2e-14,t_eval=np.linspace(0,float(phase.duration),101))
        if not sol.success:raise RuntimeError(sol.message)
        times.extend((elapsed+sol.t).tolist());records.extend(sol.y.T.tolist());elapsed+=float(phase.duration);z=sol.y[:,-1]
    xj,xi,h,delta,response,absolute=np.split(z,6)
    return dict(h=h,covariance_transport=delta,finite_dependence_error=xi-xj,response_integral=response,absolute_response_bound=absolute,
        identity_residual=float(max(abs(xi-xj-response))),leading_prediction=eps**2*delta),np.array(times),np.array(records)