"""Literal polymer source objects and count/density reactor, shared with the companion model.

Self-contained copy of the original Lixiviant companion example components,
with incidence labels retained for local attribution. No cross-package imports.
"""
from __future__ import annotations
from dataclasses import dataclass
from collections import Counter
import itertools
import math
import numpy as np
from scipy.integrate import solve_ivp

# Fixed manuscript conventions; study inputs are at the top of example.py.
MAX_CATALOGUE_SPECIES = 1022
KINETIC_MARKS = (1.0, 1.5, 2.0)
BASAL_EPSILON = 1/500_000_000
CATALYTIC_SCALE = 4.0
MARK_SEED = 30092026
WINDOW_START, WINDOW_END = 1.0, 100.0
MASS_CEILING, EXPORT_THRESHOLD = 11.0, 0.1
MAX_EVENTS = 200_000
@dataclass(frozen=True)
class Channel:
    left: int
    right: int
    product: int


class PolymerCatalogue:
    def __init__(self, n):
        if not isinstance(n,int) or n < 2 or 2**(n+1)-2 > MAX_CATALOGUE_SPECIES:
            raise ValueError('Catalogue length outside the explicit enumeration guard.')
        self.n = n
        self.words = tuple(''.join(bits) for size in range(1,n+1)
                           for bits in itertools.product('01',repeat=size))
        self.index = {word:i for i,word in enumerate(self.words)}
        self.channels = tuple(Channel(self.index[word[:cut]],self.index[word[cut:]],self.index[word])
                              for word in self.words for cut in range(1,len(word)))
        self.lengths = np.array([len(word) for word in self.words],dtype=np.int64)
        self.nonfood = self.lengths*(self.lengths > 2)
        self.food = frozenset(np.flatnonzero(self.lengths <= 2).tolist())
        self.channel_index = {(self.words[c.left],self.words[c.right]):i for i,c in enumerate(self.channels)}

    def closure(self, selected):
        available = set(self.food)
        while True:
            previous = len(available)
            for i in selected:
                c = self.channels[i]
                if c.left in available and c.right in available:
                    available.add(c.product)
                if c.product in available:
                    available.update((c.left,c.right))
            if len(available) == previous:
                return frozenset(available)


@dataclass(frozen=True)
class CatalyticEnvironment:
    catalogue: PolymerCatalogue
    rows: tuple[frozenset[int], ...]

    def __post_init__(self):
        if len(self.rows) != len(self.catalogue.words) or any(
                not 0 <= r < len(self.catalogue.channels) for row in self.rows for r in row):
            raise ValueError('Catalytic incidence outside catalogue.')

    @classmethod
    def specified(cls, catalogue, assignments=()):
        rows = [set() for _ in catalogue.words]
        for catalyst,left,right in assignments:
            rows[catalogue.index[catalyst]].add(catalogue.channel_index[left,right])
        return cls(catalogue,tuple(frozenset(row) for row in rows))

    def max_raf(self, container=None):
        active = set(range(len(self.catalogue.channels))) if container is None else set(container)
        while active:
            closure = self.catalogue.closure(active)
            catalysed = set().union(*(self.rows[x] for x in closure))
            retained = {i for i in active if i in catalysed and all(v in closure for v in
                (self.catalogue.channels[i].left,self.catalogue.channels[i].right,self.catalogue.channels[i].product))}
            if retained == active:
                return frozenset(active)
            active = retained
        return frozenset()

    def irreducible_raf(self):
        """Deletion-minimal witness; deliberately NOT a minimum-size solver."""
        active = self.max_raf()
        for r in sorted(active):
            smaller = self.max_raf(active-{r})
            if smaller:
                active = smaller
        return active

    def witness_event(self):
        c = self.catalogue
        if c.n < 4:
            return False
        return (all(not self.rows[f] for f in c.food)
                and c.channel_index['00','11'] in self.rows[c.index['0011']])


@dataclass(frozen=True)
class KineticParameters:
    basal: tuple[float, ...]
    catalytic: dict[tuple[int,int],float]

    @classmethod
    def sample(cls, environment, seed=MARK_SEED, epsilon=BASAL_EPSILON,
               catalytic_scale=CATALYTIC_SCALE, reproduce_host_marks=False):
        rng = np.random.default_rng(seed)
        R = len(environment.catalogue.channels)
        marks = rng.choice(KINETIC_MARKS,size=(3 if reproduce_host_marks else 2,R))
        basal = tuple(epsilon*marks[0]*marks[1])
        coefficients = {}
        incidences = sorted((r,x) for x,row in enumerate(environment.rows) for r in row)
        if reproduce_host_marks and len({r for r,x in incidences}) != len(incidences):
            raise ValueError('Figure-mark adapter assumes at most one catalyst per channel.')
        for r,x in incidences:
            h = marks[2,r] if reproduce_host_marks else rng.choice(KINETIC_MARKS)
            coefficients[r,x] = float(catalytic_scale*marks[0,r]*h)
        return cls(basal,coefficients)


@dataclass(frozen=True)
class ReactionEvent:
    inputs: tuple[int, ...]
    outputs: tuple[int, ...]
    coefficient: float
    kind: str
    incidence: tuple[int,int] | None = None

    def propensity(self, counts, volume):
        rate = self.coefficient*volume**(1-len(self.inputs))
        for x,multiplicity in Counter(self.inputs).items():
            for offset in range(multiplicity):
                rate *= max(0,int(counts[x])-offset)
        return rate


class FedReactor:
    def __init__(self, environment, parameters, enabled=True):
        self.environment,self.parameters,self.enabled = environment,parameters,enabled
        self.catalogue = c = environment.catalogue
        if len(parameters.basal) != len(c.channels) or min(parameters.basal) <= 0:
            raise ValueError('Each channel needs a positive basal coefficient.')
        events = []
        for r,channel in enumerate(c.channels):
            left,right,product = channel.left,channel.right,channel.product
            events.extend((ReactionEvent((left,right),(product,),parameters.basal[r],'basal'),
                           ReactionEvent((product,),(left,right),parameters.basal[r],'basal')))
            if enabled:
                for x,row in enumerate(environment.rows):
                    if r in row:
                        k = parameters.catalytic[r,x]
                        if k <= 0:
                            raise ValueError('Positive catalytic coefficient required.')
                        events.extend((ReactionEvent((left,right,x),(product,x),k,'catalytic',(r,x)),
                                       ReactionEvent((product,x),(left,right,x),k,'catalytic',(r,x))))
        self.events = tuple(events)
        count = len(c.words)
        if count*len(events)>20_000_000:
            raise ValueError('Dense reactor matrix exceeds 20 million entries; use a sparse backend for this environment.')
        self.stoich = np.zeros((count,len(events)),dtype=np.int64)
        self.input_indices = np.full((len(events),3),count,dtype=np.int64)
        self.offsets = np.zeros((len(events),3),dtype=np.int64)
        self.orders = np.array([len(e.inputs) for e in events])
        self.coefficients = np.array([e.coefficient for e in events])
        for j,e in enumerate(events):
            used = Counter()
            for k,x in enumerate(e.inputs):
                self.input_indices[j,k] = x
                self.offsets[j,k] = used[x]
                used[x] += 1
                self.stoich[x,j] -= 1
            for x in e.outputs:
                self.stoich[x,j] += 1
        assert np.all(c.lengths@self.stoich == 0)
        self.nf_changes = c.nonfood@self.stoich
        self.is_catalytic = np.array([e.kind == 'catalytic' for e in events])
        self.feed = (c.lengths <= 2).astype(float)

    def rates(self, values, volume=None):
        padded = np.r_[values,1]
        factors = padded[self.input_indices]
        if volume is not None:
            factors = np.maximum(factors-self.offsets,0)
        rates = self.coefficients*np.prod(factors,axis=1)
        return rates if volume is None else rates*np.power(float(volume),1-self.orders)

    def rhs(self, time, state):
        c = self.catalogue
        x = state[:len(c.words)]
        rates = self.rates(x)
        signed = self.nf_changes*rates
        return np.r_[self.feed-x+self.stoich@rates, c.nonfood@x,
                     np.sum(signed[self.is_catalytic]),np.sum(signed[~self.is_catalytic]),
                     np.maximum(self.nf_changes[~self.is_catalytic],0)@rates[~self.is_catalytic]]

    def deterministic(self, times, rtol=2e-12, method='DOP853'):
        times = np.array(times,dtype=float)
        if times[0] != 0 or np.any(np.diff(times) <= 0):
            raise ValueError('Increasing output times must start at zero.')
        solution = solve_ivp(self.rhs,(0,times[-1]),np.r_[self.feed,0,0,0,0],
                             t_eval=times,method=method,rtol=rtol,atol=1e-18)
        if not solution.success:
            raise RuntimeError(solution.message)
        if np.min(solution.y[:len(self.catalogue.words)]) < -1e-12:
            raise ArithmeticError('Materially negative concentration: refine integration.')
        return solution.y.T

    def stochastic(self, volume, seed, end=WINDOW_END, start=WINDOW_START,max_events=MAX_EVENTS):
        if (not isinstance(volume,int) or not 0 < volume < 10**12
                or not 0 <= start < end or not math.isfinite(end) or max_events < 1):
            raise ValueError('Positive integer volume and valid output window required.')
        c, rng = self.catalogue,np.random.default_rng(seed)
        counts = (self.feed*volume).astype(np.int64)
        food = sorted(c.food)
        time, peak = 0.0,int(c.lengths@counts)
        export_total = export_window = catalytic_input = basal_input = basal_positive = 0
        feed_mass = outflow_mass = 0
        trace = [(0.0,10.0,0.0,0.0)]
        completed = False
        for step in range(max_events):
            rates = np.r_[self.rates(counts,volume),np.full(len(food),volume),counts]
            total = rates.sum()
            waiting = rng.exponential(1/total)
            if time+waiting > end:
                time = end
                completed = True
                break
            time += waiting
            event = int(np.searchsorted(np.cumsum(rates),rng.random()*total,side='right'))
            if event < len(self.events):
                counts += self.stoich[:,event]
                delta = int(self.nf_changes[event])
                if self.is_catalytic[event]:
                    catalytic_input += delta
                else:
                    basal_input += delta
                    basal_positive += max(0,delta)
            elif event < len(self.events)+len(food):
                x = food[event-len(self.events)]
                counts[x] += 1
                feed_mass += int(c.lengths[x])
            else:
                x = event-len(self.events)-len(food)
                counts[x] -= 1
                outflow_mass += int(c.lengths[x])
                reward = int(c.nonfood[x])
                export_total += reward
                if time > start:
                    export_window += reward
            assert np.min(counts) >= 0
            mass = int(c.lengths@counts)
            peak = max(peak,mass)
            assert mass == 10*volume+feed_mass-outflow_mass
            assert int(c.nonfood@counts)+export_total == catalytic_input+basal_input
            if step % 100 == 0:
                trace.append((time,mass/volume,int(c.nonfood@counts)/volume,export_window/volume))
        trace.append((time,int(c.lengths@counts)/volume,int(c.nonfood@counts)/volume,export_window/volume))
        return {'completed':completed,'events':step if completed else max_events,'end_time':time,
                'volume':volume,'maximum_total_mass':peak/volume,'window_export':export_window/volume,
                'productive':bool(peak/volume <= MASS_CEILING and export_window/volume > EXPORT_THRESHOLD) if completed else None,
                'signed_catalytic_input':catalytic_input/volume,'signed_basal_input':basal_input/volume,
                'positive_basal_input':basal_positive/volume,'final_nonfood_mass':int(c.nonfood@counts)/volume,
                'total_export':export_total/volume,'integer_balances_exact':True,'trace':trace}