"""Exact interfaces, usable independently: prices, passive tails, stars, join trees.

Run: python example.py --output outputs. See README.md for units and scope.
Scientific calculations use Fraction; plotting is the only floating conversion.
"""
from __future__ import annotations

# EDITABLE INPUTS -------------------------------------------------------------
# Stoichiometric coefficients and amplification q are dimensionless.
AMPLIFICATION_THRESHOLDS = ("1", "2", "3", "4", "5")
# A reversible unary chain X0 <-> X1 <-> ... <-> X4.
# Rate constants and internal degradation: inverse time, e.g. h^-1.
TAIL_FORWARD = ("2", "3", "1", "4")
TAIL_BACKWARD = ("1", "2", "2", "1")
TAIL_DEGRADATION = ("1/5", "1/2", "1/10")
LEFT_CONCENTRATION = "1"       # arbitrary consistent concentration unit
RIGHT_CONCENTRATION = "1/5"
# Paper's star counterexample; entries have units inverse time.
ROOT_DIAGONAL = "-1"
BRANCH_DIAGONAL = "-15"
ROOT_TO_BRANCH = "3"
BRANCH_TO_ROOT = "3"
BRANCH_COUNTS = (1, 2, 3, 4)
# Finite examples use exact binary relations; no discretization of chemistry.
FINITE_DOMAIN = (0, 1)  # the worked Boolean gluing example requires these values
# ---------------------------------------------------------------------------

import argparse
import csv
from dataclasses import dataclass
from fractions import Fraction as F
import hashlib
from itertools import combinations, product
import json
from pathlib import Path
import platform
from typing import Hashable, Iterable


def rational(value):
    """Decimal strings retain their intended decimal, including supplied floats."""
    return F(str(value))


def dot(a, b):
    return sum((x * y for x, y in zip(a, b, strict=True)), F(0))


def solve(matrix, rhs):
    """Exact square Gaussian elimination; singular systems return None."""
    n = len(rhs)
    if len(matrix) != n or any(len(row) != n for row in matrix):
        raise ValueError("Expected a square system")
    a = [[rational(x) for x in row] + [rational(y)]
         for row, y in zip(matrix, rhs, strict=True)]
    for j in range(n):
        pivot = next((i for i in range(j, n) if a[i][j]), None)
        if pivot is None:
            return None
        a[j], a[pivot] = a[pivot], a[j]
        scale = a[j][j]
        a[j] = [x / scale for x in a[j]]
        for i in range(n):
            if i != j:
                scale = a[i][j]
                a[i] = [x - scale * y for x, y in zip(a[i], a[j])]
    return tuple(row[-1] for row in a)


@dataclass(frozen=True)
class RationalPolytope:
    """Small exact LP: inequalities row*x >= bound plus one normalization.

    Vertex enumeration is deliberately bounded to small models. It is exponential;
    use a larger LP backend with exact certificate verification for large networks.
    """
    normalization: tuple[F, ...]
    inequalities: tuple[tuple[tuple[F, ...], F], ...]

    def vertices(self):
        n = len(self.normalization)
        found = set()
        for active in combinations(self.inequalities, n - 1):
            x = solve([self.normalization] + [r for r, _ in active],
                      [F(1)] + [b for _, b in active])
            if x is not None and all(dot(row, x) >= b for row, b in self.inequalities):
                found.add(x)
        return tuple(sorted(found))


@dataclass(frozen=True)
class PriceProfile:
    """The full threshold-dependent strict halfspaces, not a scalar MAF."""
    inputs: tuple[tuple[F, ...], ...]   # species rows, reaction columns
    outputs: tuple[tuple[F, ...], ...]

    def __post_init__(self):
        for name in ("inputs", "outputs"):
            object.__setattr__(self, name, tuple(tuple(rational(x) for x in row)
                                                for row in getattr(self, name)))
        if (not self.inputs or not self.inputs[0] or len(self.inputs) != len(self.outputs)
                or any(len(row) != len(self.inputs[0]) for row in self.inputs + self.outputs)
                or any(x < 0 for row in self.inputs + self.outputs for x in row)):
            raise ValueError("Nonempty equally shaped nonnegative source matrices required")

    def compose(self, other):
        """Shared-species composition; both profiles use the same species order."""
        if len(self.inputs) != len(other.inputs):
            raise ValueError("Shared species dimensions differ")
        return PriceProfile(tuple(a + b for a, b in zip(self.inputs, other.inputs)),
                            tuple(a + b for a, b in zip(self.outputs, other.outputs)))

    def residual(self, q):
        q = rational(q)
        return tuple(tuple(b - q * a for a, b in zip(ar, br))
                     for ar, br in zip(self.inputs, self.outputs))

    def contains(self, q, prices):
        prices = tuple(map(rational, prices))
        m = self.residual(q)
        return (len(prices) == len(m) and all(p >= 0 for p in prices)
                and all(dot(column, prices) < 0 for column in zip(*m)))

    def decide(self, q):
        """Exactly return either normalized feasible flux or STRICT price witness.

        Weak primal inequalities retain the threshold endpoint. A weak zero price
        residual is never accepted as an impossibility certificate.
        """
        m = self.residual(q)
        ns, nr = len(m), len(m[0])
        positive = tuple((tuple(F(i == j) for i in range(nr)), F(0)) for j in range(nr))
        primal = RationalPolytope((F(1),) * nr, tuple((r, F(0)) for r in m) + positive)
        vertices = primal.vertices()
        if vertices:
            x = vertices[0]
            return {"kind": "feasible", "flux": x, "residual": tuple(dot(r, x) for r in m)}
        # Maximize epsilon with p>=0, sum(p)=1 and -(B-qA)^T p >= epsilon.
        bounds = tuple((tuple(-v for v in col) + (F(-1),), F(0)) for col in zip(*m))
        bounds += tuple((tuple(F(i == j) for i in range(ns + 1)), F(0)) for j in range(ns + 1))
        candidates = RationalPolytope((F(1),) * ns + (F(0),), bounds).vertices()
        best = max(candidates, key=lambda x: x[-1], default=None)
        if best is None or best[-1] <= 0 or not self.contains(q, best[:-1]):
            raise ArithmeticError("No verified alternative; never infer infeasibility from a missing witness")
        return {"kind": "infeasible", "prices": best[:-1], "strict_margin": best[-1]}


@dataclass(frozen=True)
class TwoPort:
    """j_left=(forward+left_leak)*X-backward*Y;
    j_right=forward*X-(backward+right_leak)*Y.
    """
    forward: F
    backward: F
    left_leak: F = F(0)
    right_leak: F = F(0)

    def currents(self, left, right):
        left, right = rational(left), rational(right)
        return ((self.forward + self.left_leak) * left - self.backward * right,
                self.forward * left - (self.backward + self.right_leak) * right)

    def extend(self, forward, backward, degradation):
        forward, backward, degradation = map(rational, (forward, backward, degradation))
        loss = self.backward + self.right_leak + forward + degradation
        return TwoPort(self.forward * forward / loss, self.backward * backward / loss,
                       self.left_leak + self.forward * (self.right_leak + degradation) / loss,
                       backward * (self.right_leak + degradation) / loss)


@dataclass(frozen=True)
class PassiveTail:
    forward: tuple[F, ...]
    backward: tuple[F, ...]
    degradation: tuple[F, ...]

    def __post_init__(self):
        for name in ("forward", "backward", "degradation"):
            object.__setattr__(self, name, tuple(map(rational, getattr(self, name))))
        if (not self.forward or len(self.backward) != len(self.forward)
                or len(self.degradation) != len(self.forward) - 1
                or min(self.forward + self.backward) <= 0
                or any(d < 0 for d in self.degradation)):
            raise ValueError("Positive reversible rates and one nonnegative loss per internal node required")

    def prefixes(self):
        ports = [TwoPort(self.forward[0], self.backward[0])]
        for c, b, d in zip(self.forward[1:], self.backward[1:], self.degradation):
            ports.append(ports[-1].extend(c, b, d))
        return tuple(ports)

    def port(self):
        return self.prefixes()[-1]

    def reconstruct(self, left, right):
        """Unique internal steady concentrations by reversing exact elimination."""
        left, right = rational(left), rational(right)
        if min(left, right) < 0:
            raise ValueError("Concentrations must be nonnegative")
        ports = self.prefixes()
        states = [right]
        for i in range(len(self.forward) - 1, 0, -1):
            p = ports[i - 1]
            loss = p.backward + p.right_leak + self.forward[i] + self.degradation[i - 1]
            states.append((p.forward * left + self.backward[i] * states[-1]) / loss)
        return (left,) + tuple(reversed(states))

    def literal_currents(self, states):
        if len(states) != len(self.forward) + 1:
            raise ValueError("Wrong number of chain states")
        return tuple(c * rational(x) - b * rational(y)
                     for c, b, x, y in zip(self.forward, self.backward, states[:-1], states[1:]))

    def audit(self, left, right):
        states = self.reconstruct(left, right)
        currents = self.literal_currents(states)
        residuals = tuple(currents[i] - currents[i + 1] - d * states[i + 1]
                          for i, d in enumerate(self.degradation))
        boundary = self.port().currents(left, right)
        if any(residuals) or boundary != (currents[0], currents[-1]):
            raise ArithmeticError("Exact two-port audit failed")
        return {"states": states, "edge_currents": currents, "internal_residuals": residuals,
                "boundary_currents": boundary, "internal_loss": sum(d*x for d, x in zip(self.degradation, states[1:-1]))}


@dataclass(frozen=True)
class Branch:
    internal: tuple[tuple[F, ...], ...]
    to_root: tuple[F, ...]
    from_root: tuple[F, ...]

    def __post_init__(self):
        object.__setattr__(self, "internal", tuple(tuple(map(rational, r)) for r in self.internal))
        for name in ("to_root", "from_root"):
            object.__setattr__(self, name, tuple(map(rational, getattr(self, name))))
        n = len(self.internal)
        if (not n or any(len(r) != n for r in self.internal)
                or len(self.to_root) != n or len(self.from_root) != n
                or any(x < 0 for x in self.to_root) or any(x <= 0 for x in self.from_root)
                or any(self.internal[i][j] < 0 for i in range(n) for j in range(n) if i != j)):
            raise ValueError("Metzler branch, nonnegative return and strictly positive source coupling required")
        self.response()

    def response(self):
        u = solve(self.internal, tuple(-c for c in self.from_root))
        if u is None or any(x <= 0 for x in u):
            raise ValueError("Branch lacks a positive solve witness A*u=-c")
        return u

    def load(self):
        return dot(self.to_root, self.response())


@dataclass(frozen=True)
class DegradationStar:
    root_diagonal: F
    branches: tuple[Branch, ...]

    def __post_init__(self):
        object.__setattr__(self, "root_diagonal", rational(self.root_diagonal))
        if not self.branches:
            raise ValueError("Supply at least one branch")

    def matrix(self):
        n = 1 + sum(len(b.internal) for b in self.branches)
        m = [[F(0) for _ in range(n)] for _ in range(n)]
        m[0][0] = self.root_diagonal
        offset = 1
        for b in self.branches:
            for i, row in enumerate(b.internal):
                m[0][offset+i] = b.to_root[i]
                m[offset+i][0] = b.from_root[i]
                m[offset+i][offset:offset+len(row)] = row
            offset += len(b.internal)
        return tuple(map(tuple, m))

    def certificate(self):
        kappa = self.root_diagonal + sum((b.load() for b in self.branches), F(0))
        sign = (kappa > 0) - (kappa < 0)
        delta = sign * min(F(1, 2), abs(kappa) / (2 * (1 + abs(self.root_diagonal))))
        v = (1 + delta,) + tuple(x for b in self.branches for x in b.response())
        residual = tuple(dot(r, v) for r in self.matrix())
        if not all((x > 0) - (x < 0) == sign for x in residual):
            raise ArithmeticError("Failed strict sign certificate")
        return {"state": {1: "growth", 0: "critical", -1: "extinction"}[sign],
                "schur_load": kappa, "positive_vector": v, "matrix_times_vector": residual}


@dataclass(frozen=True)
class FiniteRelation:
    """Extensional relation: out-of-scope inspection is impossible by construction."""
    scope: tuple[str, ...]
    rows: frozenset[tuple[Hashable, ...]]

    def __post_init__(self):
        object.__setattr__(self, "scope", tuple(self.scope))
        object.__setattr__(self, "rows", frozenset(tuple(row) for row in self.rows))
        if len(set(self.scope)) != len(self.scope) or any(len(r) != len(self.scope) for r in self.rows):
            raise ValueError("Unique variables and full rows required")

    def accepts(self, assignment):
        return tuple(assignment[v] for v in self.scope) in self.rows

    def project(self, boundary):
        boundary = tuple(boundary)
        indices = [self.scope.index(v) for v in boundary]
        return FiniteRelation(boundary, frozenset(tuple(row[i] for i in indices) for row in self.rows))

    def compatible(self, context):
        if self.scope != context.scope:
            raise ValueError("Context must use the same ordered boundary")
        return bool(self.rows & context.rows)

    def separating_context(self, other):
        if self.scope != other.scope:
            raise ValueError("Different boundaries")
        difference = self.rows ^ other.rows
        return None if not difference else FiniteRelation(self.scope, frozenset((min(difference),)))


@dataclass(frozen=True)
class JoinNode:
    name: str
    bag: tuple[str, ...]
    separator: tuple[str, ...]
    local: FiniteRelation
    children: tuple[JoinNode, ...] = ()

    def __post_init__(self):
        if not set(self.local.scope) <= set(self.bag) or not set(self.separator) <= set(self.bag):
            raise ValueError("Local scope and separator must lie in the bag")
        if len(set(self.bag)) != len(self.bag) or len(set(self.separator)) != len(self.separator):
            raise ValueError("Duplicate variable")

    def all_variables(self):
        return set(self.bag).union(*(c.all_variables() for c in self.children))

    def validate(self):
        for i, child in enumerate(self.children):
            others = set(self.bag).union(*(c.all_variables() for j, c in enumerate(self.children) if i != j))
            if not child.all_variables() & others <= set(child.separator):
                raise ValueError(f"Running intersection fails at {child.name}")
            if not set(child.separator) <= set(self.bag):
                raise ValueError("This finite implementation requires child separators in the parent bag")
            child.validate()

    def message(self, domain):
        """Exact finite separator relation, refusing invalid decompositions."""
        self.validate()
        return self._message(tuple(domain))

    def _message(self, domain):
        children = [c._message(domain) for c in self.children]
        rows = set()
        for values in product(domain, repeat=len(self.bag)):
            assignment = dict(zip(self.bag, values))
            if self.local.accepts(assignment) and all(c.accepts(assignment) for c in children):
                rows.add(tuple(assignment[v] for v in self.separator))
        return FiniteRelation(self.separator, frozenset(rows))

    def global_trace(self, domain):
        """Independent exhaustive oracle; one assignment for the whole subtree."""
        variables = tuple(sorted(self.all_variables()))
        def holds(node, assignment):
            return node.local.accepts(assignment) and all(holds(c, assignment) for c in node.children)
        rows = set()
        for values in product(domain, repeat=len(variables)):
            assignment = dict(zip(variables, values))
            if holds(self, assignment):
                rows.add(tuple(assignment[v] for v in self.separator))
        return FiniteRelation(self.separator, frozenset(rows))

    def width(self):
        return max([len(self.separator)] + [c.width() for c in self.children])


def paper_profiles():
    return (PriceProfile(((1, 0), (0, 1)), ((0, 4), (1, 0))),
            PriceProfile(((1, 0), (0, 1)), ((0, 1), (4, 0))))


def gluing_examples(domain=FINITE_DOMAIN):
    if set(domain) != {0, 1}:
        raise ValueError("The worked gluing example requires domain {0,1}; JoinNode accepts other finite domains")
    truth = FiniteRelation((), frozenset(((),)))
    left = FiniteRelation(("x", "a"), frozenset((x, x) for x in domain))
    right = FiniteRelation(("x", "b"), frozenset((x, 1-x) for x in domain))
    good = JoinNode("root", ("x",), ("x",), truth,
                    (JoinNode("left", ("x", "a"), ("x",), left),
                     JoinNode("right", ("x", "b"), ("x",), right)))
    bad = JoinNode("root", (), (), truth,
                  (JoinNode("left", ("x",), (), FiniteRelation(("x",), frozenset(((0,),)))),
                   JoinNode("right", ("x",), (), FiniteRelation(("x",), frozenset(((1,),))))))
    # A corrected decomposition routes the conflicting x through both separators.
    repaired = JoinNode("root", ("x",), (), truth,
                       tuple(JoinNode(c.name, c.bag, ("x",), c.local) for c in bad.children))
    return good, bad, repaired


def serialize(value):
    if isinstance(value, F):
        return str(value)
    if isinstance(value, dict):
        return {str(k): serialize(v) for k, v in value.items()}
    if isinstance(value, (tuple, list, set, frozenset)):
        return [serialize(v) for v in value]
    return value


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)
    n1, n2 = paper_profiles()
    profiles = {"N1": n1, "N2": n2, "N1+N1": n1.compose(n1), "N2+N1": n2.compose(n1)}
    decisions = {name: {q: model.decide(q) for q in AMPLIFICATION_THRESHOLDS}
                 for name, model in profiles.items()}
    tail = PassiveTail(TAIL_FORWARD, TAIL_BACKWARD, TAIL_DEGRADATION)
    audit = tail.audit(LEFT_CONCENTRATION, RIGHT_CONCENTRATION)
    branch = Branch(((BRANCH_DIAGONAL,),), (BRANCH_TO_ROOT,), (ROOT_TO_BRANCH,))
    stars = {n: DegradationStar(ROOT_DIAGONAL, (branch,) * n).certificate() for n in BRANCH_COUNTS}
    good, bad, repaired = gluing_examples()
    message = good.message(FINITE_DOMAIN)
    assert message == good.global_trace(FINITE_DOMAIN)
    try:
        bad.message(FINITE_DOMAIN)
        raise AssertionError("Invalid tree accepted")
    except ValueError as error:
        rejection = str(error)
    assert not repaired.message(FINITE_DOMAIN).rows and not bad.global_trace(FINITE_DOMAIN).rows
    # Trace cardinality loses information; a singleton context separates these.
    t0 = FiniteRelation(("x",), frozenset(((0,),)))
    t1 = FiniteRelation(("x",), frozenset(((1,),)))
    context = t0.separating_context(t1)
    results = {"amplification": decisions, "tail": {"coefficients": vars(tail.port()), **audit},
               "stars": stars, "finite_traces": {"width": good.width(), "message": sorted(message.rows),
               "invalid_tree_rejection": rejection, "repaired_conflict_trace": [],
               "same_cardinality": len(t0.rows) == len(t1.rows), "separating_context": sorted(context.rows),
               "context_verdicts": [t0.compatible(context), t1.compatible(context)]}}
    (args.output / "results.json").write_text(json.dumps(serialize(results), indent=2) + "\n")
    sweep = []
    for numerator in range(41):
        scale = F(numerator, 10)
        candidate = PassiveTail(tail.forward, tail.backward, tuple(scale*d for d in tail.degradation))
        result = candidate.audit(LEFT_CONCENTRATION, RIGHT_CONCENTRATION)
        sweep.append((scale, *result["boundary_currents"], result["internal_loss"],
                      candidate.port().left_leak, candidate.port().right_leak))
    with (args.output / "tail_sweep.csv").open("w", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(("degradation_multiplier", "left_current", "right_current", "internal_loss", "left_leak", "right_leak"))
        writer.writerows(sweep)
    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": "exact-interfaces-v1"})
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.8), layout="constrained")
    ratio = [F(i, 100) for i in range(401)]
    q = F(3)
    for label, model, y, color in (("N1", n1, 1, "#187f91"), ("N2", n2, 0, "#ae6539")):
        valid = [r for r in ratio if model.contains(q, (1, r))]
        axes[0].plot([float(r) for r in valid], [y] * len(valid), lw=8, color=color)
    axes[0].scatter([F(4, 3), 3, F(1, 3), F(3, 4)], [1, 1, 0, 0], s=65,
                    facecolors="white", edgecolors=["#187f91"]*2+["#ae6539"]*2, zorder=3)
    axes[0].set(yticks=[0, 1], yticklabels=["N2", "N1"], ylim=(-.5, 1.5), xlim=(0, 3.5),
                xlabel="Price ratio pB / pA", title="At q = 3: separate certificates, no intersection")
    axes[0].text(1.7, .45, "N2 + N1 is feasible", ha="center")
    colors = ["#ae6539" if s["schur_load"] < 0 else "#187f91" for s in stars.values()]
    axes[1].bar(list(stars), [float(s["schur_load"]) for s in stars.values()], color=colors)
    axes[1].axhline(0, color="#333333", lw=1)
    axes[1].set(xlabel="Branches sharing one root", ylabel="Schur load (inverse time)",
                xticks=list(stars), title="Branch loads add: extinction can become growth")
    for extension in ("png", "svg"):
        fig.savefig(args.output / f"interfaces.{extension}", dpi=165)
    plt.close(fig)
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.8), layout="constrained")
    axes[0].plot(range(len(audit["states"])), list(map(float, audit["states"])), "o-", color="#187f91")
    axes[0].set(xlabel="Chain node (endpoints fixed)", ylabel="Steady concentration", title="Steady concentrations along the reaction\nchain")
    axes[1].plot([float(row[0]) for row in sweep], [float(row[1]) for row in sweep], label="Current leaving left", color="#187f91")
    axes[1].plot([float(row[0]) for row in sweep], [float(row[2]) for row in sweep], label="Current entering right", color="#ae6539")
    axes[1].axhline(0, color="#444444", lw=.8)
    axes[1].set(xlabel="Internal degradation multiplier", ylabel="Current (concentration / time)", title="Boundary currents and internal degradation")
    axes[1].legend(frameon=False)
    for extension in ("png", "svg"):
        fig.savefig(args.output / f"passive_tail.{extension}", dpi=165)
    plt.close(fig)
    console = "\n".join([
        "Exact interfaces for modular autocatalytic networks",
        "All certificates, chain balances and finite messages use rational/finite arithmetic.",
        "At q=4: N1+N1 " + profiles["N1+N1"].decide(4)["kind"] + "; N2+N1 " + profiles["N2+N1"].decide(4)["kind"],
        "Two-port coefficients: " + str(serialize(vars(tail.port()))),
        "Internal steady residuals: " + str(audit["internal_residuals"]),
        "Star states: " + str({n: r["state"] for n, r in stars.items()}),
        "Valid finite message: " + str(sorted(message.rows)),
        "Invalid join tree: " + rejection,
        "Cardinality-summary collision, context verdicts: " + str(results["finite_traces"]["context_verdicts"]),
        "Not a re-execution of the Lean proofs; not a nonlinear kinetic model."])
    print(console)
    (args.output / "console.txt").write_text(console + "\n")
    outputs = {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"}
    metadata = {"source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                "manuscript_sha256": "62106305bc8d8f4652be7989329b13ed300007da7596d596e7e0cdc1ea52ff61",
                "python": platform.python_version(), "matplotlib": matplotlib.__version__,
                "inputs": {"thresholds": AMPLIFICATION_THRESHOLDS, "tail_forward": TAIL_FORWARD,
                           "tail_backward": TAIL_BACKWARD, "tail_degradation": TAIL_DEGRADATION,
                           "left_concentration": LEFT_CONCENTRATION, "right_concentration": RIGHT_CONCENTRATION,
                           "root_diagonal": ROOT_DIAGONAL, "branch_diagonal": BRANCH_DIAGONAL,
                           "root_to_branch": ROOT_TO_BRANCH, "branch_to_root": BRANCH_TO_ROOT,
                           "branch_counts": BRANCH_COUNTS, "finite_domain": FINITE_DOMAIN},
                "output_sha256": outputs, "arithmetic": "exact Fraction / finite sets; float only for plots"}
    (args.output / "run_metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")


if __name__ == "__main__":
    main()