"""Literal ordered-split polymer chemistry; feed, washout and marked rewards."""
from __future__ import annotations
from dataclasses import dataclass
from collections import Counter
from itertools import product
import math
import numpy as np
from scipy.integrate import solve_ivp

FOOD=('0','1','00','01','10','11')
SELECTED=('00','11','0011')
EPS=1/500000000


@dataclass(frozen=True)
class Channel:
    label: str
    inputs: tuple
    outputs: tuple
    coefficient: float
    kind: str = 'internal'
    selected_birth: bool = False

    @property
    def change(self):
        a,b=Counter(self.inputs),Counter(self.outputs)
        return {z:b[z]-a[z] for z in a.keys()|b.keys() if b[z]!=a[z]}

    def rate(self, counts, volume):
        value=self.coefficient*volume
        for word,multiplicity in Counter(self.inputs).items():
            for i in range(multiplicity):value*=max(counts[word]-i,0)/volume
        return value


@dataclass(frozen=True)
class CatalyticAssignment:
    basal: float = EPS
    selected: float = 4.
    background: tuple = ()  # (catalyst, (left,right,product), coefficient), shared by both directions
    weak_food: float = 0.

    def validate(self, catalogue, eta):
        if not EPS <= self.basal <= 4*EPS or not 4 <= self.selected <= 16:
            raise ValueError('Basal or selected coefficient outside uncertainty class')
        if not 0 <= self.weak_food <= eta <= .03:raise ValueError('Invalid weak food bound')
        seen=set()
        for word,split,k in self.background:
            if word not in catalogue.words or split not in catalogue.splits:raise ValueError('Unknown catalyst/split')
            if (word,split) in seen or (word=='0011' and split==SELECTED):raise ValueError('Duplicate incidence')
            seen.add((word,split))
            if word in FOOD:
                if self.weak_food or not 0<=k<=eta:raise ValueError('Duplicate/out-of-class food incidence')
            elif not (k==0 or 4<=k<=16):raise ValueError('Nonfood coefficients are zero or in [4,16]')


class PolymerCatalogue:
    def __init__(self,n=4):
        if not isinstance(n,int) or not 4<=n<=10:raise ValueError('Explicit catalogue supports n=4..10; analytic certificate has no catalogue cutoff')
        self.n=n
        self.words=tuple(''.join(s) for k in range(1,n+1) for s in product('01',repeat=k))
        self.splits=tuple((z[:j],z[j:],z) for z in self.words for j in range(1,len(z)))
        self.index={z:i for i,z in enumerate(self.words)}

    def channels(self,assignment,eta):
        assignment.validate(self,eta)
        incidences={split:[] for split in self.splits}
        incidences[SELECTED].append(('0011',assignment.selected))
        for word,split,k in assignment.background:
            if k:incidences[split].append((word,k))
        channels=[]
        for split in self.splits:
            left,right,joined=split
            cats=[(None,assignment.basal)]+incidences[split]
            if assignment.weak_food:cats += [(z,assignment.weak_food) for z in FOOD]
            for cat,k in cats:
                forward=(left,right)+( (cat,) if cat else ())
                reverse=(joined,)+( (cat,) if cat else ())
                name=f'{left}|{right};{cat or "basal"}'
                channels.append(Channel(name+';forward',forward,reverse,k,selected_birth=split==SELECTED and cat=='0011'))
                channels.append(Channel(name+';reverse',reverse,forward,k))
        channels += [Channel('feed:'+z,(),(z,),1.,'feed') for z in FOOD]
        channels += [Channel('wash:'+z,(z,),(),1.,'wash') for z in self.words]
        return tuple(channels)


class PolymerReactor:
    def __init__(self,catalogue,assignment=CatalyticAssignment(),eta=1e-5):
        self.catalogue=catalogue;self.assignment=assignment;self.eta=eta
        self.channels=catalogue.channels(assignment,eta)
        self.stoich=np.zeros((len(catalogue.words),len(self.channels)))
        self.roles=np.full((len(self.channels),3),len(catalogue.words),dtype=int)
        self.offsets=np.zeros_like(self.roles)
        for j,ch in enumerate(self.channels):
            for word,change in ch.change.items():self.stoich[catalogue.index[word],j]=change
            used=Counter()
            for i,word in enumerate(ch.inputs):
                self.roles[j,i]=catalogue.index[word];self.offsets[j,i]=used[word];used[word]+=1
        self.coeff=np.array([ch.coefficient for ch in self.channels])
        self.length=np.array(list(map(len,catalogue.words)))
        self.nonfood=np.where(self.length>2,self.length,0)
        self.output=np.array([len(ch.inputs[0]) if ch.kind=='wash' and len(ch.inputs[0])>2 else 0 for ch in self.channels])
        self.arrival=np.array([ch.kind=='feed' for ch in self.channels],dtype=int)
        self.feed_mass=np.array([len(ch.outputs[0]) if ch.kind=='feed' else 0 for ch in self.channels])

    def initial(self,volume=1):
        return np.array([volume if z in FOOD else 0 for z in self.catalogue.words])

    def flux(self,x,volume=None):
        """Density flux; if volume supplied, use falling factorial count corrections."""
        values=np.r_[x,1.][self.roles]
        if volume is not None:values=values-self.offsets/volume
        return self.coeff*np.prod(np.maximum(values,0),axis=1)

    def deterministic(self,times):
        size=len(self.catalogue.words)
        # Q, food molecule arrivals, food monomer arrivals, signed synthesis.
        def rhs(t,y):
            rates=self.flux(y[:size]);dx=self.stoich@rates
            internal=np.array([ch.kind=='internal' for ch in self.channels])
            return np.r_[dx,self.output@rates,self.arrival@rates,self.feed_mass@rates,self.nonfood@(self.stoich[:,internal]@rates[internal])]
        sol=solve_ivp(rhs,(0,float(times[-1])),np.r_[self.initial(),np.zeros(4)],t_eval=times,method='LSODA',rtol=2e-9,atol=1e-12)
        if not sol.success:raise RuntimeError(sol.message)
        return sol.y.T

    def simulate(self,volume=20,horizon=199.,seed=55092026,event_budget=100000):
        """Direct marked SSA. Small count diagnostics only; never restart at time 100."""
        if not isinstance(volume,int) or not 1<=volume<=10**9:raise ValueError('Direct SSA volume must be an integer in [1,1e9]')
        if horizon<=0 or event_budget<1:raise ValueError('Positive horizon and event budget required')
        rng=np.random.default_rng(seed);counts=self.initial(volume).astype(np.int64);t=0.;Q=Af=Bf=J=0
        windows=[0,0];observations={};maxmass=int(self.length@counts);events=0
        observations_due=[v for v in [1.,100.,199.] if v<=horizon]
        while t<horizon and events<event_budget:
            rates=volume*self.flux(counts/volume,volume);total=rates.sum()
            next_t=t+rng.exponential(1/total)
            for obs in observations_due:
                if obs not in observations and t<=obs<next_t:
                    observations[obs]=dict(mass=int(self.length@counts),largest_tetramer=max(int(counts[i]) for i,z in enumerate(self.catalogue.words) if len(z)==4))
            if next_t>horizon:t=horizon;break
            mark=min(int(np.searchsorted(np.cumsum(rates),rng.random()*total)),len(rates)-1)
            ch=self.channels[mark];change=self.stoich[:,mark].astype(np.int64);counts+=change;t=next_t;events+=1
            if np.any(counts<0):raise AssertionError('Negative population')
            Q+=int(self.output[mark]);Af+=int(self.arrival[mark]);Bf+=int(self.feed_mass[mark])
            if ch.kind=='internal':J+=int(self.nonfood@change)
            if 1<t<=100:windows[0]+=int(self.output[mark])
            if 100<t<=199:windows[1]+=int(self.output[mark])
            maxmass=max(maxmass,int(self.length@counts))
        complete=t>=horizon
        residual=int(self.nonfood@counts)+Q-J
        if residual:raise AssertionError('Nonfood material ledger mismatch')
        mission=(complete and horizon==199 and len(observations)==3 and maxmass<=11*volume and Af<=1195*volume
                 and all(o['mass']*2<=21*volume and o['largest_tetramer']*3*10**18>=volume for o in observations.values())
                 and all(q*10>volume for q in windows))
        return dict(completed=complete,time=t,events=events,seed=seed,volume=volume,observations=observations,
                    collected=Q,window_outputs=windows,food_arrivals=Af,food_monomers=Bf,signed_synthesis=J,
                    final_nonfood=int(self.nonfood@counts),ledger_residual=residual,maximum_mass=maxmass,
                    mission=mission if complete and horizon==199 else None,
                    scope='One literal small-count trajectory, not a reliability estimate or theorem-scale simulation')


@dataclass(frozen=True)
class RetainedReward:
    volume: int

    def increments(self,counts,channel):
        K=counts['0011'];change=channel.change
        if K<=2:raise ValueError('Logarithmic reward requires pre-jump K>2')
        dk=change.get('0011',0)
        logpart=math.log1p(dk/K)
        kept=logpart if dk<0 or channel.selected_birth else 0.
        # Integer difference avoids losing tiny count jumps at macroscopic V.
        food=0.
        for word in ['00','11']:
            a=max(self.volume-counts[word],0);b=max(self.volume-counts[word]-change.get(word,0),0)
            food += -4*((b-a)*(b+a)/self.volume**2)
        return kept+food,logpart+food,logpart-kept

    def diagnostic(self,reactor,counts):
        mass=sum(len(z)*v for z,v in counts.items())/self.volume
        nonfood=sum(len(z)*v for z,v in counts.items() if len(z)>2)/self.volume
        sums=dict(drift=0.,variance=0.,full_variance=0.,omitted_log_variance=0.,copy_loss=0.,max_jump=0.,mass_drift=0.,output_rate=0.)
        for ch in reactor.channels:
            rate=ch.rate(counts,self.volume)
            if not rate:continue
            r,full,omitted=self.increments(counts,ch)
            if omitted < -1e-14:raise AssertionError('Reward exceeds potential increment')
            sums['drift']+=rate*r;sums['variance']+=rate*r*r;sums['full_variance']+=rate*full*full
            sums['omitted_log_variance']+=rate*omitted*omitted
            sums['copy_loss']+=rate*max(-ch.change.get('0011',0),0)
            sums['max_jump']=max(sums['max_jump'],abs(r))
            sums['mass_drift']+=rate*sum(len(z)*v for z,v in ch.change.items())/self.volume
            if ch.kind=='wash' and len(ch.inputs[0])>2:sums['output_rate']+=rate*len(ch.inputs[0])/self.volume
        return dict(mass=mass,nonfood=nonfood,K=counts['0011'],**sums)