"""Supported-trace SAT enumeration and the paper's exact SAT-to-RAF family source."""
# EDITABLE INPUTS -------------------------------------------------------------
DECLARED_VARIABLES = 3
SATISFIABLE_CLAUSES = ((1,), (2,))
UNSATISFIABLE_CLAUSES = ((1,), (-1,))
SOLVER_NAME = "g3"                 # installed PySAT Glucose3 backend
MAX_SOLVER_CALLS = 1000            # includes the final UNSAT completion call
CONFLICT_BUDGET_PER_CALL = 100000  # exhausted solve returns unknown, not complete
MAX_FORMULA_LITERALS = 2000000     # guard compilation before allocating clauses
MAX_TRUTH_TABLE_ASSIGNMENTS = 4096 # independent source-check limit, not an SAT-solver limit
# ---------------------------------------------------------------------------

import argparse
from dataclasses import dataclass
import hashlib
from itertools import combinations, product
import json
import math
from pathlib import Path
import platform
import importlib.metadata

@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 FormulaLimit(RuntimeError):
    pass


class SupportedTraceCNF:
    """Literal one-way support formula from the paper, without stronger axioms."""
    def __init__(self, system, known=(), container=None, max_literals=None):
        self.system = system
        self.reactions = tuple(r.name for r in system.reactions)
        self.species = tuple(sorted(system.food.union(*(r.reactants | r.products | r.catalysts for r in system.reactions))))
        self.d, self.r = len(self.species),len(self.reactions)
        self.container = system.ids if container is None else frozenset(container)
        system.select(self.container)
        self.known = tuple(map(frozenset,known))
        if any(not member <= system.ids for member in self.known):
            raise ValueError("Unknown identifier in avoidance family")
        self.variable_count = self.r+self.d*(self.d+1)+self.d*self.r
        self.expected_counts = self.counts()
        if max_literals is not None and self.expected_counts["literals"] > max_literals:
            raise FormulaLimit("Supported-trace formula exceeds the configured literal limit")
        self.species_index = {x:i for i,x in enumerate(self.species)}
        self.reaction_index = {x:i for i,x in enumerate(self.reactions)}
        clauses = [tuple(self.selection(a) for a in range(self.r))]
        clauses += [(-self.selection(a),) for a,name in enumerate(self.reactions) if name not in self.container]
        clauses += [(-self.row(0,x),) for x,name in enumerate(self.species) if name not in system.food]
        for i in range(self.d):
            for a,reaction in enumerate(system.reactions):
                clauses.append((-self.firing(i,a),self.selection(a)))
                clauses += [(-self.firing(i,a),self.row(i,self.species_index[x])) for x in sorted(reaction.reactants)]
            for x,name in enumerate(self.species):
                producers = tuple(self.firing(i,a) for a,reaction in enumerate(system.reactions) if name in reaction.products)
                clauses.append((-self.row(i+1,x),self.row(i,x))+producers)
        for a,reaction in enumerate(system.reactions):
            clauses += [(-self.selection(a),self.row(self.d,self.species_index[x])) for x in sorted(reaction.reactants)]
            clauses.append((-self.selection(a),)+tuple(self.row(self.d,self.species_index[x]) for x in sorted(reaction.catalysts)))
        clauses += [self.block(member) for member in self.known]
        self.clauses = tuple(clauses)
        if len(clauses) != self.expected_counts["clauses"] or sum(map(len,clauses)) != self.expected_counts["literals"]:
            raise ArithmeticError("Compiler clause accounting differs from the paper")

    def selection(self,a):
        return a+1

    def row(self,i,x):
        return self.r+i*self.d+x+1

    def firing(self,i,a):
        return self.r+self.d*(self.d+1)+i*self.r+a+1

    def block(self,member):
        return tuple(-self.selection(self.reaction_index[a]) for a in sorted(member))

    def counts(self):
        d,r = self.d,self.r
        ar = sum(len(a.reactants) for a in self.system.reactions)
        ap = sum(len(a.products) for a in self.system.reactions)
        ac = sum(len(a.catalysts) for a in self.system.reactions)
        h,t = len(self.system.food),len(self.system.ids-self.container)
        g,hg = len(self.known),sum(map(len,self.known))
        return {"variables":r+d*(d+1)+d*r,
                "clauses":1+t+(d-h)+d*r+d*ar+d*d+ar+r+g,
                "literals":2*r+t+(d-h)+2*d*r+2*(d+1)*ar+2*d*d+d*ap+ac+hg}

    def dimacs(self,extra=()):
        clauses = self.clauses+tuple(extra)
        return f"p cnf {self.variable_count} {len(clauses)}\n"+"".join(" ".join(map(str,c))+" 0\n" for c in clauses)

    def witness(self,selected):
        """Canonical satisfying assignment from a RAF's actual staged closure."""
        selected = frozenset(selected)
        stages = self.system.closure_stages(selected)
        rows = [stages[min(i,len(stages)-1)] for i in range(self.d+1)]
        true = {self.selection(a) for a,name in enumerate(self.reactions) if name in selected}
        true |= {self.row(i,x) for i in range(self.d+1) for x,name in enumerate(self.species) if name in rows[i]}
        true |= {self.firing(i,a) for i in range(self.d) for a,reaction in enumerate(self.system.reactions)
                 if reaction.name in selected and reaction.reactants <= rows[i]}
        return tuple(v if v in true else -v for v in range(1,self.variable_count+1))

    def check_model(self,model,extra=()):
        true = frozenset(v for v in model if v > 0)
        if not all(any((literal > 0) == (abs(literal) in true) for literal in clause)
                   for clause in self.clauses+tuple(extra)):
            raise ValueError("SAT model fails a compiled clause")
        selected = frozenset(name for a,name in enumerate(self.reactions) if self.selection(a) in true)
        if not self.system.is_raf(selected) or not selected <= self.container:
            raise ValueError("Selected reactions fail independent RAF validation")
        if any(member <= selected for member in self.known):
            raise ValueError("Model failed known-family avoidance")
        # Check the soundness invariant, allowing supported rows to omit molecules.
        stages = self.system.closure_stages(selected)
        for i in range(self.d+1):
            supported = {name for x,name in enumerate(self.species) if self.row(i,x) in true}
            if not supported <= stages[min(i,len(stages)-1)]:
                raise ValueError("Unsupported molecule in the SAT trace")
        return selected


class PySATBackend:
    """Incremental solver adapter. UNSAT trusts the backend, not a Lean certificate."""
    def __init__(self,clauses,name=SOLVER_NAME,conflict_budget=CONFLICT_BUDGET_PER_CALL):
        from pysat.solvers import Solver
        self.solver = Solver(name=name,bootstrap_with=clauses)
        self.conflict_budget = conflict_budget

    def add_clause(self,clause):
        self.solver.add_clause(list(clause))

    def solve(self):
        if self.conflict_budget is None:
            status = self.solver.solve()
        else:
            self.solver.conf_budget(self.conflict_budget)
            status = self.solver.solve_limited()
        if status is None:
            return "unknown",None
        return ("sat",self.solver.get_model()) if status else ("unsat",None)

    def close(self):
        self.solver.delete()


class ModelEnumerator:
    def __init__(self,system,backend_factory=PySATBackend):
        self.system,self.backend_factory = system,backend_factory

    def run(self,known=(),max_calls=MAX_SOLVER_CALLS,max_literals=MAX_FORMULA_LITERALS):
        family = list(sorted(set(map(frozenset,known)),key=lambda s:tuple(sorted(s))))
        initial = len(family)
        validator = RAFOracle(self.system)
        if any(not validator.irreducible(i) for i in family):
            raise ValueError("Every initial member must be an irrRAF")
        try:
            formula = SupportedTraceCNF(self.system,family,max_literals=max_literals)
        except FormulaLimit as error:
            return {"status":"incomplete","reason":str(error),"family":[sorted(i) for i in family],"solver_calls":0}
        backend = self.backend_factory(formula.clauses)
        events,blocks,calls,min_calls = [],[],0,0
        status,reason = "incomplete","Solver-call limit reached before a final UNSAT result"
        try:
            while max_calls is None or calls < max_calls:
                verdict,model = backend.solve()
                calls += 1
                if verdict == "unknown":
                    reason = "Solver returned unknown or reached its resource limit"
                    break
                if verdict == "unsat":
                    status,reason = "complete","Final supported-trace query is UNSAT"
                    break
                if verdict != "sat" or model is None:
                    raise ValueError("Invalid solver response")
                selected = formula.check_model(model,blocks)
                if any(i <= selected for i in family):
                    raise ValueError("Backend returned a blocked reaction set")
                oracle = RAFOracle(self.system)
                current = selected
                # The whole sweep stays inside the avoiding SAT witness.
                for reaction in self.system.reactions:
                    if reaction.name in current:
                        smaller = oracle.maximum(current-{reaction.name})
                        if smaller:
                            current = smaller
                min_calls += oracle.calls
                if current in family or not validator.irreducible(current):
                    raise ArithmeticError("Emitted set is not a new irreducible RAF")
                family.append(current)
                clause = formula.block(current)
                blocks.append(clause)
                backend.add_clause(clause)
                events.append({"selected":sorted(selected),"emitted":sorted(current),"minimization_calls":oracle.calls})
        finally:
            backend.close()
        if status == "complete" and calls != len(family)-initial+1:
            raise ArithmeticError("Completed solver-call identity failed")
        return {"status":status,"reason":reason,"initial_members":initial,"family":[sorted(i) for i in family],
                "solver_calls":calls,"minimization_maxraf_calls":min_calls,"events":events,
                "base_formula":formula.expected_counts,"final_blocking_clauses":len(family)}


@dataclass(frozen=True)
class BooleanFormula:
    variables: int
    clauses: tuple[tuple[int,...],...]

    def __post_init__(self):
        if self.variables < 2:
            raise ValueError("Declare at least two variables; padding changes the model count")
        if any(v == 0 or abs(v) > self.variables for c in self.clauses for v in c):
            raise ValueError("Literals must use signed one-based variable indices")

    def evaluate(self,assignment):
        return all(any(bool(assignment[abs(v)-1]) == (v > 0) for v in c) for c in self.clauses)

    def satisfying_assignments(self):
        if 2**self.variables > MAX_TRUTH_TABLE_ASSIGNMENTS:
            raise ValueError("Source truth-table audit exceeds its limit; use ModelEnumerator on your network or explicitly raise the audit limit")
        return tuple(a for a in product((False,True),repeat=self.variables) if self.evaluate(a))

    def accepts_choices(self,choices):
        choices = frozenset(choices)
        conflict = any({(i,False),(i,True)} <= choices for i in range(self.variables))
        coverage = all((i,False) in choices or (i,True) in choices for i in range(self.variables))
        hits = all(any((abs(v)-1,v>0) in choices for v in c) for c in self.clauses)
        return conflict or (coverage and hits)


class SATFamilySource:
    """Whole-family source: mandatory auxiliary cycle plus reset and literal inputs."""
    def __init__(self,formula):
        self.formula = formula
        n,m = formula.variables,len(formula.clauses)
        def literal(i,b): return f"literal:{i}:{int(b)}"
        wires = {literal(i,b) for i in range(n) for b in (False,True)}
        wires |= {f"covered:{i}" for i in range(n)}|{f"clause:{j}" for j in range(m)}|{"out"}
        rules = []
        for i in range(n):
            rules.append(({literal(i,False),literal(i,True)},{"out"}))
        for i,b in product(range(n),(False,True)):
            rules.append(({literal(i,b)},{f"covered:{i}"}))
        for j,clause in enumerate(formula.clauses):
            for i,b in product(range(n),(False,True)):
                signed = i+1 if b else -(i+1)
                rules.append(({literal(i,b)},{f"clause:{j}"} if signed in clause else set()))
        rules.append(({f"covered:{i}" for i in range(n)}|{f"clause:{j}" for j in range(m)},{"out"}))
        q = len(rules)
        self.auxiliary = frozenset(f"aux:{j}" for j in range(q+1))
        self.inputs = {(i,b):f"input:{i}:{int(b)}" for i,b in product(range(n),(False,True))}
        reactions = [Reaction(name,{"food"},{literal(i,b)},{f"marker:{q}"}) for (i,b),name in self.inputs.items()]
        reactions += [Reaction(f"aux:{j}",required,produced|{f"marker:{j}"},{f"marker:{j+1}"}) for j,(required,produced) in enumerate(rules)]
        reactions.append(Reaction(f"aux:{q}",{"out"},wires|{f"marker:{q}"},{"marker:0"}))
        self.system = ReactionSystem({"food"},tuple(reactions))
        self.reset = f"aux:{q}"
        self.baseline = tuple(self.canonical(((i,False),(i,True))) for i in range(n))

    def canonical(self,choices):
        return self.auxiliary|frozenset(self.inputs[c] for c in choices)

    def predicted_family(self):
        return frozenset(self.baseline)|frozenset(self.canonical(enumerate(a)) for a in self.formula.satisfying_assignments())

    def dimensions(self):
        n,m = self.formula.variables,len(self.formula.clauses)
        d,r = 6*n+2*n*m+m+4,5*n+2*n*m+2
        return {"molecules":d,"reactions":r,"input_bits":2*d+r+2+3*d*r,"baseline_output_bits":(r+1)*n+1}


def output_bits(reaction_order,family):
    family = tuple(map(frozenset,family))
    if len(set(family)) != len(family) or any(not s <= set(reaction_order) for s in family):
        raise ValueError("Count-first output requires distinct valid reaction masks")
    return "1"*len(family)+"0"+"".join("1" if r in member else "0" for member in family for r in reaction_order)


def essential_reactions(system):
    oracle = RAFOracle(system)
    if not oracle.maximum():
        return {"status":"no RAF","reactions":[]}
    return {"status":"RAF exists","reactions":sorted(r for r in system.ids if not oracle.maximum(system.ids-{r}))}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output",type=Path,default=Path("outputs"))
    parser.add_argument("--input",type=Path,help="JSON source with food, reactions and optional known_family")
    args = parser.parse_args()
    args.output.mkdir(parents=True,exist_ok=True)
    if args.input:
        data = json.loads(args.input.read_text())
        report = ModelEnumerator(ReactionSystem.from_dict(data)).run(data.get("known_family",()))
        (args.output/"enumeration.json").write_text(json.dumps(report,indent=2)+"\n")
        print(report["status"],len(report["family"]))
        return
    cases = {"satisfiable":BooleanFormula(DECLARED_VARIABLES,SATISFIABLE_CLAUSES),
             "unsatisfiable":BooleanFormula(DECLARED_VARIABLES,UNSATISFIABLE_CLAUSES)}
    for boolean in cases.values():
        boolean.satisfying_assignments()  # check the small-source audit limit before solver work
    results = {}
    for name,boolean in cases.items():
        source = SATFamilySource(boolean)
        report = ModelEnumerator(source.system).run(source.baseline)
        expected = source.predicted_family()
        found = frozenset(map(frozenset,report["family"]))
        if not found <= expected or (report["status"] == "complete" and found != expected):
            raise ArithmeticError("Enumerated family differs from the source correspondence")
        order = tuple(r.name for r in source.system.reactions)
        encoded = output_bits(order,report["family"])
        results[name] = {"dimensions":source.dimensions(),"satisfying_assignment_count":len(boolean.satisfying_assignments()),
                         "predicted_family_count":len(expected),"enumeration":report,"output_bits":len(encoded),
                         "output_sizes":sorted(map(len,found)),"essential":essential_reactions(source.system)}
        (args.output/f"{name}_masks.txt").write_text(encoded+"\n")
        query = SupportedTraceCNF(source.system,source.baseline)
        # .txt suffix permits a standalone DIMACS artifact in the site's schema.
        (args.output/f"{name}_completion_dimacs.txt").write_text(query.dimacs())
        if name == "satisfiable":
            source_data = {**source.system.to_dict(),"known_family":[sorted(i) for i in source.baseline]}
            (args.output/"model_input.json").write_text(json.dumps(source_data,indent=2)+"\n")
            # Low-order input deletions leave a conflict pair; at order n the
            # surviving assignment exposes the SAT distinction.
            by_order = []
            oracle = RAFOracle(source.system)
            for size in range(DECLARED_VARIABLES+1):
                deletions = tuple(combinations(source.inputs.values(),size))
                alive = sum(bool(oracle.maximum(source.system.ids-set(d))) for d in deletions)
                by_order.append({"deleted_inputs":size,"survivors":alive,"total":len(deletions)})
            results[name]["input_deletion_profile"] = by_order
    (args.output/"results.json").write_text(json.dumps(results,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":"sat-enumeration-v1"})
    fig,axes = plt.subplots(1,2,figsize=(12,4.8),layout="constrained")
    names = tuple(results)
    baseline = [DECLARED_VARIABLES]*len(names)
    assignments = [results[n]["satisfying_assignment_count"] for n in names]
    axes[0].bar(names,baseline,label="Known conflict-pair cores",color="#187f91")
    axes[0].bar(names,assignments,bottom=baseline,label="Satisfying-assignment cores",color="#ae6539")
    axes[0].set(ylabel="Exact number of irreducible RAFs",title="Irreducible RAF counts by Boolean formula")
    axes[0].legend(frameon=False)
    profile = results["satisfiable"]["input_deletion_profile"]
    axes[1].bar([r["deleted_inputs"] for r in profile],[r["survivors"]/r["total"] for r in profile],color="#187f91")
    axes[1].set(xlabel="Number of deleted input reactions",ylabel="Fraction of deletion sets leaving a RAF",
                ylim=(0,1.15),xticks=[r["deleted_inputs"] for r in profile],title="Small deletion tests can hide the SAT choice")
    for row in profile:
        axes[1].text(row["deleted_inputs"],row["survivors"]/row["total"]+.025,f"{row['survivors']}/{row['total']}",ha="center")
    for ext in ("png","svg"):
        fig.savefig(args.output/f"families.{ext}",dpi=165)
    plt.close(fig)
    console = "\n".join(["Supported-trace SAT enumeration; all emitted RAFs independently checked"]+
        [f"{name}: {r['enumeration']['status']}, {len(r['enumeration']['family'])} cores, {r['enumeration']['solver_calls']} SAT calls from {DECLARED_VARIABLES} known cores, output {r['output_bits']} bits" for name,r in results.items()]+
        ["Completion trusts the PySAT backend's UNSAT answer. Solver-call count is not a runtime guarantee."])
    print(console)
    (args.output/"console.txt").write_text(console+"\n")
    metadata = {"paper_sha256":"d3f3798ee1f02275a8ecbd97805982690d1d0f793ba8eb90a028fbc51d5d9004",
                "source_sha256":hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),"python":platform.python_version(),
                "python-sat":importlib.metadata.version("python-sat"),"matplotlib":matplotlib.__version__,
                "inputs":{"variables":DECLARED_VARIABLES,"satisfiable_clauses":SATISFIABLE_CLAUSES,
                          "unsatisfiable_clauses":UNSATISFIABLE_CLAUSES,"solver":SOLVER_NAME,"max_calls":MAX_SOLVER_CALLS,
                          "conflict_budget":CONFLICT_BUDGET_PER_CALL,"formula_literal_limit":MAX_FORMULA_LITERALS,
                          "truth_table_assignment_limit":MAX_TRUTH_TABLE_ASSIGNMENTS},
                "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()


