"""Checked ranked witnesses, local deletion queries, portfolios and robustness.

All queries are independent deletions from a fixed baseline. No kinetics implied.
"""
from __future__ import annotations
import argparse
import csv
from dataclasses import dataclass
from fractions import Fraction as F
from itertools import combinations,product
import hashlib
import json
from math import comb
from pathlib import Path
import platform
import time

# USER INPUTS ---------------------------------------------------------------
PAIR_COUNT = 3                   # two alternative producers per required input
DELETION = (0,3)                  # mixed deletion defeats complementary witnesses
GREEDY_PORTFOLIO_BUDGET = 2        # trained on every singleton, queried on any deletion
RETENTION_GRID = tuple(F(i,50) for i in range(51))
GATEWAY_SIZE = 12
FUNCTIONAL_PARENTS = (0,0,1,2,3)
CATALYST_ADDITION = (2,2)          # add producer 2 as catalyst of reaction 2
ADDITION_RETENTION = F(4,5)
EXHAUSTIVE_REACTION_CAP = 14
WITNESS_POOL_CAP = 4096
PAPER_SHA256 = '6482ca887a9fa90c1a44bf1950c6a89d08f45b2390bfc7ed0ed7b726fc9a2999'
# Inputs are structural incidence choices and independent availability probabilities.
# --------------------------------------------------------------------------


def subsets(values):
    values=tuple(values)
    for size in range(len(values)+1):
        for s in combinations(values,size): yield frozenset(s)


@dataclass(frozen=True)
class Reaction:
    identifier: int
    reactants: frozenset[str]
    products: frozenset[str]
    catalysts: frozenset[str]

    def __post_init__(self):
        for name in ('reactants','products','catalysts'): object.__setattr__(self,name,frozenset(getattr(self,name)))


class ReactionSystem:
    def __init__(self,reactions,food=('f',)):
        reactions=tuple(reactions);self.reactions={r.identifier:r for r in reactions};self.food=frozenset(food)
        if len(self.reactions)!=len(reactions): raise ValueError('Reaction identifiers must be distinct.')
        if any(not isinstance(r.identifier,int) for r in reactions): raise ValueError('Reaction identifiers must be integers.')
        self.identifiers=frozenset(self.reactions)

    def selected(self,ids):
        ids=frozenset(ids)
        if not ids<=self.identifiers: raise ValueError('Unknown reaction identifier.')
        return ids

    def closure(self,ids,food=None):
        ids=self.selected(ids);pool=set(self.food if food is None else food);schedule=[];pending=set(ids)
        while True:
            enabled=sorted(r for r in pending if self.reactions[r].reactants<=pool)
            if not enabled: return frozenset(pool),tuple(schedule)
            for r in enabled: pool.update(self.reactions[r].products);pending.remove(r);schedule.append(r)

    def prune(self,ids,pool):
        return frozenset(r for r in ids if self.reactions[r].reactants<=pool and self.reactions[r].catalysts&pool)

    def maximum(self,ids=None,food=None):
        ids=self.identifiers if ids is None else self.selected(ids)
        while True:
            pool,_=self.closure(ids,food);kept=self.prune(ids,pool)
            if kept==ids: return ids
            ids=kept

    def is_raf(self,ids):
        ids=self.selected(ids);pool,_=self.closure(ids)
        return bool(ids) and self.prune(ids,pool)==ids

    def pruning_certificate(self,ids,food):
        ids=self.selected(ids);rounds=[]
        while True:
            pool,schedule=self.closure(ids,food);rounds.append(dict(available=sorted(ids),schedule=list(schedule)))
            kept=self.prune(ids,pool)
            if kept==ids: return rounds
            ids=kept

    def check_certificate(self,ids,food,rounds):
        ids=self.selected(ids)
        for index,round_data in enumerate(rounds):
            if set(round_data['available'])!=ids: raise ValueError('Pruning round availability mismatch.')
            pool=set(food)
            for r in round_data['schedule']:
                if r not in ids or not self.reactions[r].reactants<=pool: raise ValueError('Invalid closure schedule.')
                pool.update(self.reactions[r].products)
            if any(self.reactions[r].reactants<=pool and not self.reactions[r].products<=pool for r in ids):
                raise ValueError('Proposed closure is not closed under the available reactions.')
            kept=self.prune(ids,pool)
            if kept==ids:
                if index!=len(rounds)-1: raise ValueError('Extra certificate rounds after fixed point.')
                return kept
            ids=kept
        raise ValueError('Certificate has no terminal fixed point.')


class RankedWitness:
    def __init__(self,baseline,parents,ranks):
        self.baseline=frozenset(baseline);self.parents={r:frozenset(p) for r,p in parents.items()};self.ranks=dict(ranks)

    def validate(self,system):
        s=self.baseline
        if not s<=system.identifiers or set(self.parents)!=s or set(self.ranks)!=s: raise ValueError('Witness domain mismatch.')
        for r in s:
            if not isinstance(self.ranks[r],int) or self.ranks[r]<0 or not self.parents[r]<=s: raise ValueError('Invalid rank or parent.')
            reaction=system.reactions[r]
            for x in reaction.reactants-system.food:
                if not any(x in system.reactions[t].products and self.ranks[t]<self.ranks[r] for t in self.parents[r]):
                    raise ValueError('Nonfood reactant lacks a lower-ranked producer.')
            if not reaction.catalysts&system.food and not any(reaction.catalysts&system.reactions[t].products for t in self.parents[r]):
                raise ValueError('Reaction lacks a food or selected product catalyst.')
        return True

    @classmethod
    def build(cls,system,baseline=None,retain_first=()):
        s=system.maximum() if baseline is None else system.selected(baseline);retained=frozenset(retain_first)
        if not retained<=s or system.maximum(retained)!=retained or system.maximum(s)!=s:
            raise ValueError('Construction requires fixed baseline and retained sets.')
        pool=set(system.food);order=[]
        for block in (retained,s-retained):
            pending=set(block)
            while pending:
                enabled=sorted(r for r in pending if system.reactions[r].reactants<=pool)
                if not enabled: raise ValueError('Selected set cannot be generated from food.')
                for r in enabled: order.append(r);pool.update(system.reactions[r].products);pending.remove(r)
        ranks={r:i for i,r in enumerate(order)};parents={r:set() for r in s}
        for r in s:
            eligible=retained if r in retained else s
            for x in sorted(system.reactions[r].reactants-system.food):
                candidates=[t for t in eligible if ranks[t]<ranks[r] and x in system.reactions[t].products]
                parents[r].add(min(candidates))
            if not system.reactions[r].catalysts&system.food:
                candidates=[t for t in eligible if system.reactions[r].catalysts&system.reactions[t].products]
                parents[r].add(min(candidates))
        witness=cls(s,parents,ranks);witness.validate(system);return witness

    def cone(self,deletions):
        reached=set(deletions)&self.baseline;pending=list(reached);children={r:set() for r in self.baseline}
        for r,parents in self.parents.items():
            for p in parents: children[p].add(r)
        while pending:
            r=pending.pop()
            for child in children[r]-reached: reached.add(child);pending.append(child)
        return frozenset(reached)

    def record(self): return dict(parents={str(r):sorted(p) for r,p in self.parents.items()},ranks=self.ranks)


class DeletionEngine:
    def __init__(self,system,witnesses=(),availability=None):
        self.system=system;self.baseline=system.maximum(availability)
        self.witnesses=tuple(RankedWitness(w.baseline,w.parents,w.ranks) for w in witnesses)
        for witness in self.witnesses:
            if witness.baseline!=self.baseline: raise ValueError('Witness baseline does not match engine baseline.')
            witness.validate(system)

    def query(self,deletions,proposed_certificate=None):
        deleted=self.system.selected(deletions)&self.baseline;region=self.baseline
        for witness in self.witnesses: region&=witness.cone(deleted)
        retained=self.baseline-region;local=region-deleted;fallback=False;certificate=[]
        if region==deleted: surviving=retained;route='closed seed'
        else:
            food=self.system.food.union(*(self.system.reactions[r].products for r in retained))
            certificate=self.system.pruning_certificate(local,food) if proposed_certificate is None else proposed_certificate
            try:
                surviving=retained|self.system.check_certificate(local,food,certificate);route='checked residual'
            except (ValueError,KeyError,TypeError):
                surviving=self.system.maximum(self.baseline-deleted);fallback=True;route='fresh fallback'
        return dict(deleted=sorted(deleted),region=sorted(region),retained=sorted(retained),surviving=sorted(surviving),
            loss=sorted(self.baseline-surviving),excess=len(region-(self.baseline-surviving)),route=route,fallback=fallback,certificate=certificate)


def greedy_portfolio(pool,queries,budget):
    if not pool or not isinstance(budget,int) or budget<1: raise ValueError('Nonempty witness pool and positive budget required.')
    queries=tuple(frozenset(k) for k in queries);events=[{(i,r) for i,k in enumerate(queries) for r in w.baseline-w.cone(k)} for w in pool]
    selected=[];covered=set()
    for _ in range(min(budget,len(pool))):
        i=max((i for i in range(len(pool)) if i not in selected),key=lambda i:len(events[i]-covered))
        selected.append(i);covered|=events[i]
    return tuple(pool[i] for i in selected),dict(pool_indices=selected,rescued_training_events=len(covered))


def paired_source(k):
    if not isinstance(k,int) or k<1 or 2**k>WITNESS_POOL_CAP: raise ValueError('Pair count is positive and must fit the complete witness budget.')
    reactions=[Reaction(2*i+j,{'f'},{f'x{i}',f'private{2*i+j}'},{'f'}) for i in range(k) for j in (0,1)]
    reactions.append(Reaction(2*k,{f'x{i}' for i in range(k)},{'z'},{'f'}));system=ReactionSystem(reactions);pool=[]
    for choice in product((0,1),repeat=k):
        parents={r:set() for r in system.identifiers};parents[2*k]={2*i+choice[i] for i in range(k)}
        witness=RankedWitness(system.identifiers,parents,{r:int(r==2*k) for r in system.identifiers});witness.validate(system);pool.append(witness)
    return system,tuple(pool)


class RobustnessPolynomial:
    """Exact availability enumeration; coefficient order is ascending powers of p."""
    def __init__(self,system):
        self.system=system;self.baseline=system.maximum();n=len(self.baseline)
        if n>EXHAUSTIVE_REACTION_CAP: raise ValueError('Availability enumeration budget exceeded.')
        self.mean=[0]*(n+1);self.second=[0]*(n+1);self.nonempty=[0]*(n+1)
        for available in subsets(sorted(self.baseline)):
            size=len(system.maximum(available));k=len(available)
            for j in range(n-k+1):
                coefficient=(-1)**j*comb(n-k,j)
                self.mean[k+j]+=size*coefficient;self.second[k+j]+=size*size*coefficient;self.nonempty[k+j]+=bool(size)*coefficient

    @staticmethod
    def evaluate(coefficients,p):
        p=F(p);result=F(0)
        for coefficient in reversed(coefficients): result=result*p+coefficient
        return result

    def moments(self,p):
        p=F(p)
        if not 0<=p<=1: raise ValueError('Retention probability must be between zero and one.')
        mean=self.evaluate(self.mean,p)
        return dict(mean=mean,nonempty=self.evaluate(self.nonempty,p),variance=self.evaluate(self.second,p)-mean*mean)

    def sensitivity(self):
        s=self.baseline;losses={r:s-self.system.maximum(s-{r}) for r in s};pairs=[]
        for i,j in combinations(sorted(s),2):
            joint=s-self.system.maximum(s-{i,j});overlap=losses[i]&losses[j];cooperative=joint-(losses[i]|losses[j])
            pairs.append(dict(first=i,second=j,overlap=len(overlap),cooperative=len(cooperative)))
        first=sum(map(len,losses.values()));second=2*sum(r['overlap']-r['cooperative'] for r in pairs)
        if first!=sum(i*c for i,c in enumerate(self.mean)) or second!=sum(i*(i-1)*c for i,c in enumerate(self.mean)):
            raise ArithmeticError('Loss-derived derivatives differ from the independent polynomial.')
        return dict(first_derivative=first,second_derivative=second,singleton_losses={str(r):sorted(v) for r,v in losses.items()},pair_terms=pairs)


class FunctionalSource:
    def __init__(self,parents):
        self.parents=tuple(parents);n=len(parents)
        if n<1 or any(not isinstance(r,int) or not 0<=r<n for r in parents): raise ValueError('Parent map must stay in a nonempty reaction set.')
        self.exposures=tuple(self.orbit(r) for r in range(n))

    def orbit(self,r):
        visited=set()
        while r not in visited: visited.add(r);r=self.parents[r]
        return frozenset(visited)

    def network(self,alternative=None):
        if alternative is not None and len(alternative.parents)!=len(self.parents): raise ValueError('Alternative map has wrong size.')
        return ReactionSystem(Reaction(r,{'f'},{f'x{r}'},{f'x{self.parents[r]}'}|({f'x{alternative.parents[r]}'} if alternative else set())) for r in range(len(self.parents)))

    def moments(self,p):
        p=F(p)
        if not 0<=p<=1: raise ValueError('Retention probability must be between zero and one.')
        exposures=self.exposures;mean=sum((p**len(e) for e in exposures),F(0))
        variance=sum((p**len(e|g)-p**(len(e)+len(g)) for e in exposures for g in exposures),F(0))
        return dict(mean=mean,variance=variance)

    def addition(self,site,parent):
        if not 0<=site<len(self.parents) or not 0<=parent<len(self.parents): raise ValueError('Invalid catalyst-addition index.')
        other=list(self.parents);other[site]=parent;return OneSiteExposure(self,FunctionalSource(other))


class OneSiteExposure:
    def __init__(self,original,alternative):
        if len(original.parents)!=len(alternative.parents) or sum(a!=b for a,b in zip(original.parents,alternative.parents))>1:
            raise ValueError('Two-exposure law requires maps differing at at most one site.')
        self.original,self.alternative=original,alternative

    def surviving(self,available):
        available=frozenset(available)
        if not available<=set(range(len(self.original.parents))): raise ValueError('Unknown available reaction.')
        return frozenset(r for r,(e,g) in enumerate(zip(self.original.exposures,self.alternative.exposures)) if e<=available or g<=available)

    def probability(self,target,p):
        p=F(p)
        if not isinstance(target,int) or not 0<=target<len(self.original.parents): raise ValueError('Unknown target reaction.')
        if not 0<=p<=1: raise ValueError('Retention probability must be between zero and one.')
        e,g=self.original.exposures[target],self.alternative.exposures[target]
        return p**len(e)+p**len(g)-p**len(e|g)

    def gain(self,target,p):
        return self.probability(target,p)-F(p)**len(self.original.exposures[target])

    def cuts(self,target):
        if not isinstance(target,int) or not 0<=target<len(self.original.parents): raise ValueError('Unknown target reaction.')
        e,g=self.original.exposures[target],self.alternative.exposures[target]
        return frozenset([*(frozenset({r}) for r in (e&g)-{target}),*(frozenset({r,s}) for r in e-g for s in g-e)])


def catalyst_candidates(source,p):
    rows=[]
    for site in range(len(source.parents)):
        for parent in range(len(source.parents)):
            addition=source.addition(site,parent)
            rows.append(dict(site=site,parent=parent,total_gain=str(sum((addition.gain(r,p) for r in range(len(source.parents))),F(0)))))
    return sorted(rows,key=lambda r:(-F(r['total_gain']),r['site'],r['parent']))


def write_csv(path,rows):
    with path.open('w',newline='',encoding='utf-8') as f:
        writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)


def plot(profiles,curves,output):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    import numpy as np
    with plt.rc_context({'font.size':11,'figure.facecolor':'white','savefig.facecolor':'white','svg.fonttype':'none'}):
        keys=list(dict.fromkeys(r['deleted'] for r in profiles));x=np.arange(len(keys));fig,ax=plt.subplots(figsize=(8,4.5),layout='constrained')
        for i,(strategy,color) in enumerate([('single','#0072B2'),('two','#D55E00'),('full','#009E73')]):
            rows=[r for r in profiles if r['strategy']==strategy]
            ax.bar(x+(i-1)*.23,[r['region_size'] for r in rows],width=.23,label=f'{strategy} witness region',color=color)
        rows=[r for r in profiles if r['strategy']=='single'];ax.scatter(x,[r['loss_size'] for r in rows],marker='_',s=130,color='black',label='Actual loss')
        ax.set_xticks(x,keys);ax.set(xlabel='Deleted reaction identifiers',ylabel='Reactions in region or actual loss')
        ax.legend(fontsize=8);ax.spines[['top','right']].set_visible(False)
        fig.savefig(output/'regions.png',dpi=220);fig.savefig(output/'regions.svg');plt.close(fig)
        fig,axes=plt.subplots(1,2,figsize=(9,4.2),layout='constrained')
        for name,color,style in [('cooperative','#0072B2','-'),('cycle','#D55E00','--'),('gateway','#009E73','-.')]:
            rows=[r for r in curves if r['source']==name];p=[float(F(r['p'])) for r in rows]
            axes[0].plot(p,[float(F(r['mean_fraction'])) for r in rows],style,color=color,label=name)
            axes[1].plot(p,[float(F(r['nonempty'])) for r in rows],style,color=color,label=name)
        for ax in axes:
            ax.set(xlabel='Independent reaction retention p',xlim=(0,1),ylim=(0,1));ax.spines[['top','right']].set_visible(False);ax.grid(color='#e3e6e8');ax.legend(fontsize=8)
        axes[0].set_ylabel('Expected surviving fraction');axes[1].set_ylabel('Probability of any RAF')
        fig.savefig(output/'robustness.png',dpi=220);fig.savefig(output/'robustness.svg');plt.close(fig)


def main():
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--output',type=Path,default=Path('outputs'))
    args=parser.parse_args();args.output.mkdir(parents=True,exist_ok=True);start=time.perf_counter()
    inputs=dict(pair_count=PAIR_COUNT,deletion=DELETION,greedy_budget=GREEDY_PORTFOLIO_BUDGET,retention_grid=list(map(str,RETENTION_GRID)),
        gateway_size=GATEWAY_SIZE,functional_parents=FUNCTIONAL_PARENTS,catalyst_addition=CATALYST_ADDITION,
        addition_retention=str(ADDITION_RETENTION),exhaustive_reaction_cap=EXHAUSTIVE_REACTION_CAP,witness_pool_cap=WITNESS_POOL_CAP)
    intro='Resolved inputs: '+json.dumps(inputs);print(intro,flush=True)
    system,pool=paired_source(PAIR_COUNT);s=system.identifiers
    if len(s)>EXHAUSTIVE_REACTION_CAP: raise ValueError('Default complete query audit exceeds reaction budget.')
    queries=[{r} for r in sorted(s)]+[set(DELETION)];strategies={'single':pool[:1],'two':(pool[0],pool[-1]),'full':pool}
    selected,selection=greedy_portfolio(pool,({r} for r in sorted(s)),GREEDY_PORTFOLIO_BUDGET)
    profiles=[];audits=[];query_receipts={}
    for name,witnesses in strategies.items():
        engine=DeletionEngine(system,witnesses);excess_count=0
        for deleted in subsets(sorted(s)):
            result=engine.query(deleted);fresh=system.maximum(s-deleted)
            if set(result['surviving'])!=fresh: raise ArithmeticError('Localized result differs from fresh maximum RAF.')
            excess_count+=bool(result['excess'])
        audits.append(dict(strategy=name,queries=2**len(s),inexact_regions=excess_count,all_answers_exact=True))
        for deleted in queries:
            result=engine.query(deleted);profiles.append(dict(strategy=name,deleted=' '.join(map(str,sorted(deleted))),region_size=len(result['region']),loss_size=len(result['loss']),excess=result['excess']))
        query_receipts[name]=engine.query(DELETION)
    write_csv(args.output/'regions.csv',profiles)
    (args.output/'query_certificates.json').write_text(json.dumps(query_receipts,indent=2)+'\n',encoding='utf-8')
    (args.output/'witnesses.json').write_text(json.dumps([w.record() for w in pool],indent=2)+'\n',encoding='utf-8')
    (args.output/'model.json').write_text(json.dumps(dict(food=sorted(system.food),reactions=[dict(id=r.identifier,reactants=sorted(r.reactants),products=sorted(r.products),catalysts=sorted(r.catalysts)) for r in system.reactions.values()]),indent=2)+'\n',encoding='utf-8')
    cooperative,_=paired_source(1);cycle=FunctionalSource((1,2,0));gateway=FunctionalSource((0,)*GATEWAY_SIZE)
    polynomials={'cooperative':RobustnessPolynomial(cooperative),'cycle':RobustnessPolynomial(cycle.network())}
    sensitivities={name:poly.sensitivity() for name,poly in polynomials.items()};curves=[]
    for p in RETENTION_GRID:
        for name,poly in polynomials.items():
            moments=poly.moments(p);curves.append(dict(source=name,p=str(p),mean_fraction=str(moments['mean']/len(poly.baseline)),nonempty=str(moments['nonempty']),variance=str(moments['variance'])))
        gm=gateway.moments(p);curves.append(dict(source='gateway',p=str(p),mean_fraction=str(gm['mean']/GATEWAY_SIZE),nonempty=str(p),variance=str(gm['variance'])))
    write_csv(args.output/'robustness.csv',curves)
    source=FunctionalSource(FUNCTIONAL_PARENTS);addition=source.addition(*CATALYST_ADDITION);network=source.network(addition.alternative);n=len(source.parents)
    if n>EXHAUSTIVE_REACTION_CAP: raise ValueError('Exposure audit exceeds availability enumeration budget.')
    for available in subsets(range(n)):
        if addition.surviving(available)!=network.maximum(available): raise ArithmeticError('Two-exposure law differs from literal RAF evaluation.')
    exposures=[dict(target=r,original=sorted(source.exposures[r]),alternative=sorted(addition.alternative.exposures[r]),
        probability=str(addition.probability(r,ADDITION_RETENTION)),gain=str(addition.gain(r,ADDITION_RETENTION)),
        external_cuts=[sorted(c) for c in sorted(addition.cuts(r),key=lambda s:(len(s),sorted(s)))]) for r in range(n)]
    candidates=catalyst_candidates(source,ADDITION_RETENTION);write_csv(args.output/'catalyst_candidates.csv',candidates)
    (args.output/'exposures.json').write_text(json.dumps(exposures,indent=2)+'\n',encoding='utf-8')
    controls=[]
    for parents in ((0,2,0,4,0),(0,2,1,1,1)):
        control=FunctionalSource(parents);moments=control.moments(F(1,2));literal=RobustnessPolynomial(control.network()).moments(F(1,2))
        if any(moments[k]!=literal[k] for k in moments): raise ArithmeticError('Exposure moment formula failed.')
        controls.append(dict(parents=parents,orbit_sizes=sorted(map(len,control.exposures)),mean=str(moments['mean']),variance=str(moments['variance'])))
    plot(profiles,curves,args.output)
    summary=dict(inputs=inputs,localization_audits=audits,greedy_selection=selection,greedy_query=DeletionEngine(system,selected).query(DELETION),
        robustness_coefficients={name:dict(mean=p.mean,nonempty=p.nonempty,second_moment=p.second) for name,p in polynomials.items()},
        sensitivities=sensitivities,overlap_controls=controls,best_single_catalyst_addition=candidates[0],exposure_queries_checked=2**n,
        evidence='Exact finite sets, checked closure/pruning schedules, rational availability laws. No Lean compilation, biological calibration, kinetic-persistence claim, or replay of the paper charge-model/runtime theorem.')
    transcript=json.dumps(summary,indent=2)
    (args.output/'summary.json').write_text(transcript+'\n',encoding='utf-8');(args.output/'console.txt').write_text(intro+'\n'+transcript+'\n',encoding='utf-8')
    import matplotlib
    metadata=dict(paper_sha256=PAPER_SHA256,source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),python=platform.python_version(),
        matplotlib=matplotlib.__version__,platform=platform.platform(),elapsed_seconds=time.perf_counter()-start,command='python example.py --output outputs',seed_policy='No random sampling.',
        output_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(args.output.iterdir()) if p.is_file() and p.name!='run_metadata.json'})
    (args.output/'run_metadata.json').write_text(json.dumps(metadata,indent=2)+'\n',encoding='utf-8');print(transcript)


if __name__=='__main__': main()