"""Validate a supplied irrRAF family, certify completeness, or return a missing member."""
# EDITABLE INPUTS -------------------------------------------------------------
STATEMENTS = ("0", "1", "2")
IMPLICATIONS = ((("0",), "1"), (("1",), "2"), (("2",), "0"))
AXIOM_SLOTS = 1
# Demonstrate a complete family as well: 3 independent statements need 3 axioms.
CONTROL_STATEMENTS = ("0", "1", "2")
CONTROL_IMPLICATIONS = ()
CONTROL_SLOTS = 2
# Optional graph-to-RAF reduction, reproducing the paper's second mechanism.
GRAPH_VERTICES = 3
GRAPH_EDGES = ((0, 1), (1, 2), (0, 2))
CLIQUE_SLOTS = 3
# Safe, explicit search budget: exhausting it gives INCONCLUSIVE, never complete.
MAX_TRANSVERSALS = 100000
# ---------------------------------------------------------------------------

import argparse
from dataclasses import dataclass
import hashlib
from itertools import combinations, product
import json
import math
from pathlib import Path
import platform


@dataclass(frozen=True)
class Reaction:
    name: str
    reactants: frozenset[str]
    products: frozenset[str]
    catalysts: frozenset[str]

    def __post_init__(self):
        for field in ("reactants", "products", "catalysts"):
            object.__setattr__(self, field, frozenset(getattr(self, field)))


@dataclass(frozen=True)
class ReactionSystem:
    food: frozenset[str]
    reactions: tuple[Reaction, ...]

    def __post_init__(self):
        object.__setattr__(self, "food", frozenset(self.food))
        object.__setattr__(self, "reactions", tuple(self.reactions))
        if len(self.ids) != len(self.reactions):
            raise ValueError("Reaction identifiers must be unique")

    @property
    def ids(self):
        return frozenset(r.name for r in self.reactions)

    def select(self, allowed):
        allowed = frozenset(allowed)
        if not allowed <= self.ids:
            raise ValueError("Unknown reaction identifier")
        return tuple(r for r in self.reactions if r.name in allowed)

    def closure_stages(self, allowed):
        selected = self.select(allowed)
        stages = [self.food]
        while True:
            next_state = stages[-1] | frozenset(x for r in selected if r.reactants <= stages[-1] for x in r.products)
            if next_state == stages[-1]:
                return tuple(stages)
            stages.append(next_state)

    def is_raf(self, allowed):
        selected = self.select(allowed)
        closure = self.closure_stages(allowed)[-1]
        return bool(selected) and all(r.reactants <= closure and r.catalysts & closure for r in selected)

    def to_dict(self):
        return {"food": sorted(self.food), "reactions": [{"name": r.name,
                "reactants": sorted(r.reactants), "products": sorted(r.products),
                "catalysts": sorted(r.catalysts)} for r in self.reactions]}

    @classmethod
    def from_dict(cls, data):
        return cls(frozenset(data["food"]), tuple(Reaction(**row) for row in data["reactions"]))


class RAFOracle:
    def __init__(self, system):
        self.system = system
        self.calls = 0

    def maximum(self, allowed=None):
        self.calls += 1
        remaining = self.system.ids if allowed is None else frozenset(allowed)
        self.system.select(remaining)
        while remaining:
            closure = self.system.closure_stages(remaining)[-1]
            keep = frozenset(r.name for r in self.system.select(remaining)
                             if r.reactants <= closure and r.catalysts & closure)
            if keep == remaining:
                break
            remaining = keep
        return remaining

    def irreducible(self, allowed):
        allowed = frozenset(allowed)
        return self.system.is_raf(allowed) and all(not self.maximum(allowed-{r}) for r in sorted(allowed))

    def extract(self, allowed):
        """One inclusion-minimal RAF; no minimum-cardinality claim."""
        current = self.maximum(allowed)
        for r in sorted(current):
            if r in current:
                smaller = self.maximum(current-{r})
                if smaller:
                    current = smaller
        return current


class CompletenessChecker:
    def __init__(self, system):
        self.system = system

    def check(self, supplied, max_transversals=None):
        # A family is a SET of reaction sets. Duplicate entries do not increase k.
        family = tuple(sorted(set(map(frozenset, supplied)), key=lambda s: tuple(sorted(s))))
        oracle = RAFOracle(self.system)
        base = {"family": [sorted(s) for s in family], "k": len(family),
                "largest_listed_size": max(map(len,family),default=1),
                "transversal_bound": math.prod(map(len,family))}
        for member in family:
            if not member <= self.system.ids or not oracle.irreducible(member):
                return {**base, "status": "invalid", "invalid_entry": sorted(member),
                        "reason": "Each distinct supplied member must be an irreducible RAF"}
        validity_calls = oracle.calls
        if max_transversals is not None and max_transversals < 0:
            raise ValueError("Search budget must be nonnegative or None")
        trace, seen = [], {}
        for index, choices in enumerate(product(*(sorted(s) for s in family))):
            if max_transversals is not None and index >= max_transversals:
                return {**base, "status": "inconclusive", "validity_oracle_calls": validity_calls,
                        "tested_tuples": len(trace), "distinct_residual_queries": len(seen), "trace": trace}
            deleted = frozenset(choices)
            cached = deleted in seen
            if not cached:
                seen[deleted] = oracle.maximum(self.system.ids-deleted)
            residual = seen[deleted]
            trace.append({"choices": choices, "deleted": sorted(deleted),
                          "residual_maximum": sorted(residual), "cached": cached})
            if residual:
                missing = oracle.extract(residual)
                if missing in family or not oracle.irreducible(missing):
                    raise ArithmeticError("Missing-member certificate failed")
                return {**base, "status": "incomplete", "missing": sorted(missing),
                        "omissions": choices, "validity_oracle_calls": validity_calls,
                        "tested_tuples": len(trace), "distinct_residual_queries": len(seen), "trace": trace,
                        "missing_closure_stages": [sorted(s) for s in self.system.closure_stages(missing)]}
        return {**base, "status": "complete", "validity_oracle_calls": validity_calls,
                "tested_tuples": len(trace), "distinct_residual_queries": len(seen), "trace": trace}

    def verify_report(self, report):
        """Recompute the finite claim; no trust in a saved success label or trace."""
        status = report.get("status")
        family = tuple(map(frozenset, report.get("family", ())))
        oracle = RAFOracle(self.system)
        if status == "incomplete":
            try:
                return (all(oracle.irreducible(s) for s in family)
                        and frozenset(report["missing"]) not in family
                        and oracle.irreducible(report["missing"]))
            except (KeyError, ValueError):
                return False
        if status == "complete":
            return self.check(family)["status"] == "complete"
        return False


@dataclass(frozen=True)
class Implication:
    premises: frozenset[str]
    conclusion: str

    def __post_init__(self):
        object.__setattr__(self, "premises", frozenset(self.premises))


@dataclass(frozen=True)
class ImplicationSystem:
    statements: tuple[str, ...]
    rules: tuple[Implication, ...]

    def __post_init__(self):
        if len(set(self.statements)) != len(self.statements):
            raise ValueError("Duplicate statement")
        universe = set(self.statements)
        if any(not r.premises <= universe or r.conclusion not in universe for r in self.rules):
            raise ValueError("Rule refers to an unknown statement")

    def stages(self, axioms):
        current = frozenset(axioms)
        if not current <= set(self.statements):
            raise ValueError("Unknown axiom")
        stages = [current]
        while True:
            new = current | frozenset(r.conclusion for r in self.rules if r.premises <= current)
            if new == current:
                return tuple(stages)
            stages.append(new)
            current = new

    def generators(self, budget):
        return tuple(frozenset(s) for size in range(min(budget,len(self.statements))+1)
                     for s in combinations(self.statements,size) if self.stages(s)[-1] == frozenset(self.statements))

    def normalized(self):
        statements, rules = list(self.statements), list(self.rules)
        while len(statements) < 2:
            new = "_padding_"+str(len(statements))
            while new in statements:
                new += "_"
            statements.append(new)
            rules.append(Implication(frozenset(),new))
        return ImplicationSystem(tuple(statements),tuple(rules))


class AxiomReduction:
    """Literal selector / decoder / implication / closing construction."""
    def __init__(self, source, slots):
        if not isinstance(slots,int) or slots < 0:
            raise ValueError("Nonnegative integer slot count required")
        self.source, self.slots = source.normalized(), slots
        m = len(self.source.statements)
        reactions = []
        self.guards = tuple(frozenset(f"s:{i}:{u}" for u in range(m)) for i in range(slots))
        for i in range(slots):
            for u in range(m):
                reactions.append(Reaction(f"s:{i}:{u}", {"f"}, {f"x:{i}:{u}",f"c:{i}:{u}"}, {"z",f"c:{i}:{(u+1)%m}"}))
                reactions.append(Reaction(f"d:{i}:{u}", {f"x:{i}:{v}" for v in range(m) if v != u}, {f"b:{i}",f"A:{u}"}, {"z"}))
        lookup = {s:i for i,s in enumerate(self.source.statements)}
        for j, rule in enumerate(self.source.rules):
            reactions.append(Reaction(f"p:{j}", {"f"}|{f"A:{lookup[v]}" for v in rule.premises}, {f"A:{lookup[rule.conclusion]}"}, {"z"}))
        reactions.append(Reaction("close", {f"b:{i}" for i in range(slots)}|{f"A:{u}" for u in range(m)}, {"z"}, {"z"}))
        self.system = ReactionSystem(frozenset(("f",)), tuple(reactions))

    def decode_missing(self, missing):
        missing = frozenset(missing)
        if missing in self.guards or "close" not in missing or not RAFOracle(self.system).irreducible(missing):
            raise ValueError("Supply an unlisted irrRAF of this constructed system")
        choices = []
        for guard in self.guards:
            omitted = guard-missing
            if len(omitted) != 1:
                raise ArithmeticError("Omission pattern failed")
            u = int(next(iter(omitted)).split(":")[-1])
            choices.append(self.source.statements[u])
        stages = self.source.stages(choices)
        if stages[-1] != frozenset(self.source.statements):
            raise ArithmeticError("Extracted axioms do not generate the universe")
        return {"slot_choices": choices, "axioms": sorted(set(choices)),
                "derivation_stages": [sorted(s) for s in stages]}

    def source_witness(self, choices):
        if len(choices) != self.slots or self.source.stages(choices)[-1] != frozenset(self.source.statements):
            raise ValueError("Supply one generating choice per slot")
        indices = [self.source.statements.index(s) for s in choices]
        result = {r for i,g in enumerate(self.guards) for r in g if r != f"s:{i}:{indices[i]}"}
        result |= {f"d:{i}:{u}" for i,u in enumerate(indices)}
        result |= {f"p:{j}" for j in range(len(self.source.rules))}|{"close"}
        if not self.system.is_raf(result):
            raise ArithmeticError("Forward witness failed")
        return frozenset(result)


class CliqueReduction:
    """Paper's independent ETH reduction, retaining all ordered pair gates."""
    def __init__(self, vertices, edges, slots):
        if vertices < 2 or slots < 0:
            raise ValueError("At least two vertices and nonnegative slots required")
        self.vertices, self.slots = vertices, slots
        self.edges = frozenset(frozenset(e) for e in edges)
        if any(len(e) != 2 or not e <= set(range(vertices)) for e in self.edges):
            raise ValueError("A simple graph on the stated vertices is required")
        # Guards are identical to the axiom construction; only decoder products
        # and the verifier differ. Build this network independently for clarity.
        reactions = []
        self.guards = tuple(frozenset(f"s:{i}:{u}" for u in range(vertices)) for i in range(slots))
        for i in range(slots):
            for u in range(vertices):
                reactions.append(Reaction(f"s:{i}:{u}", {"f"}, {f"x:{i}:{u}",f"c:{i}:{u}"}, {"z",f"c:{i}:{(u+1)%vertices}"}))
                reactions.append(Reaction(f"d:{i}:{u}", {f"x:{i}:{v}" for v in range(vertices) if v != u}, {f"b:{i}"}, {"z"}))
        positions = tuple(product(range(slots),range(vertices)))
        pairs = tuple(product(positions,repeat=2))
        for index, ((i,u),(j,v)) in enumerate(pairs):
            forbidden = i != j and frozenset((u,v)) not in self.edges
            for side, position in (("l",(i,u)),("r",(j,v))):
                inputs = {f"x:{position[0]}:{position[1]}"} if forbidden else {"f"}
                reactions.append(Reaction(f"{side}:{index}", inputs, {f"y:{index}"}, {"z"}))
        reactions.append(Reaction("close", {f"b:{i}" for i in range(slots)}|{f"y:{i}" for i in range(len(pairs))}, {"z"}, {"z"}))
        self.system = ReactionSystem(frozenset(("f",)),tuple(reactions))

    def decode_missing(self, missing):
        missing = frozenset(missing)
        if missing in self.guards or not RAFOracle(self.system).irreducible(missing):
            raise ValueError("Expected an unlisted irreducible RAF")
        vertices = []
        for guard in self.guards:
            omitted = guard-missing
            if len(omitted) != 1:
                raise ArithmeticError("Expected one omitted selector per slot")
            vertices.append(int(next(iter(omitted)).split(":")[-1]))
        if len(set(vertices)) != self.slots or any(frozenset(pair) not in self.edges for pair in combinations(vertices,2)):
            raise ArithmeticError("Decoded choices are not a clique")
        return vertices


def literal_family(system):
    """Independent all-subsets predicate enumeration for small examples only."""
    if len(system.reactions) > 18:
        raise ValueError("Literal enumeration is limited to 18 reactions")
    rafs, irreducibles = [], []
    for size in range(1,len(system.reactions)+1):
        for subset in combinations(sorted(system.ids),size):
            subset = frozenset(subset)
            if system.is_raf(subset):
                rafs.append(subset)
                if not any(i < subset for i in irreducibles):
                    irreducibles.append(subset)
    return tuple(rafs), tuple(irreducibles)


def coverage_example():
    return ReactionSystem(frozenset(("f",)),tuple(Reaction(s,{"f"},{s},set("abc")-{s}) for s in "abc"))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output",type=Path,default=Path("outputs"))
    parser.add_argument("--input",type=Path,help="JSON with food, reactions, and supplied_family; skips built-in demonstrations")
    args = parser.parse_args()
    args.output.mkdir(parents=True,exist_ok=True)
    if args.input:
        data = json.loads(args.input.read_text())
        system = ReactionSystem.from_dict(data)
        report = CompletenessChecker(system).check(data["supplied_family"],MAX_TRANSVERSALS)
        (args.output/"certification.json").write_text(json.dumps(report,indent=2)+"\n")
        print(report["status"])
        return
    source = ImplicationSystem(STATEMENTS,tuple(Implication(frozenset(p),c) for p,c in IMPLICATIONS))
    reduction = AxiomReduction(source,AXIOM_SLOTS)
    checker = CompletenessChecker(reduction.system)
    report = checker.check(reduction.guards,MAX_TRANSVERSALS)
    decoded = reduction.decode_missing(report["missing"]) if report["status"] == "incomplete" else None
    control_source = ImplicationSystem(CONTROL_STATEMENTS,tuple(Implication(frozenset(p),c) for p,c in CONTROL_IMPLICATIONS))
    control = AxiomReduction(control_source,CONTROL_SLOTS)
    control_report = CompletenessChecker(control.system).check(control.guards,MAX_TRANSVERSALS)
    overlap = coverage_example()
    overlap_report = CompletenessChecker(overlap).check(({"a","b"},{"a","c"}),MAX_TRANSVERSALS)
    clique = CliqueReduction(GRAPH_VERTICES,GRAPH_EDGES,CLIQUE_SLOTS)
    clique_report = CompletenessChecker(clique.system).check(clique.guards,MAX_TRANSVERSALS)
    decoded_clique = clique.decode_missing(clique_report["missing"]) if clique_report["status"] == "incomplete" else None
    # Literal enumeration is an optional independent check, never needed by the
    # reusable completeness checker. Larger edited examples remain runnable.
    if len(reduction.system.reactions) <= 18:
        rafs, irr = literal_family(reduction.system)
        counts = {"rafs":len(rafs),"irreducibles":len(irr),"members":[sorted(s) for s in irr]}
    else:
        counts = {"skipped":"More than 18 reactions; literal enumeration would be exponential"}
    results = {"axiom_instance": report,"decoded_axioms":decoded,"literal_check":counts,
               "complete_control":control_report,"coverage_is_not_completeness":overlap_report,
               "clique_instance":clique_report,"decoded_clique":decoded_clique}
    (args.output/"results.json").write_text(json.dumps(results,indent=2)+"\n")
    source_data = {**reduction.system.to_dict(),"supplied_family":[sorted(g) for g in reduction.guards]}
    (args.output/"model_input.json").write_text(json.dumps(source_data,indent=2)+"\n")
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    plt.rcParams.update({"font.size":11,"axes.spines.top":False,"axes.spines.right":False,"svg.hashsalt":"irrraf-completeness-v1"})
    fig, axes = plt.subplots(1,2,figsize=(12,4.8),layout="constrained")
    # Show the small coverage loss with its independently enumerated exact family.
    _, overlap_irr = literal_family(overlap)
    members = [sorted(s) for s in overlap_irr]
    axes[0].imshow([[int(r in member) for r in "abc"] for member in members],cmap="Blues",vmin=0,vmax=1,aspect="auto")
    axes[0].set(xticks=range(3),xticklabels=["a","b","c"],yticks=range(3),
                yticklabels=["Listed {a,b}","Listed {a,c}","Missing {b,c}"],title="Covering every reaction can still miss an irrRAF")
    for row,member in enumerate(members):
        for col,r in enumerate("abc"):
            axes[0].text(col,row,"present" if r in member else "absent",ha="center",va="center",color="white" if r in member else "#333333")
    for largest in (2,3,5,10):
        axes[1].plot(range(1,9),[largest**k for k in range(1,9)],"o-",label=f"Largest listed size {largest}")
    axes[1].set(yscale="log",xlabel="Number k of distinct listed irrRAFs",ylabel="Upper bound on deletion tuples",title="Upper bound on completeness-search\ncombinations")
    axes[1].legend(frameon=False,fontsize=9)
    for ext in ("png","svg"):
        fig.savefig(args.output/f"completeness.{ext}",dpi=165)
    plt.close(fig)
    console = "\n".join(["Completeness of irreducible RAF families",
        f"Default source: {len(source.statements)} statements, {len(source.rules)} rules, {AXIOM_SLOTS} slots",
        "Literal check: "+str({k:v for k,v in counts.items() if k != "members"}),
        "Supplied guard family: "+report["status"],"Decoded source witness: "+str(decoded),
        "Independent-statement control: "+control_report["status"]+f" after {control_report.get('tested_tuples',0)} tuples",
        "Overlapping coverage list: "+overlap_report["status"]+"; missing "+str(overlap_report.get("missing")),
        "Graph reduction: "+clique_report["status"]+"; decoded clique "+str(decoded_clique),
        "Tuple bounds are combinatorial counts, not measured runtimes. No Lean re-execution claimed."])
    print(console)
    (args.output/"console.txt").write_text(console+"\n")
    metadata = {"paper_sha256":"c89e78a1ffadcaaca3ac047d71817feef3b32d26f6c0af1b1d245082499e35ea",
                "source_sha256":hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                "python":platform.python_version(),"matplotlib":matplotlib.__version__,
                "inputs":{"statements":STATEMENTS,"implications":IMPLICATIONS,"slots":AXIOM_SLOTS,
                          "control_statements":CONTROL_STATEMENTS,"control_implications":CONTROL_IMPLICATIONS,"control_slots":CONTROL_SLOTS,
                          "graph_vertices":GRAPH_VERTICES,"graph_edges":GRAPH_EDGES,"clique_slots":CLIQUE_SLOTS,"max_transversals":MAX_TRANSVERSALS},
                "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")


if __name__ == "__main__":
    main()
