"""Correlated catalytic rows, singleton candidates and productive-input attribution."""
# EDITABLE STUDY INPUTS. Concentrations/time use the paper's dimensionless units.
MAX_WORD_LENGTH = 4
LOCAL_INCIDENCE = ('0011','00','11')  # catalyst, left food, right food
BACKGROUND_INCIDENCES = ()          # changing this can leave local-uniqueness hypotheses
MARK_SEED = 30092026
STOCHASTIC_VOLUME = 20              # exploratory only; not the theorem envelope
STOCHASTIC_SEED = 22092026
MAX_EVENTS = 200000
ANALYTIC_LENGTHS = (4,8,16,32,64,128,256,1024)
FINITE_CANDIDATE_LENGTHS = (4,5,6,7,8)
MAX_DEGREE_TABLE = 10000
WINDOW_START, WINDOW_END = 1.0,100.0
MANUSCRIPT_SHA256 = 'a93db4a9bd559971dc0dcc2654e5c205bbdf684cd53d9b44db7c722db9e5efaf'

import argparse
from dataclasses import dataclass
from fractions import Fraction as Q
from functools import lru_cache
import csv
import hashlib
import itertools
import json
from pathlib import Path
import platform

import mpmath as mp
import numpy as np
from scipy.integrate import solve_ivp
from polymer import PolymerCatalogue, CatalyticEnvironment, KineticParameters, FedReactor

mp.mp.dps=90


def hurwitz_tail(s,R):
    """Numerical Euler--Maclaurin tail for huge R, avoiding a large-prefix scan.

    Twenty corrections are ample at R>1e6 and the exponents used here. This
    evaluator is high-precision numerics, not an outward-rounded certificate.
    """
    if R<=10**6:return mp.zeta(s,R)
    a=mp.mpf(R)
    return a**(1-s)/(s-1)+a**(-s)/2+mp.fsum(
        mp.bernoulli(2*k)/mp.factorial(2*k)*mp.rf(s,2*k-1)*a**(1-s-2*k) for k in range(1,21))


class ZipfLaw:
    """D=min(K,R)-1, retaining the entire Zipf tail at R-1."""
    def __init__(self,n):
        if not isinstance(n,int) or not 4<=n<=10000:raise ValueError('Analytic evaluator supports integer 4<=n<=10000.')
        self.n=n;self.R=(n-2)*2**(n+1)+4;self.X=2**(n+1)-2
        self.a=2-mp.mpf(2)/n;self.zeta=mp.zeta(self.a)
        tail=hurwitz_tail(self.a,self.R)
        H=lambda s:mp.zeta(s)-hurwitz_tail(s,self.R)
        self.mean=(H(self.a-1)-H(self.a)+(self.R-1)*tail)/self.zeta
        self.second=(H(self.a-2)-2*H(self.a-1)+H(self.a)+(self.R-1)**2*tail)/self.zeta
        self.p=self.mean/self.R;self.q2=(self.second-self.mean)/(self.R*(self.R-1));self.q0=1/self.zeta

    @lru_cache(maxsize=None)
    def weights(self,budget=MAX_DEGREE_TABLE):
        if self.R>budget:raise ValueError('Explicit degree table exceeds budget; use moments instead.')
        return tuple([mp.power(d+1,-self.a)/self.zeta for d in range(self.R-1)]+[mp.zeta(self.a,self.R)/self.zeta])

    @lru_cache(maxsize=None)
    def row_hit(self,J):
        if not 0<=J<=self.R:raise ValueError('Invalid channel count.')
        # Stable product for the exact without-replacement row-miss probability.
        result=mp.mpf(0)
        for d,w in enumerate(self.weights()):
            miss=mp.mpf(1)
            if d>self.R-J:miss=mp.mpf(0)
            else:
                for j in range(J):miss*=mp.mpf(self.R-d-j)/(self.R-j)
            result+=w*(1-miss)
        return result

    def candidate_union(self,group_sizes):
        """Exact finite source formula, numerically evaluated; groups are distinct rows."""
        log_miss=mp.fsum(mp.log1p(-self.row_hit(J)) for J in group_sizes)
        return -mp.expm1(log_miss)

    def rectangle_2_by_2(self):
        """Stable exact expressions for two rows and two columns, not the huge K* box."""
        p,q=self.p,self.q2;one=2*p-2*q;hit=2*p-q
        return {'actual_two_or_more':2*q-q*q+one*one,'excess_count_bound':2*q+hit*hit,
            'pair_count_bound':2*q+4*p*p,'first_moment_bound':4*p,
            'wrong_independent_pair_bound':6*p*p}

    def sample(self,catalogue,rng,conditioned_incidence=None,empty_food=False):
        if catalogue.n!=self.n:raise ValueError('Catalogue/source mismatch.')
        weights=np.array([float(w) for w in self.weights()]);weights/=weights.sum()
        degrees=rng.choice(self.R,size=len(catalogue.words),p=weights)
        if empty_food:degrees[list(catalogue.food)]=0
        forced=None
        if conditioned_incidence is not None:
            z,u,v=conditioned_incidence;forced=(catalogue.index[z],catalogue.channel_index[u,v])
            if empty_food and forced[0] in catalogue.food:raise ValueError('Conditioned food row cannot also be empty.')
            biased=np.arange(self.R)*weights;biased/=biased.sum();degrees[forced[0]]=rng.choice(self.R,p=biased)
        rows=[]
        for i,d in enumerate(degrees):
            if forced is not None and i==forced[0]:
                draw=rng.choice(self.R-1,int(d)-1,replace=False)
                row={int(k)+(k>=forced[1]) for k in draw}|{forced[1]}
            else:row=set(rng.choice(self.R,int(d),replace=False).tolist())
            rows.append(frozenset(row))
        return CatalyticEnvironment(catalogue,tuple(rows))


def candidates(catalogue,productive=True):
    result=[]
    for r,c in enumerate(catalogue.channels):
        if c.left in catalogue.food and c.right in catalogue.food and (not productive or catalogue.nonfood[c.product]>0):
            result.extend((z,r) for z in sorted(catalogue.food|{c.product}))
    return tuple(result)


def local_classification(catalogue,z,r):
    c=catalogue.channels[r];m=catalogue.nonfood;gain=int(m[c.product]-m[c.left]-m[c.right])
    if gain==0:kind='neutral'
    elif c.left not in catalogue.food or c.right not in catalogue.food:kind='mixed'
    elif z in catalogue.food or z==c.product:kind='productive'
    else:kind='outsider'
    credit=m.astype(np.int64).copy()
    if kind=='mixed':credit[c.product]-=gain
    if kind=='outsider':credit[z]+=2000
    return kind,gain,credit


def erase_incidence(environment,incidence):
    z,u,v=incidence;c=environment.catalogue;zi=c.index[z];ri=c.channel_index[u,v]
    if ri not in environment.rows[zi]:raise ValueError('Incidence to erase is absent.')
    rows=list(environment.rows);rows[zi]=rows[zi]-{ri}
    return CatalyticEnvironment(c,tuple(rows))


class AttributedReactor(FedReactor):
    """Same chemical kernel; six integrated rewards identify signed local input."""
    ledger=('total_export','signed_local','signed_other','signed_basal','positive_basal','positive_other')

    def __init__(self,environment,parameters,local_incidence):
        super().__init__(environment,parameters)
        c=self.catalogue;z,u,v=local_incidence;key=(c.channel_index[u,v],c.index[z])
        self.local=np.array([e.incidence==key for e in self.events])
        self.other=self.is_catalytic&~self.local

    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,signed[self.local].sum(),signed[self.other].sum(),
            signed[~self.is_catalytic].sum(),np.maximum(self.nf_changes[~self.is_catalytic],0)@rates[~self.is_catalytic],
            np.maximum(self.nf_changes[self.other],0)@rates[self.other]]

    def deterministic(self,times,method='DOP853'):
        times=np.array(times,float)
        if times[0]!=0 or np.any(np.diff(times)<=0):raise ValueError('Increasing times starting at zero required.')
        solution=solve_ivp(self.rhs,(0,times[-1]),np.r_[self.feed,np.zeros(6)],t_eval=times,method=method,rtol=2e-11,atol=1e-18)
        if not solution.success:raise RuntimeError(solution.message)
        if np.min(solution.y[:len(self.catalogue.words)]) < -1e-11:raise ArithmeticError('Negative concentration; refine numerical solution.')
        return solution.y.T


@dataclass(frozen=True)
class ConditioningLaw:
    """Reweight a supplied finite environment law; reliabilities must be supplied."""
    weights: tuple
    success: tuple

    def __post_init__(self):
        if len(self.weights)!=len(self.success) or sum(self.weights)!=1 or min(self.weights,default=-1)<0:
            raise ValueError('A normalized finite environment law is required.')
        if any(not 0<=p<=1 for p in self.success):raise ValueError('Success probabilities must be in [0,1].')

    def successful_run(self):
        total=sum(w*p for w,p in zip(self.weights,self.success))
        if total==0:return None
        return tuple(w*p/total for w,p in zip(self.weights,self.success))

    def reliable(self,eta):
        if not 0<eta<1:raise ValueError('Fixed reliability threshold must lie in (0,1).')
        selected=tuple(w if p>=eta else 0 for w,p in zip(self.weights,self.success));total=sum(selected)
        return tuple(w/total for w in selected) if total else None

    def structural(self,flags):
        if len(flags)!=len(self.weights):raise ValueError('One structural flag per environment required.')
        selected=tuple(w if flag else 0 for w,flag in zip(self.weights,flags));total=sum(selected)
        return tuple(w/total for w in selected) if total else None


class ProofBudgets:
    """Exact rational kinetic budgets, with explicit theorem-scale applicability."""
    epsilon=Q(1,500000000);K=10**12;k0=4*10**8;d=Q(1,6*10**20);c=Q(1,144*10**47)

    @classmethod
    def algebra(cls):
        b=1936*cls.epsilon;t=Q(85184,cls.K);loss=1012*cls.epsilon+Q(44528,cls.K)
        creation=Q(550,3)*cls.epsilon+Q(29040,cls.K)
        mixed=(b+t+2*loss)/2;outsider=100*(b+t+2000*creation)+Q(1001,100000)
        other=100*(b+t)+Q(2,100)
        assert mixed<Q(1,200000) and outsider<Q(1,10) and other<Q(1,40)
        return {'mixed_half_drift':str(mixed),'mixed_noise_plus_drift':str(Q(21,2000)),
            'mixed_export_contradiction_threshold':str(Q(3,100)),'outsider_export_budget':str(outsider),
            'outsider_export_threshold':str(Q(1,10)),'other_input_budget':str(other),
            'other_input_limit':str(Q(1,40)),'washout_margin':2000-1936}

    @classmethod
    def evaluate(cls,n,V):
        if not isinstance(V,int) or V<=0:raise ValueError('Positive integer volume required.')
        law=ZipfLaw(n)
        if V<2*n/cls.d or Q(n,V)>Q(1,200000):
            return {'applicable':False,'reason':'Requires V>=2n/d and n/V<=1/200000; small simulations do not satisfy these.'}
        exponent=mp.mpf(cls.c.numerator)*V/(cls.c.denominator*n)
        Z=mp.mpf(2)**(min(n,cls.K)+1)-2
        jn=min(n,cls.K+2);J=(jn-2)*mp.mpf(2)**(jn+1)+4
        # Fixed K* rectangle, using actual finite-n row/column counts. Never enumerate it.
        pair=Z*J*(J-1)/2*law.q2+Z*(Z-1)/2*J**2*law.p**2
        multiplicity=min(mp.mpf(1),pair)
        log_lower=6*mp.log(law.q0)+mp.log(law.p)+mp.log1p(-24*mp.exp(-exponent)) if exponent>mp.log(24) else None
        posterior=min(mp.mpf(1),(8*mp.exp(-exponent)+multiplicity)/mp.exp(log_lower)) if log_lower is not None else None
        return {'applicable':True,'at_sufficient_quadratic_envelope':V>=10**60*(n+1)**2,
            'log10_deletion_probability_upper':mp.nstr(min(0,mp.log(4)-exponent)/mp.log(10),30),
            'log10_success_lower':mp.nstr(log_lower/mp.log(10),30) if log_lower is not None else None,
            'fixed_Kstar_multiplicity_upper':mp.nstr(multiplicity,30),
            'posterior_no_productive_candidate_upper':mp.nstr(posterior,30) if posterior is not None else None,
            'scope':'Evaluated theorem inequalities, not interval-rounded probabilities. The huge fixed rectangle makes this finite-n posterior bound vacuous in the displayed range.'}


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'));parser.add_argument('--simulate',action='store_true');args=parser.parse_args()
    out=args.output;out.mkdir(parents=True,exist_ok=True);catalogue=PolymerCatalogue(MAX_WORD_LENGTH)
    environment=CatalyticEnvironment.specified(catalogue,(LOCAL_INCIDENCE,*BACKGROUND_INCIDENCES));parameters=KineticParameters.sample(environment,seed=MARK_SEED)
    reactor=AttributedReactor(environment,parameters,LOCAL_INCIDENCE)
    deleted_environment=erase_incidence(environment,LOCAL_INCIDENCE);deleted=AttributedReactor(deleted_environment,parameters,LOCAL_INCIDENCE)
    times=np.unique(np.r_[np.linspace(0,100,1001),1]);enabled=reactor.deterministic(times);off=deleted.deterministic(times);independent=reactor.deterministic(times,method='Radau')
    N=len(catalogue.words);start=int(np.searchsorted(times,1));export=enabled[-1,N]-enabled[start,N]
    result={'specified_environment_not_posterior_sample':True,'productive_candidates':len(candidates(catalogue)),
        'singleton_witness_candidates':len(candidates(catalogue,False)),
        'food_catalyst_candidates':sum(z in catalogue.food for z,r in candidates(catalogue)),
        'product_catalyst_candidates':sum(z==catalogue.channels[r].product for z,r in candidates(catalogue)),
        'deterministic':{'window_export':float(export),'deleted_window_export':float(off[-1,N]-off[start,N]),
            'signed_local_input_0_to_100':float(enabled[-1,N+1]),'local_input_over_window_export':float(enabled[-1,N+1]/export),
            'nonfood_balance_error':float(np.max(np.abs(enabled[:,:N]@catalogue.nonfood+enabled[:,N]-enabled[:,N+1:N+4].sum(axis=1)))),
            'mass_error':float(np.max(np.abs(enabled[:,:N]@catalogue.lengths-10))),
            'independent_solver_max_difference':float(np.max(np.abs(enabled-independent)))},
        'proof_budgets':ProofBudgets.algebra(),'bounds_at_simulation_volume':ProofBudgets.evaluate(MAX_WORD_LENGTH,STOCHASTIC_VOLUME)}
    rows=[]
    for n in ANALYTIC_LENGTHS:
        law=ZipfLaw(n);rectangle=law.rectangle_2_by_2()
        rows.append({'n':n,'X_times_p':mp.nstr(law.X*law.p,35),'q2_over_p':mp.nstr(law.q2/law.p,35),
            'log10_q2_over_p_squared':mp.nstr(mp.log10(law.q2/law.p**2),35),
            'mean_degree':mp.nstr(law.mean,35),'conditional_mean_degree':mp.nstr(law.second/law.mean,35),
            **{k+'_over_p':mp.nstr(v/law.p,35) for k,v in rectangle.items()},
            'theorem_bounds':ProofBudgets.evaluate(n,10**60*(n+1)**2)})
    result['source_moments']=rows;finite=[]
    for n in FINITE_CANDIDATE_LENGTHS:
        law=ZipfLaw(n);cat=PolymerCatalogue(n)
        groups=lambda incidences:[sum(z==i for z,r in incidences) for i in sorted({z for z,r in incidences})]
        probability=law.candidate_union(groups(candidates(cat)));singleton=law.candidate_union(groups(candidates(cat,False)))
        finite.append({'n':n,'p':mp.nstr(law.p,35),'productive_union_over_p':mp.nstr(probability/law.p,35),
            'singleton_union_over_p':mp.nstr(singleton/law.p,35),'cap_atom':mp.nstr(law.weights()[-1],35)})
    result['finite_candidate_probabilities']=finite
    if args.simulate:
        result['exploratory_count_paths']={name:model.stochastic(STOCHASTIC_VOLUME,STOCHASTIC_SEED,max_events=MAX_EVENTS) for name,model in [('original',reactor),('deleted',deleted)]}
        result['count_path_scope']='Small-volume original kernel. Completed=False has productive=None. Catalytic ledger sums all incidences; density ledger isolates the named incidence.'
    (out/'results.json').write_text(json.dumps(result,indent=2)+'\n')
    def csv_file(name,headers,data):
        with (out/name).open('w',newline='') as f:
            w=csv.writer(f);w.writerow(headers);w.writerows(data)
    csv_file('deterministic_original.csv',['time',*catalogue.words,*reactor.ledger],[[t,*values] for t,values in zip(times,enabled)])
    csv_file('deterministic_deleted.csv',['time',*catalogue.words,*reactor.ledger],[[t,*values] for t,values in zip(times,off)])
    keys=[k for k in rows[0] if k!='theorem_bounds'];csv_file('source_moments.csv',keys,[[r[k] for k in keys] for r in rows])
    keys=list(finite[0]);csv_file('candidate_probabilities.csv',keys,[[r[k] for k in keys] for r in finite])
    csv_file('candidate_incidences.csv',['catalyst','left','right','product','productive_type'],
        [[catalogue.words[z],catalogue.words[catalogue.channels[r].left],catalogue.words[catalogue.channels[r].right],catalogue.words[catalogue.channels[r].product],int(catalogue.nonfood[catalogue.channels[r].product]>0)] for z,r in candidates(catalogue,False)])
    lines=[f'Candidate incidences: {result["productive_candidates"]} productive; {result["singleton_witness_candidates"]} singleton witnesses.',
        f'Productive split: {result["food_catalyst_candidates"]} food-catalyzed, {result["product_catalyst_candidates"]} product-catalyzed; not posterior proportions.',
        f'Specified ODE export: {export:.9g}; deleting only the named incidence: {result["deterministic"]["deleted_window_export"]:.9g}.',
        f'Named signed input / window export: {result["deterministic"]["local_input_over_window_export"]:.9g}; numerical illustration, not a conditional limit.',
        'Correlated-row source formulas include the full cap tail and the incidence-conditioned size bias.',
        'The 2x2 rectangle is a small source-law demonstration. The K*=1e12 posterior bound remains vacuous at these finite n.']
    (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=(10,4.2),layout='constrained')
    for key,label in [('actual_two_or_more_over_p','Actual P(N >= 2) / p'),('pair_count_bound_over_p','Pair-count upper bound / p'),('first_moment_bound_over_p','First-moment bound / p')]:
        axes[0].plot(ANALYTIC_LENGTHS,[float(r[key]) for r in rows],'.-',label=label)
    axes[0].set(xscale='log',yscale='log',xlabel='Maximum word length n',ylabel='Probability or bound divided by incidence p',title='Fixed two-row, two-column rectangle');axes[0].legend(fontsize=8)
    axes[1].plot(ANALYTIC_LENGTHS,[float(r['log10_q2_over_p_squared']) for r in rows],'.-',color='#b95024')
    axes[1].set(xscale='log',xlabel='Maximum word length n',ylabel='log10(same-row pair probability / p squared)',title='Correlation between one catalyst’s channel\nassignments')
    for ax in axes:ax.grid(alpha=.2)
    fig.savefig(out/'correlated_rows.png',dpi=180);fig.savefig(out/'correlated_rows.svg');plt.close(fig)
    fig,axes=plt.subplots(1,2,figsize=(10,4.2),layout='constrained')
    axes[0].plot(times,enabled[:,:N]@catalogue.nonfood,label='Original nonfood mass');axes[0].plot(times,off[:,:N]@catalogue.nonfood,'--',label='Named incidence deleted')
    axes[0].set(xlabel='Time in dilution units',ylabel='Nonfood mass / volume',title='Specified density-limit reactor');axes[0].legend(fontsize=8)
    axes[1].plot(times[start:],enabled[start:,N]-enabled[start,N],label='Collected export since time 1')
    axes[1].plot(times[start:],enabled[start:,N+1],label='Signed local input since time 0')
    axes[1].plot(times[start:],enabled[start:,N+4]+enabled[start:,N+5],label='Other positive input since time 0')
    axes[1].set(xlabel='Time in dilution units',ylabel='Mass / volume',title='Cumulative signed input and collected\nexport');axes[1].legend(fontsize=8)
    for ax in axes:ax.grid(alpha=.2)
    fig.savefig(out/'catalytic_deletion.png',dpi=180);fig.savefig(out/'catalytic_deletion.svg');plt.close(fig)
    digest=lambda p:hashlib.sha256(p.read_bytes()).hexdigest();base=Path(__file__).resolve().parent
    (out/'run_metadata.json').write_text(json.dumps({'manuscript_sha256':MANUSCRIPT_SHA256,'source_sha256':digest(Path(__file__)),
        'input_sha256':{'polymer.py':digest(base/'polymer.py')},'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()