"""A shared polymer catalogue, catalytic source, and fed mass-action reactor.

Finite-copy runs and deterministic trajectories explore the literal model;
neither estimates the theorem's asymptotic success probability by default.
"""
from __future__ import annotations

from dataclasses import dataclass
from collections import Counter
from pathlib import Path
import argparse
import csv
import hashlib
import itertools
import json
import math
import platform
import mpmath as mp
import numpy as np
from scipy.integrate import solve_ivp

# EDITABLE INPUTS. Time is in dilution-time units; volume converts concentration
# to molecule counts. These are manuscript parameters, not measured constants.
MAX_WORD_LENGTH = 4
BASAL_EPSILON = 1 / 500_000_000
CATALYTIC_SCALE = 4.0
KINETIC_MARKS = (1.0, 1.5, 2.0)
MARK_SEED = 7601
TRAJECTORY_SEED = 20260922
STOCHASTIC_VOLUME = 20
MAX_EVENTS = 200_000
WINDOW_START, WINDOW_END = 1.0, 100.0
MASS_CEILING, EXPORT_THRESHOLD = 11.0, 0.1
OUTPUT_POINTS = 1001
SOURCE_LENGTHS = (4, 5, 6, 7, 8)
SOURCE_TRIALS = 48
SOURCE_SEED = 16092026
ANALYTIC_LENGTHS = (4, 8, 16, 32, 64, 128, 256, 1024)
MAX_CATALOGUE_SPECIES = 1022
MANUSCRIPT_SHA256 = '76c780846acad0324c5ac76cc00fb7d42ca5c5bf7060f9b9a2d13b59f75f5f3f'
# Specified host: product self-catalysis plus four nonfood assignments.
# Each entry is (catalyst, left reactant, right reactant).
HOST_ASSIGNMENTS = (('0011','00','11'), ('010','0','01'), ('101','1','10'),
                    ('011','0','11'), ('110','1','00'))


@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=HOST_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']])


class CappedZipfSource:
    def __init__(self, catalogue, exponent=None):
        self.catalogue = catalogue
        self.exponent = 2-2/catalogue.n if exponent is None else exponent
        if self.exponent <= 1:
            raise ValueError('Zipf exponent must exceed one.')
        R = len(catalogue.channels)
        with mp.workdps(70):
            a = mp.mpf(str(self.exponent))
            weights = [mp.power(k,-a)/mp.zeta(a) for k in range(1,R)]
            weights.append(mp.zeta(a,R)/mp.zeta(a))
            self.weights = np.array([float(v) for v in weights])
        self.weights /= self.weights.sum()  # only floating-point normalization

    def sample(self, rng, condition_on_witness=False):
        """Capped law including its tail atom, then uniform subsets per row.

        Witness conditioning uses a size-biased product-row degree; simply
        adding the required incidence to an ordinary sample has the wrong law.
        """
        c, R = self.catalogue,len(self.catalogue.channels)
        if condition_on_witness and c.n < 4:
            raise ValueError('Witness needs maximum word length >= 4.')
        degrees = rng.choice(R,size=len(c.words),p=self.weights)
        forced = None
        if condition_on_witness:
            degrees[list(c.food)] = 0
            forced = (c.index['0011'],c.channel_index['00','11'])
            biased = np.arange(R)*self.weights
            biased /= biased.sum()
            degrees[forced[0]] = rng.choice(R,p=biased)
        rows = []
        for x,degree in enumerate(degrees):
            if forced is not None and x == forced[0]:
                other = rng.choice(R-1,size=int(degree)-1,replace=False)
                row = {int(v)+(v >= forced[1]) for v in other}
                row.add(forced[1])
            else:
                row = set(rng.choice(R,size=int(degree),replace=False).tolist())
            rows.append(frozenset(row))
        return CatalyticEnvironment(c,tuple(rows))

    @staticmethod
    def statistics(n, digits=80):
        """No catalogue allocation; Hurwitz zeta includes the whole cap tail."""
        if n < 4:
            raise ValueError('Critical sequence reported for n >= 4.')
        with mp.workdps(digits):
            a = 2-mp.mpf(2)/n
            R, X = (n-2)*mp.power(2,n+1)+4,mp.power(2,n+1)-2
            mean = (mp.zeta(a-1)-mp.zeta(a-1,R)+R*mp.zeta(a,R))/mp.zeta(a)-1
            p,q0 = mean/R,1/mp.zeta(a)
            return {key:mp.nstr(value,35) for key,value in {
                'exponent':a,'mean_degree':mean,'mean_degree_over_n':mean/n,
                'X_times_p':X*p,'log10_incidence':mp.log10(p),
                'log10_witness_event':mp.log10(q0**6*p),'zero_degree_probability':q0,
                'cap_probability':mp.zeta(a,R)/mp.zeta(a)}.items()}


@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

    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'),
                                       ReactionEvent((product,x),(left,right,x),k,'catalytic')))
        self.events = tuple(events)
        count = len(c.words)
        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):
        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='DOP853',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}


class TheoremBounds:
    """Conservative finite inequalities; no huge catalogue is allocated."""
    @staticmethod
    def structural_budget(n):
        if not isinstance(n,int) or n < 4:
            raise ValueError('Integer n >= 4 required.')
        h = math.isqrt(math.isqrt(n**3))
        L = math.isqrt(64*n)
        L += L*L < 64*n
        with mp.workdps(80):
            # A_n=(n+1)*2^(L+1)+n^4*2^h; B_n=A_n^2.
            A = (n+1)*mp.power(2,L+1)+n**4*mp.power(2,h)
            log2B = 2*mp.log(A,2)
            return {'n':n,'h':h,'L':L,'log2_budget':mp.nstr(log2B,30),
                    'log2_budget_over_n':mp.nstr(log2B/n,30)}

    @staticmethod
    def evaluate(n, volume):
        with mp.workdps(80):
            d = mp.mpf(1)/(6*10**20)
            c = d*d/(4*10**7)
            if n < 4 or volume < 2*n/d:
                return {'applicable':False,'reason':'requires n >= 4 and V >= 2n/d'}
            stats = CappedZipfSource.statistics(n)
            logp = mp.mpf(stats['log10_incidence'])*mp.log(10)
            logwitness = mp.mpf(stats['log10_witness_event'])*mp.log(10)
            logdelta = -c*mp.mpf(volume)/n
            k = 4*10**8
            logC = ((k+1)*mp.log(2)+mp.log1p(-mp.power(2,-k))
                    +mp.log(k)+(k+3)*mp.log(2)+mp.log1p(4/(k*mp.power(2,k+3))))
            lower_factor = 1-24*mp.exp(logdelta)
            upper_log = max(logC+logp,mp.log(4)+logdelta)
            upper_log += mp.log1p(mp.exp(min(logC+logp,mp.log(4)+logdelta)-upper_log))
            return {'applicable':True,'log10_noise':mp.nstr(logdelta/mp.log(10),25),
                    'log10_positive_lower_bound':mp.nstr((logwitness+mp.log(lower_factor))/mp.log(10),25) if lower_factor > 0 else None,
                    'log10_upper_bound_capped_at_one':mp.nstr(min(0,upper_log)/mp.log(10),25),
                    'log10_disabled_upper_bound_capped_at_one':mp.nstr(min(0,mp.log(2)+logdelta)/mp.log(10),25),
                    'log10_catalogue_constant':mp.nstr(logC/mp.log(10),25)}


def write_csv(path, headers, rows):
    with path.open('w',newline='') as stream:
        writer = csv.writer(stream)
        writer.writerow(headers)
        writer.writerows(rows)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=Path('outputs'))
    args = parser.parse_args()
    out = args.output
    out.mkdir(parents=True,exist_ok=True)
    catalogue = PolymerCatalogue(MAX_WORD_LENGTH)
    environment = CatalyticEnvironment.specified(catalogue)
    parameters = KineticParameters.sample(environment,reproduce_host_marks=True)
    enabled,disabled = FedReactor(environment,parameters),FedReactor(environment,parameters,False)
    times = np.unique(np.r_[np.linspace(0,WINDOW_END,OUTPUT_POINTS),WINDOW_START])
    fine,coarse,off = enabled.deterministic(times),enabled.deterministic(times,1e-10),disabled.deterministic(times)
    N = len(catalogue.words)
    start = int(np.searchsorted(times,WINDOW_START))
    deterministic = {'specified_host_not_source_sample':True,'witness_event':environment.witness_event(),
        'max_raf_channels':len(environment.max_raf()),'irreducible_witness_channels':len(environment.irreducible_raf()),
        'enabled_export':float(fine[-1,N]-fine[start,N]),'disabled_export':float(off[-1,N]-off[start,N]),
        'tolerance_max_difference':float(np.max(np.abs(fine-coarse))),
        'maximum_mass_error':float(np.max(np.abs(fine[:,:N]@catalogue.lengths-10))),
        'nonfood_balance_error':float(np.max(np.abs(fine[:,:N]@catalogue.nonfood+fine[:,N]-fine[:,N+1]-fine[:,N+2])))}
    write_csv(out/'deterministic.csv',['time',*catalogue.words,'export','signed_catalytic','signed_basal','positive_basal'],zip(times,*fine.T))
    write_csv(out/'disabled_deterministic.csv',['time',*catalogue.words,'export','signed_catalytic','signed_basal','positive_basal'],zip(times,*off.T))
    stochastic = {}
    for name,model in [('enabled',enabled),('disabled',disabled)]:
        run = model.stochastic(STOCHASTIC_VOLUME,TRAJECTORY_SEED)
        write_csv(out/f'{name}_stochastic.csv',['time','total_mass','nonfood_mass','window_export'],run.pop('trace'))
        stochastic[name] = run
    rng = np.random.default_rng(SOURCE_SEED)
    census = []
    for n in SOURCE_LENGTHS:
        cat = PolymerCatalogue(n)
        source = CappedZipfSource(cat)
        witnesses = [source.sample(rng).max_raf() for _ in range(SOURCE_TRIALS)]
        census.append({'n':n,'trials':SOURCE_TRIALS,'raf_count':sum(bool(w) for w in witnesses),
                       'maximum_raf_sizes':[len(w) for w in witnesses]})
    analytic = [{'n':n,**CappedZipfSource.statistics(n),
                 'bounds_at_sufficient_volume':TheoremBounds.evaluate(n,10**60*(n+1)**2)} for n in ANALYTIC_LENGTHS]
    result = {'deterministic':deterministic,'finite_copy_runs':stochastic,'source_census':census,
              'structural_cutoff':[TheoremBounds.structural_budget(n) for n in (4,16,256,4096,65536,1000000)],
              'source_law':analytic,'bounds_at_simulated_volume':TheoremBounds.evaluate(MAX_WORD_LENGTH,STOCHASTIC_VOLUME)}
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    lines = [f'Specified host: {N} species, {len(catalogue.channels)} reversible channels; witness event = {environment.witness_event()}.',
             f'Deterministic enabled/disabled export: {deterministic["enabled_export"]:.8g} / {deterministic["disabled_export"]:.8g}.',
             f'ODE mass error: {deterministic["maximum_mass_error"]:.3g}; balance error: {deterministic["nonfood_balance_error"]:.3g}.',
             'Finite-copy runs are exploratory; the theorem volume hypothesis is NOT met.']
    lines += [f'{name}: completed={r["completed"]}, output event={r["productive"]}, export={r["window_export"]}.' for name,r in stochastic.items()]
    lines += [f'Source n={r["n"]}: {r["raf_count"]}/{r["trials"]} RAFs (finite sample; not the limiting survival probability).' for r in census]
    (out/'console.txt').write_text('\n'.join(lines)+'\n')
    print('\n'.join(lines))
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axes = plt.subplots(1,2,figsize=(11,4.4),layout='constrained')
    axes[0].plot(times,fine[:,catalogue.index['0011']],label='Self-catalysed product 0011')
    axes[0].plot(times,fine[:,:N]@catalogue.nonfood,label='Total nonfood mass')
    axes[0].plot(times,off[:,:N]@catalogue.nonfood,'--',label='Nonfood mass, disabled')
    axes[0].set(xlabel='Time (dilution units)',ylabel='Concentration / mass per volume',title='Concentrations in the specified polymer\nreactor')
    for col,label in [(N,'Export'),(N+1,'Signed catalytic input'),(N+3,'Positive basal input')]:
        axes[1].plot(times[start:],fine[start:,col]-fine[start,col],label=label)
    axes[1].plot(times[start:],off[start:,N]-off[start,N],'--',label='Export, disabled')
    axes[1].set(xlabel='Time (dilution units)',ylabel='Mass per volume since time 1',title='Production and collected output')
    for ax in axes:
        ax.legend(fontsize=8);ax.grid(alpha=.2)
    fig.savefig(out/'reactor.png',dpi=180);fig.savefig(out/'reactor.svg');plt.close(fig)
    fig,axes = plt.subplots(1,2,figsize=(11,4.4),layout='constrained')
    for key,label in [('X_times_p','X_n p_n'),('mean_degree_over_n','Mean degree / n')]:
        axes[0].plot(ANALYTIC_LENGTHS,[float(r[key]) for r in analytic],'.-',label=label)
    axes[0].axhline(9/math.pi**2,color='black',ls='--',label='Limit 9 / pi^2')
    axes[0].set(xscale='log',xlabel='Maximum word length n',ylabel='Normalized source intensity',title='Capped-Zipf source law')
    for key,label in [('log10_incidence','One specified incidence p_n'),('log10_witness_event','Witness environment q0^6 p_n')]:
        axes[1].plot(ANALYTIC_LENGTHS,[float(r[key]) for r in analytic],'.-',label=label)
    axes[1].set(xlabel='Maximum word length n',ylabel='log10 probability',title='Catalytic-assignment and witness\nprobabilities')
    for ax in axes:
        ax.legend(fontsize=8);ax.grid(alpha=.2)
    fig.savefig(out/'source.png',dpi=180);fig.savefig(out/'source.svg');plt.close(fig)
    digest = lambda p:hashlib.sha256(p.read_bytes()).hexdigest()
    (out/'run_metadata.json').write_text(json.dumps({'manuscript_sha256':MANUSCRIPT_SHA256,
        'source_sha256':digest(Path(__file__)),'python':platform.python_version(),
        'output_sha256':{p.name:digest(p) for p in sorted(out.iterdir()) if p.is_file() and p.name != 'run_metadata.json'}},indent=2)+'\n')


if __name__ == '__main__':
    main()
