"""Latched amplification source, exact finite enclosures, and identity observation.

The uniformization construction follows Appendix A/check_paper.py of manuscript 57;
it is parameterized here and separated from the observation/decision contracts.
"""
from dataclasses import dataclass
from fractions import Fraction as F
from math import factorial, ceil
import numpy as np
from scipy.linalg import expm
from scipy.stats import poisson
from scipy.optimize import brentq


def rational(x):
    return x if isinstance(x,F) else F(str(x))


def exp_negative(x, degree=200):
    x=rational(x)
    if not 0<=x<degree+2: raise ValueError('Increase exponential series budget.')
    term=total=F(1)
    for k in range(1,degree+1):
        term*=x/k;total+=term
    tail=term*x/(degree+1)/(1-x/F(degree+2))
    return 1/(total+tail),1/total


@dataclass(frozen=True)
class ErrorBounds:
    blank: tuple
    miss: tuple
    empty: tuple

    def deadline_status(self, alpha=F(1,100), beta=F(1,20)):
        if self.blank[1]<=alpha and self.miss[1]<=beta:return 'usable'
        if self.blank[0]>alpha or self.miss[0]>beta:return 'excluded deadline'
        return 'unresolved'

    def separates_all_timing(self, alpha=F(1,100), beta=F(1,20)):
        # Uses the manuscript's deadline-optimality theorem, not a grid search.
        return self.blank[0]>alpha and self.miss[0]>beta


@dataclass(frozen=True)
class AmplificationSource:
    threshold: int = 5
    capacity: F = F(5)
    background: F = F(1,100)
    growth: F = F(1)
    loading: F = F(4)

    def __post_init__(self):
        for name in ['capacity','background','growth','loading']:
            object.__setattr__(self,name,rational(getattr(self,name)))
        if type(self.threshold)!=int or not 1<=self.threshold<=200:raise ValueError('Finite-chain threshold budget: 1..200.')
        if self.capacity<=self.threshold-1 or self.background<=0 or self.growth<=0 or self.loading<0:
            raise ValueError('Require positive transient rates and nonnegative loading.')

    def rates(self):
        return tuple((self.background+self.growth*z)*(1-F(z)/self.capacity) for z in range(self.threshold))

    def numerical_errors(self,t):
        if not np.isfinite(t) or t<0:raise ValueError('Nonnegative finite time required.')
        q=np.array(list(map(float,self.rates())))
        Q=np.diag(-q)+np.diag(q[:-1],1)
        survival=expm(t*Q)@np.ones(self.threshold)
        return float(1-survival[0]),float(poisson.pmf(np.arange(self.threshold),float(self.loading))@survival)

    def initial_hit(self):
        return float(poisson.sf(self.threshold-1,float(self.loading)))

    def numerical_optimum(self,alpha=.01):
        if not 0<alpha<1:raise ValueError('Interior blank limit required.')
        hi=1.
        while self.numerical_errors(hi)[0]<alpha:
            hi*=2
            if hi>1e12:raise ArithmeticError('Root-bracketing budget exceeded.')
        t=brentq(lambda t:self.numerical_errors(t)[0]-alpha,0,hi,xtol=1e-12)
        return dict(deadline=t,blank=alpha,miss=self.numerical_errors(t)[1],evidence='numerical; deadline optimality is a manuscript theorem')

    def certify(self,t,degree=200):
        """Exact rational Poisson uniformization, including an explicit tail.

        Loading at/above threshold is a t=0 hit; transient weights are not
        normalized. The method also supports repeated transition rates.
        """
        t=rational(t);q=self.rates();h=self.threshold
        if t<0 or degree<1 or degree>1000:raise ValueError('Invalid time or degree budget.')
        nu=F(ceil(max(q)));mu=nu*t
        elo,ehi=exp_negative(mu,degree)
        if h*degree>12000:raise ValueError('Exact state/degree budget exceeded.')
        s=[F(1)]*h+[F(0)];sums=[F(0)]*(h+1);weight=F(1)
        for k in range(degree+1):
            sums=[a+weight*b for a,b in zip(sums,s)]
            s=[(1-q[z]/nu)*s[z]+q[z]/nu*s[z+1] for z in range(h)]+[F(0)]
            weight*=mu/(k+1)
        tail=ehi*weight/(1-mu/F(degree+2))
        lo=[elo*x for x in sums];hi=[min(F(1),ehi*x+tail) for x in sums]
        p0=exp_negative(self.loading,degree)
        weights=[self.loading**n/factorial(n) for n in range(h)]
        return ErrorBounds((1-hi[0],1-lo[0]),
                           (p0[0]*sum(w*v for w,v in zip(weights,lo)),min(F(1),p0[1]*sum(w*v for w,v in zip(weights,hi)))),p0)

    def sample_hit_times(self,n,rng,clock=1.,loaded=True):
        """Shared constant clock per well; separate exponential innovations."""
        if type(n)!=int or n<1:raise ValueError('Positive sample size required.')
        clock=np.broadcast_to(np.asarray(clock,float),(n,))
        if np.any(~np.isfinite(clock)) or np.any(clock<=0):raise ValueError('Positive clocks required.')
        initial=np.minimum(rng.poisson(float(self.loading),n),self.threshold) if loaded else np.zeros(n,int)
        times=np.zeros(n)
        for z,q in enumerate(self.rates()):
            wait=rng.exponential(1/float(q),n)/clock
            times+=np.where(initial<=z,wait,0.)
        return initial,times


@dataclass(frozen=True)
class IdentityChannel:
    false_positive: F = F(1,100)
    occupied_sensitivity: F = F(99,100)

    def __post_init__(self):
        for name in ['false_positive','occupied_sensitivity']:
            v=rational(getattr(self,name));object.__setattr__(self,name,v)
            if not 0<=v<=1:raise ValueError('Channel probabilities must lie in [0,1].')

    def guaranteed_errors(self,bounds):
        C=max(F(0),1-bounds.miss[1]-bounds.empty[1]*bounds.blank[1])
        return dict(blank_upper=self.false_positive*bounds.blank[1],miss_upper=1-self.occupied_sensitivity*C,occupied_hit_lower=C)

    def joint_fixture(self,source,t):
        """Actual three-category numerical law for a constant conditional channel.

        Categories: hit/identity+, hit/identity-, no hit. This fixture preserves
        the amplification source; it is not an empirical specificity claim.
        """
        B,M=source.numerical_errors(t);e=np.exp(-float(source.loading));C=1-M-e*B
        f,s=float(self.false_positive),float(self.occupied_sensitivity)
        return dict(blank=[f*B,(1-f)*B,1-B],loaded=[s*C+f*e*B,(1-s)*C+(1-f)*e*B,M])


def fixed_rate_window(source,tlo,thi,fraction=F(1,50),loading_floor=F(399,100)):
    """Holding-time coupling for fixed rates in coordinatewise +/-fraction bands."""
    tlo,thi,fraction,loading_floor=map(rational,[tlo,thi,fraction,loading_floor])
    if not 0<=tlo<=thi or not 0<=fraction<1 or loading_floor<0:raise ValueError('Invalid uncertainty box.')
    from dataclasses import replace
    floor=replace(source,loading=loading_floor)
    fast=floor.certify((1+fraction)*thi);slow=floor.certify((1-fraction)*tlo)
    return ErrorBounds((F(0),fast.blank[1]),(F(0),slow.miss[1]),slow.empty)