"""Literal donor/refined channels, common exchange, pulses, and physical ledgers."""
from dataclasses import dataclass
from fractions import Fraction as F
import math
import copy
import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import expm

SPECIES=('U','W','X','C1','C2','Z','D')
A=np.array([1,0,1,2,2,2,1]);B=np.array([0,1,1,1,2,2,1])
I=np.array([0,0,1,1,1,2,0]);J=I.copy();J[6]=1
Y=np.array([0,0,1,9/8,7/5,9/5,0]);Y560=np.array([0,0,560,630,784,1008,0])
EPS=F(1,500000000);ETA=F(1,8000000000)
ACCOUNTS=('QI','QX','gross_service','all_wash_J','net_synthesis','upstream_minus_downstream','wash_D')


@dataclass(frozen=True)
class Channel:
    name: str
    inputs: tuple
    outputs: tuple
    k: F
    size: int
    pair: int=-1
    direction: int=0
    wash: int=-1
    feed: int=-1

    @property
    def jump(self):return tuple(self.outputs.count(i)-self.inputs.count(i) for i in range(self.size))
    def propensity(self,N,V):
        value=self.k*F(V)**(1-len(self.inputs));used={}
        for i in self.inputs:
            n=N[i]-used.get(i,0)
            if n<=0:return F(0)
            value*=n;used[i]=used.get(i,0)+1
        return value
    def marks(self,collect):
        washj=int(J[self.wash]) if self.wash>=0 else 0
        netj=sum(int(w)*v for w,v in zip(J,self.jump)) if self.pair>=0 else 0
        storage=(self.direction if self.pair==5 else -self.direction if self.pair==6 else 0) if self.size==7 else 0
        return (int(I[self.wash])*int(collect) if self.wash>=0 else 0,
            int(collect and self.wash==2),int(self.pair==5),washj,netj,storage,int(self.wash==6))


class Chemistry:
    def __init__(self,r='20',d='3/100',theta='1/100',refined=True):
        self.r,self.d,self.theta=map(F,(r,d,theta));self.beta=F(1,100);self.refined=refined
        if min(self.r,self.d,self.theta)<=0:raise ValueError('Positive rates and lifetime required')
        self.size=7 if refined else 6
        pairs=[((0,1),(2,),EPS,EPS/10),((2,0),(3,),F(20),F(20)),
            ((3,1),(4,),F(20),F(20)),((4,),(5,),F(20),F(2)),((5,),(2,2),self.r,self.r)]
        if refined:
            pairs += [((2,),(6,),self.d*(1+self.beta),self.d*self.beta/self.theta),
                ((6,),(0,1),self.d/self.theta,self.d*ETA*(1+self.beta)/self.beta)]
        else:pairs += [((2,),(0,1),self.d,self.d*ETA)]
        channels=[]
        for j,(a,b,kf,kr) in enumerate(pairs):
            channels += [Channel(f'pair{j}+',a,b,kf,self.size,j,1),Channel(f'pair{j}-',b,a,kr,self.size,j,-1)]
        channels += [Channel(f'feed-{SPECIES[i]}',(),(i,),F(1),self.size,feed=i) for i in range(2)]
        channels += [Channel(f'wash-{SPECIES[i]}',(i,),(),F(1),self.size,wash=i) for i in range(self.size)]
        self.channels=tuple(channels);self.jumps=np.array([c.jump for c in channels])
        self.coefficients=np.array([float(c.k) for c in channels]);self.orders=np.array([len(c.inputs) for c in channels])
        self.indices=np.full((len(channels),2),self.size,int);self.offsets=np.zeros_like(self.indices)
        for j,c in enumerate(channels):
            for k,i in enumerate(c.inputs):self.indices[j,k]=i;self.offsets[j,k]=c.inputs[:k].count(i)
        self.mark_off=np.array([c.marks(False) for c in channels]);self.mark_on=np.array([c.marks(True) for c in channels])

    def scope(self,stochastic=False):
        return 19<=self.r<=21 and F(1,50)<=self.d<=F(1,25) and (not self.refined or
            (self.theta==F(1,100) if stochastic else self.theta<=F(1,100)))
    def rates(self,state,V=None):
        values=np.r_[state,1.][self.indices]
        if V is not None:values=values-self.offsets
        rates=self.coefficients*np.prod(values,axis=1)
        return rates if V is None else rates*np.power(float(V),1-self.orders)
    def field(self,c):return self.rates(c)@self.jumps
    def filter_equilibrium(self,c):
        if not self.refined:raise ValueError('No intermediate in donor')
        u,w,x=c[:3];d=float(self.d);b=float(self.beta);t=float(self.theta)
        return d*(1+b)*t/(t+d*(1+b))*(x+float(ETA)/b*u*w)


class CommonExchange:
    def __init__(self,weights):
        k=np.array(weights,float)
        if k.ndim!=2 or len(k)==0 or k.shape[0]!=k.shape[1] or not np.isfinite(k).all() or np.any(k<0) or np.any(k.diagonal()!=0) or not np.array_equal(k,k.T):
            raise ValueError('Finite symmetric nonnegative zero-diagonal graph required')
        self.k=k;self.n=len(k);self.D=k-np.diag(k.sum(axis=1));self.degree=float(max(k.sum(axis=1)))
        self.exact_degree=max(sum(F.from_float(float(v)) for v in row) for row in k)
    def material(self,initial,t):return 1+expm((self.D-np.eye(self.n))*t)@(initial-1)


@dataclass(frozen=True)
class Intervention:
    q: object='1/4'
    survival: tuple=('49/50',)*7
    error: tuple=('0','0')
    def __post_init__(self):
        object.__setattr__(self,'q',F(self.q));object.__setattr__(self,'survival',tuple(map(F,self.survival)));object.__setattr__(self,'error',tuple(map(F,self.error)))
        if not F(1,4)<=self.q<=F(3,4) or len(self.survival) not in (6,7) or any(not F(49,50)<=s<=1 for s in self.survival) or len(self.error)!=2 or any(abs(e)>F(1,200) for e in self.error):
            raise ValueError('Pulse outside paper bounds')
    def deterministic(self,c):
        if len(c)!=len(self.survival):raise ValueError('Survival size mismatch')
        q=float(self.q);ell=np.array(list(map(float,self.survival)))
        retained=q*ell*c;dose=np.array([float(1-self.q+e) for e in self.error]);retained[:2]+=dose
        return retained,(1-q)*c,q*(1-ell)*c,dose
    def molecular(self,N,V,rng):
        if len(N)!=len(self.survival) or any(type(int(n)) is not int or n<0 or int(n)!=n for n in N):raise ValueError('Nonnegative integer counts required')
        if max(N)>=2**53:raise ValueError('Sampler restricted to counts below 2^53')
        draws=np.array([rng.multinomial(int(n),[float(self.q*l),float(1-self.q),float(self.q*(1-l))]) for n,l in zip(N,self.survival)],dtype=np.int64)
        dose=np.array([math.floor(V*(1-self.q+e)) for e in self.error],dtype=np.int64)
        after=draws[:,0].copy();after[:2]+=dose
        return after,draws[:,1],draws[:,2],dose


class HistoryPolicy:
    """Fixed pulse by default; optional feedback reads actual completed accounts."""
    def __init__(self,size=7,feedback=False):self.size=size;self.feedback=feedback
    def choose(self,current,history):
        actions=[]
        for i in range(len(current)):
            q='3/4' if self.feedback and history and history[-1]['accounts'][i][0]<.1 else '1/4'
            actions.append(Intervention(q,('49/50',)*self.size))
        return actions


def ready(state,theta=F(1,100),V=None):
    """Exact integer inequalities for counts; tolerance only in density diagnostics."""
    state=np.asarray(state);s=state.shape[-1]
    if s not in (6,7) or np.any(state<0):return False
    if V is not None:
        if type(V) is not int or V<=0 or not np.equal(state,np.floor(state)).all():raise ValueError('Integer count state and scale required')
        for row in state.reshape(-1,s):
            vals=list(map(int,row));dot=lambda w:sum(int(a)*v for a,v in zip(w,vals))
            if not 159*V<=160*dot(A)<=161*V or not 159*V<=160*dot(B)<=161*V or dot(Y560)<28*V:return False
            if s==7 and F(vals[6],V)>F(1101,1000)*theta:return False
        return True
    return bool(np.all(state@A[:s]>=159/160-1e-12) and np.all(state@A[:s]<=161/160+1e-12)
        and np.all(state@B[:s]>=159/160-1e-12) and np.all(state@B[:s]<=161/160+1e-12)
        and np.all(state@Y[:s]>=.05-1e-12) and (s==6 or np.all(state[...,6]<=float(F(1101,1000)*theta)+1e-12)))


class ReactorNetwork:
    def __init__(self,chemistries,exchange):
        self.chemistries=tuple(chemistries);self.exchange=exchange;self.n=exchange.n
        if len(self.chemistries)!=self.n or len({c.size for c in self.chemistries})!=1:raise ValueError('One same-size chemistry per node')
        self.size=self.chemistries[0].size
    def validate(self,c):
        c=np.array(c,float)
        if c.shape!=(self.n,self.size) or not np.isfinite(c).all() or np.any(c<0):raise ValueError('Invalid initial state')
        return c
    def field(self,c):return np.array([m.field(row) for m,row in zip(self.chemistries,c)])+self.exchange.D@c
    def evolve(self,c,rtol=2e-9):
        c=self.validate(c);ns=self.n*self.size;v=np.r_[c.ravel(),np.zeros(7*self.n)];ts=[];cs=[]
        initial_material=np.column_stack((c@A[:self.size],c@B[:self.size]))
        for start,end,collect in [(0,3,False),(3,4,True)]:
            def rhs(t,v):
                state=v[:ns].reshape(self.n,self.size)
                accounts=np.array([m.rates(row)@(m.mark_on if collect else m.mark_off) for m,row in zip(self.chemistries,state)])
                return np.r_[self.field(state).ravel(),accounts.ravel()]
            grid=np.linspace(start,end,round((end-start)*60)+1)
            sol=solve_ivp(rhs,(start,end),v,method='Radau',rtol=rtol,atol=rtol/100,t_eval=grid)
            if not sol.success:raise RuntimeError(sol.message)
            v=sol.y[:,-1];skip=int(start!=0);ts.extend(grid[skip:]);cs.extend(sol.y[:ns,skip:].T.reshape(-1,self.n,self.size))
        samples=np.array(cs);times=np.array(ts);final=v[:ns].reshape(self.n,self.size);acc=v[ns:].reshape(self.n,7)
        if samples.min() < -1e-9:raise RuntimeError('Negative numerical state; no clipping')
        exact=np.array([self.exchange.material(initial_material,t) for t in times])
        residual=float(np.max(abs(np.stack((samples@A[:self.size],samples@B[:self.size]),axis=2)-exact)))
        return dict(end=final,accounts=acc,times=times,states=samples,material_error=residual)

    def donor_reduced(self,c):
        """Exact deterministic material forcing; never used as a count-path closure."""
        if self.size!=6:raise ValueError('This reduction is for the donor')
        c=self.validate(c);materials=np.column_stack((c@A[:6],c@B[:6]));phase=c[:,2:]
        def reconstruct(t,p):
            p=p.reshape(self.n,4);M=self.exchange.material(materials,t)
            u=M[:,0]-p@np.array([1,2,2,2]);w=M[:,1]-p@np.array([1,1,2,2])
            return np.column_stack((u,w,p))
        sol=solve_ivp(lambda t,p:self.field(reconstruct(t,p))[:,2:].ravel(),(0,4),phase.ravel(),method='Radau',rtol=2e-10,atol=2e-12,dense_output=True)
        if not sol.success:raise RuntimeError(sol.message)
        return reconstruct(4,sol.y[:,-1])


class DeterministicMission:
    def __init__(self,network,policy):self.network=network;self.policy=policy
    def run(self,initial,cycles,rtol=2e-9):
        if type(cycles) is not int or cycles<0:raise ValueError('Nonnegative integer cycles required')
        net=self.network;c=net.validate(initial);s=net.size;initialJ=float((c@J[:s]).sum());initialD=float(c[:,6].sum()) if s==7 else 0.
        history=[];traces=[];removedJ=lostJ=removedD=lostD=0.;tot=np.zeros(7);maxmaterial=0.
        for cycle in range(cycles):
            actions=self.policy.choose(c.copy(),copy.deepcopy(history));post=[];food=[]
            if len(actions)!=net.n:raise ValueError('One action per node')
            for row,action in zip(c,actions):
                after,removed,lost,dose=action.deterministic(row);post.append(after);food.append(dose+4)
                removedJ+=removed@J[:s];lostJ+=lost@J[:s]
                if s==7:removedD+=removed[6];lostD+=lost[6]
            flow=net.evolve(post,rtol);c=flow['end'];acc=flow['accounts'];tot+=acc.sum(axis=0);maxmaterial=max(maxmaterial,flow['material_error'])
            history.append(dict(cycle=cycle+1,endpoint=c.tolist(),accounts=acc.tolist(),food=np.array(food).tolist(),ready=ready(c,net.chemistries[0].theta)))
            traces.append(flow)
        finalJ=float((c@J[:s]).sum());finalD=float(c[:,6].sum()) if s==7 else 0.
        return dict(history=history,traces=traces,final=c,totals=tot,
            inventory_telescope_residual=float(tot[4]-(finalJ-initialJ+tot[3]+removedJ+lostJ)),
            storage_telescope_residual=float(tot[5]-(finalD-initialD+tot[6]+removedD+lostD)),
            material_error=maxmaterial,withdrawn_J=float(removedJ),lost_J=float(lostJ))


class MolecularMission:
    """Direct SSA. Event limits return unfinished paths, never successes or failures."""
    def __init__(self,network,policy,V,seed=51092026):
        if type(V) is not int or V<=0:raise ValueError('Positive integer scale required')
        self.network=network;self.policy=policy;self.V=V;self.rng=np.random.default_rng(seed)
    def run(self,initial,cycles,event_budget):
        if type(cycles) is not int or cycles<0 or type(event_budget) is not int or event_budget<0:raise ValueError('Nonnegative integer cycles/budget required')
        net=self.network;s=net.size;N=np.asarray(initial)
        if N.shape!=(net.n,s) or np.any(N<0) or not np.equal(N,np.floor(N)).all():raise ValueError('Nonnegative integer node counts required')
        N=N.astype(np.int64);history=[];events=0;initialJ=int((N@J[:s]).sum());initialD=int(N[:,6].sum()) if s==7 else 0
        totals=np.zeros(7,dtype=np.int64);remJ=lossJ=remD=lossD=0
        exchange=[(i,j,k,float(net.exchange.k[i,j])) for i in range(net.n) for j in range(net.n) if net.exchange.k[i,j]>0 for k in range(s)]
        for cycle in range(cycles):
            # Policy gets normalized prior accounts, but actual integer logs remain in history.
            observations=[dict(h,accounts=(np.array(h['accounts'])/self.V).tolist()) for h in history]
            actions=self.policy.choose(N.copy()/self.V,observations);food=np.zeros((net.n,2),dtype=np.int64);acc=np.zeros((net.n,7),dtype=np.int64)
            for i,action in enumerate(actions):
                N[i],removed,lost,food[i]=action.molecular(N[i],self.V,self.rng)
                remJ+=int(removed@J[:s]);lossJ+=int(lost@J[:s])
                if s==7:remD+=int(removed[6]);lossD+=int(lost[6])
            t=0.
            while t<4:
                local=[m.rates(row,self.V) for m,row in zip(net.chemistries,N)]
                lengths=[len(a) for a in local];rates=np.r_[np.concatenate(local),[k*N[i,h] for i,j,h,k in exchange]];total=float(rates.sum())
                if total<=0:t=4.;break
                if events>=event_budget:
                    return dict(status='unfinished',events=events,completed_cycles=len(history),time=t,counts=N.tolist(),history=history)
                dt=self.rng.exponential(1/total)
                if t+dt>4:t=4.;break
                t+=dt;label=int(np.searchsorted(np.cumsum(rates),self.rng.random()*total,side='right'));label=min(label,len(rates)-1)
                if label<sum(lengths):
                    i=0
                    while label>=lengths[i]:label-=lengths[i];i+=1
                    channel=net.chemistries[i].channels[label];N[i]+=channel.jump
                    marks=np.array(channel.marks(t>3));acc[i]+=marks;totals+=marks
                    if channel.feed>=0:food[i,channel.feed]+=1
                else:
                    i,j,h,k=exchange[label-sum(lengths)];N[i,h]-=1;N[j,h]+=1
                if N.min()<0:raise ArithmeticError('Disabled event was fired')
                events+=1
            success=ready(N,net.chemistries[0].theta,self.V) and bool(np.all(acc[:,0]>=-(-self.V//56)) and np.all(acc[:,1]>=-(-self.V//1080)) and np.all(food<=5*self.V) and np.all(acc[:,2]<=self.V//5))
            history.append(dict(cycle=cycle+1,counts=N.tolist(),accounts=acc.tolist(),food=food.tolist(),success=success))
            # Failed outcomes stay in the law and are carried into the next cycle.
        finalJ=int((N@J[:s]).sum());finalD=int(N[:,6].sum()) if s==7 else 0
        residual=int(totals[4])-(finalJ-initialJ+int(totals[3])+remJ+lossJ)
        storage=int(totals[5])-(finalD-initialD+int(totals[6])+remD+lossD)
        if residual or storage:raise ArithmeticError('Physical event telescope failed')
        return dict(status='completed',events=events,history=history,all_success=all(h['success'] for h in history),
            counts=N.tolist(),inventory_telescope_residual=residual,storage_telescope_residual=storage)