"""Finite-count chemistry, joint allocation, and the actual two-daughter ledger.

Count-rate constants use a fixed reference volume and model time. They are
constructed benchmarks, not concentration-rate constants or calibrated seconds.
"""
from dataclasses import dataclass
from fractions import Fraction as F
from math import factorial, ceil, expm1
import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import expm_multiply
import ctmc_certificate as kernel
from proportion_extras import psi_theta, coordinate_theta, clock_lower


@dataclass(frozen=True)
class SupportRates:
    replication: F = F(1)
    reverse_replication: F = F(1,100)
    binding: F = F(2)
    dissociation: F = F(10)
    release: F = F(10)
    reverse_release: F = F(1,50)

    def values(self):
        values=tuple(F(v) for v in self.__dict__.values())
        if any(v<=0 for v in values): raise ValueError('All six rates must be positive.')
        return values

    def groups(self):
        return [([r],[F(1)],v,v) for r,v in enumerate(self.values())]


@dataclass(frozen=True)
class ComplementaryAllocation:
    theta: F = F(1,2)

    def __post_init__(self):
        if not 0<F(self.theta)<1: raise ValueError('Allocation probability must lie strictly between zero and one.')

    def support_return(self, carriers):
        return F(0) if carriers<1 else 1-F(self.theta)**carriers-(1-F(self.theta))**carriers

    def split(self, objects, rng):
        a=rng.binomial(np.asarray(objects,dtype=int),float(self.theta))
        return a,np.asarray(objects,dtype=int)-a


@dataclass(frozen=True)
class SupportReactor:
    inventory: int = 32
    rates: SupportRates = SupportRates()

    def __post_init__(self):
        if not isinstance(self.inventory,int) or self.inventory<2: raise ValueError('Integer inventory >= 2 required.')
        self.rates.values()

    def food(self,state):
        n,c,p=state
        if any(int(v)!=v or v<0 for v in state): raise ValueError('Nonnegative molecule counts required.')
        f=self.inventory-n-2*c-p
        if f<0: raise ValueError('State exceeds the weighted material pool.')
        return f

    def transitions(self,state):
        n,c,p=state;f=self.food(state)
        for r,k in enumerate(self.rates.values()):
            a=k*kernel.factor(r,n,c,p,f)
            if a:
                dest=tuple(x+y for x,y in zip(state,kernel.INCR[r]))
                self.food(dest)
                yield dest,a,r

    def generator(self):
        states=kernel.states_of(self.inventory);index={z:i for i,z in enumerate(states)}
        rows=[];cols=[];values=[]
        for i,z in enumerate(states):
            for dest,a,_ in self.transitions(z):
                rows.extend([i,i]);cols.extend([index[dest],i]);values.extend([float(a),-float(a)])
        return states,coo_matrix((values,(rows,cols)),shape=(len(states),len(states))).tocsr()

    def deadline(self,time,quota,allocation=ComplementaryAllocation()):
        """Floating backward equation on the FULL finite state space; no transpose.

        Failure splits into disjoint quota failure and quota-met/allocation failure.
        These numerical values never enter an exact certificate.
        """
        if time<0 or quota<1: raise ValueError('Nonnegative time and positive quota required.')
        states,G=self.generator();H=[]
        for n,c,p in states:
            q=int(p>=quota);s=float(allocation.support_return(n+c))
            H.append([q*s,1-q,q*(1-s),n+c,self.food((n,c,p)),p])
        result=expm_multiply(time*G,np.array(H))
        return states,result

    def trajectory(self,start,time,rng,max_events=2_000_000):
        self.food(start);state=tuple(start);t=0.;events=0
        while t<time:
            transitions=list(self.transitions(state));total=sum(float(a) for _,a,_ in transitions)
            if total==0: break
            t+=rng.exponential(1/total)
            if t>time: break
            if events>=max_events: raise RuntimeError('Event budget exhausted; no completed trajectory returned.')
            chosen=rng.choice(len(transitions),p=[float(a)/total for _,a,_ in transitions])
            state=transitions[chosen][0];events+=1
        return state,events


@dataclass(frozen=True)
class IntegerEnvelope:
    inventory: int
    groups: tuple
    deadline: int = 5
    quota: int = 4
    allocation: ComplementaryAllocation = ComplementaryAllocation()

    def run(self):
        """Exact endpoint-kernel lower envelope; complementary payoff gives upper.

        An upper bound below target excludes the robust specification. A lower
        bound below target alone is inconclusive. It is not the exact robust value.
        """
        SupportReactor(self.inventory)
        if not isinstance(self.deadline,int) or self.deadline<0 or self.quota<1: raise ValueError('Integer nonnegative deadline and positive quota required by block implementation.')
        channels=[r for rs,_,_,_ in self.groups for r in rs]
        if sorted(channels)!=list(range(6)): raise ValueError('Each channel must belong to exactly one parameter group.')
        for rs,ms,lo,hi in self.groups:
            if len(rs)!=len(ms) or not 0<lo<=hi or any(F(m).denominator!=1 or m<=0 for m in ms): raise ValueError('Positive ordered rational bounds and positive integer multipliers required.')
        result=kernel.certify(self.inventory,list(self.groups),T=self.deadline,quota=self.quota,split=self.allocation.support_return)
        result.pop('seconds',None)
        return result


@dataclass(frozen=True)
class SerialProtocol:
    reactor: SupportReactor
    deadline: float = 5
    quota: int = 4
    allocation: ComplementaryAllocation = ComplementaryAllocation()

    def lineage(self,start,cycles,rng):
        """Always retain daughter A, including after failures; never select a survivor."""
        if cycles<0 or self.deadline<0 or self.quota<1: raise ValueError('Invalid protocol.')
        rows=[];state=tuple(start);K=self.reactor.inventory;supplied=K;harvest=0
        for cycle in range(1,cycles+1):
            terminal,events=self.reactor.trajectory(state,self.deadline,rng)
            n,c,p=terminal;f=self.reactor.food(terminal)
            a,b=self.allocation.split([n,c,f],rng)
            ma=int(a[0]+2*a[1]+a[2]);mb=int(b[0]+2*b[1]+b[2])
            refill_a=K-ma;refill_b=K-mb
            assert ma+mb+p==K and refill_a+refill_b==K+p
            returns=[1<=int(d[0]+d[1])<=K-1 for d in (a,b)]
            supplied+=refill_a+refill_b;harvest+=p
            rows.append(dict(cycle=cycle,start=list(state),terminal=list(terminal),events=events,
                daughter_a=[int(v) for v in a],daughter_b=[int(v) for v in b],
                refill_a=refill_a,refill_b=refill_b,joint_success=bool(p>=self.quota and all(returns)),
                total_supplied=supplied,total_harvested=harvest))
            state=(int(a[0]),int(a[1]),0)
        assert supplied==K+cycles*K+harvest
        return rows


@dataclass(frozen=True)
class Corridor:
    low: int
    high: int
    growth: int

    def __post_init__(self):
        if any(not isinstance(x,int) for x in (self.low,self.high,self.growth)) or not 1<=self.low<=self.growth<=self.high: raise ValueError('Corridor requires 1 <= low <= growth <= high.')

    @property
    def cost(self): return self.high+self.growth

    def uniform_return(self,theta=F(1,2)):
        ComplementaryAllocation(theta)
        return coordinate_theta(self.low,self.high,self.growth,F(theta))


def falling(n,r):
    return factorial(n)//factorial(n-r) if n>=r else 0


@dataclass(frozen=True)
class CooperativeGate:
    majority: Corridor
    minority: Corridor
    quota: int = 4
    forward_floor: F = F(20)
    reverse_ceiling: F = F(1,100000)

    def __post_init__(self):
        if self.majority.low<=self.minority.high or self.quota<1 or self.forward_floor<=0 or self.reverse_ceiling<=0: raise ValueError('Separated corridors, positive quota and rate bounds required.')

    @property
    def core(self): return self.majority.cost+self.minority.cost+self.quota

    @property
    def food_consumed(self): return self.majority.growth+self.minority.growth+self.quota

    @property
    def molecularity(self): return self.majority.low+self.minority.low+self.food_consumed+1

    def constants(self):
        x,y=self.majority,self.minority
        return (F(self.forward_floor,factorial(x.low)*factorial(y.low)*factorial(self.food_consumed)),
            F(self.reverse_ceiling,falling(x.cost,x.low+x.growth)*falling(y.cost,y.low+y.growth)*factorial(self.quota)))

    def pair_rates(self,x,y):
        X,Y=self.majority,self.minority
        if not X.low<=x<=X.high or not Y.low<=y<=Y.high: raise ValueError('Outside restart rectangle.')
        k,b=self.constants()
        return (k*falling(x,X.low)*falling(y,Y.low)*falling(self.core-x-y,self.food_consumed),
            b*falling(x+X.growth,X.low+X.growth)*falling(y+Y.growth,Y.low+Y.growth)*factorial(self.quota))

    def transitions(self,state):
        """All FOUR literal channels, species (X,Y,F,H,W,PX,PY).

        No symmetry projection is used to decide which channel is enabled.
        Changing corridors or adding channels calls for checking closure again.
        """
        if len(state)!=7 or any(int(n)!=n or n<0 for n in state): raise ValueError('Seven nonnegative integer counts required.')
        x,y,f,h,w,px,py=state
        if x+y+f+px+py!=self.core or h+w!=1: raise ValueError('Outside the conserved material/fuel class.')
        X,Y=self.majority,self.minority;k,b=self.constants()
        selected_reactants=(X.low,Y.low,self.food_consumed,1,0,0,0)
        selected_products=(X.low+X.growth,Y.low+Y.growth,0,0,1,self.quota,0)
        mirror=lambda v:(v[1],v[0],v[2],v[3],v[4],v[6],v[5])
        for label,reactants,products,rate in [
            ('selected_forward',selected_reactants,selected_products,k),
            ('selected_reverse',selected_products,selected_reactants,b),
            ('mirror_forward',mirror(selected_reactants),mirror(selected_products),k),
            ('mirror_reverse',mirror(selected_products),mirror(selected_reactants),b)]:
            a=rate
            for n,r in zip(state,reactants):a*=falling(n,r)
            if a:yield tuple(n-r+p for n,r,p in zip(state,reactants,products)),a,label

    def joint_lower(self,theta=F(1,2),time=F(1)):
        if time<0: raise ValueError('Nonnegative deadline required.')
        return clock_lower(F(self.forward_floor),F(self.reverse_ceiling),F(time))*self.majority.uniform_return(theta)*self.minority.uniform_return(theta)

    def actual_joint(self,x,y,theta=F(1,2),time=1.):
        """Numerical display of the exact two-state law, not its certificate."""
        a,d=map(float,self.pair_rates(x,y));clock=a/(a+d)*(-expm1(-(a+d)*time))
        X,Y=self.majority,self.minority
        return clock*float(psi_theta(x+X.growth,X.low,X.high,F(theta))*psi_theta(y+Y.growth,Y.low,Y.high,F(theta)))


def carrier_floor(quota,types,delta):
    """Exact integer version of q + m ceil(1 + log2(1/delta))."""
    delta=F(delta)
    if quota<1 or types<1 or not 0<delta<1: raise ValueError('Invalid carrier specification.')
    copies=1
    while F(1,2**(copies-1))>delta: copies+=1
    return quota+types*copies