#!/usr/bin/env python3
"""Executable companion to A Nontrivial Critical Window for RAF Emergence.

Run: python critical_window.py --output results --trials 400 --seed 20260908
Requires Python 3.10+, NumPy and Matplotlib. No input files or network needed.

The finite RAF simulations use the exact distributional reduction in Lemma 5.1.
They do not approximate the infinite-size limit by a fitted curve. Independent
reaction marks are used separately for the infinite model's finite escape events.
The paper, not a finite simulation, proves positivity, continuity and convergence.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import platform
from collections import Counter
from dataclasses import dataclass
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


@dataclass
class Polymer:
    n: int

    def __post_init__(self):
        if not 2 <= self.n <= 16:
            raise ValueError("Use an integer cap between 2 and 16.")
        self.words = [format(j, f"0{length}b")
                      for length in range(1, self.n + 1)
                      for j in range(2**length)]
        self.index = {w: i for i, w in enumerate(self.words)}
        self.reactions = [(w, cut) for w in self.words
                          for cut in range(1, len(w))]
        self.rindex = {r: i for i, r in enumerate(self.reactions)}
        self.U, self.R = len(self.words), len(self.reactions)
        self.u = np.array([self.index[w[:i]] for w, i in self.reactions])
        self.v = np.array([self.index[w[i:]] for w, i in self.reactions])
        self.w = np.array([self.index[w] for w, i in self.reactions])
        self.length = np.array([len(w) for w in self.words])
        assert self.U == 2**(self.n + 1) - 2
        assert self.R == (self.n - 2)*2**(self.n + 1) + 4

    def closure(self, active):
        """Food closure: join when BOTH factors exist; split when product exists.

        Catalysts are not prerequisites during this construction. They determine
        the active reaction set before the closure is computed.
        """
        u, v, w = self.u[active], self.v[active], self.w[active]
        have = self.length <= 2
        while True:
            before = int(have.sum())
            can_join = have[u] & have[v]
            can_split = have[w]
            have[w[can_join]] = True
            have[u[can_split]] = True
            have[v[can_split]] = True
            if int(have.sum()) == before:
                return have

    def usable(self, have):
        return have[self.w] | (have[self.u] & have[self.v])


def catalysis_probability(model, intensity):
    return min(1.0, max(0.0, intensity*model.n/model.R))


def first_catalyst_ranks(uniforms, p, U):
    """One fixed random rank per reaction, reused throughout pruning.

    P(rank <= k) = 1-(1-p)**k. Rank U+1 means no catalyst. These are
    first successes in an ordered Bernoulli column, not actual molecule labels.
    Reusing marks is essential; redrawing each pruning round is incorrect.
    """
    if p <= 0:
        return np.full(uniforms.shape, U + 1, dtype=np.int64)
    if p >= 1:
        return np.ones(uniforms.shape, dtype=np.int64)
    value = np.floor(np.log1p(-uniforms)/math.log1p(-p)) + 1
    return np.minimum(value, U + 1).astype(np.int64)


def canonical_raf(model, ranks):
    """Lemma 5.1: exact finite RAF-existence law using only pool sizes.

    A sampled rank history is NOT a witness in a particular literal catalysis
    matrix. The separate literal_raf routine returns such witnesses.
    """
    k = model.U
    while True:
        active = ranks <= k
        have = model.closure(active)
        new_k = int(have.sum())
        assert new_k <= k
        if new_k == k:
            # A food-only RAF counts: six retained molecules need not mean failure.
            return bool(np.any(active & model.usable(have))), new_k
        k = new_k


def literal_raf(model, catalyst_molecules, catalyst_reactions):
    """Prune an explicit list of (molecule, reaction) catalytic assignments."""
    pool = np.ones(model.U, dtype=bool)
    history = []
    while True:
        active = np.zeros(model.R, dtype=bool)
        active[catalyst_reactions[pool[catalyst_molecules]]] = True
        have = model.closure(active)
        history.append({"candidate_molecules": int(pool.sum()),
                        "active_reactions": int(active.sum()),
                        "reachable_molecules": int(have.sum())})
        assert np.all(~have | pool)
        if np.array_equal(pool, have):
            retained = np.flatnonzero(active & model.usable(have))
            # Recompute from the returned set, then check Definition 2.1 literally.
            keep = np.zeros(model.R, dtype=bool)
            keep[retained] = True
            verified = model.closure(keep)
            assert np.array_equal(have, verified)
            for r in retained:
                assert model.usable(verified)[r]
                assert np.any(verified[catalyst_molecules[catalyst_reactions == r]])
            return retained, have, history
        pool = have


def sample_literal(model, p, rng):
    # Conditional on their count, iid Bernoulli successes form a uniform subset.
    count = int(rng.binomial(model.U*model.R, p))
    positions = rng.choice(model.U*model.R, count, replace=False)
    return positions // model.R, positions % model.R


def wilson(successes, trials):
    """Pointwise approximate 95% binomial interval, including 0/N and N/N."""
    z = 1.959963984540054
    phat = successes/trials
    denominator = 1 + z*z/trials
    center = (phat + z*z/(2*trials))/denominator
    radius = z*math.sqrt(phat*(1-phat)/trials + z*z/(4*trials**2))/denominator
    return min(phat, max(0.0, center-radius)), max(phat, min(1.0, center+radius))


def probability_row(successes, trials, **fields):
    lo, hi = wilson(successes, trials)
    return dict(fields, successes=int(successes), trials=trials,
                estimate=successes/trials, ci_low=lo, ci_high=hi)


def exact_history_check():
    """Exhaustively check Lemma 5.1 on a small monotone closure operator.

    Teaching test only: three abstract species, food {0}, reactions 0->2, 2->1.
    All 64 catalysis matrices are enumerated. Counts agree separately for each
    number of present coordinates, proving equality for every Bernoulli p in
    THIS small example. This is not an enumeration of the full polymer model.
    """
    def phi(active):
        result = {0}
        if active & 1:
            result.add(2)
            if active & 2:
                result.add(1)
        return result

    distributions = [Counter(), Counter()]
    different_paths = 0
    for bits in range(64):
        histories = []
        for canonical in (False, True):
            pool = {0, 1, 2}
            trace = []
            for _ in range(4):
                active = sum(1 << r for r in range(2)
                             if any(bits & (1 << (2*x+r)) for x in pool))
                trace.append(active)
                actual = phi(active)
                pool = set(range(len(actual))) if canonical else actual
            key = (tuple(trace), bits.bit_count())
            distributions[int(canonical)][key] += 1
            histories.append(tuple(trace))
        different_paths += histories[0] != histories[1]
    assert distributions[0] == distributions[1]
    return {"matrices_checked": 64, "history_polynomials_equal": True,
            "individual_matrix_histories_differ": different_paths}


def worked_examples():
    model = Polymer(4)
    def run(assignments):
        xs = np.array([model.index[x] for x, r in assignments], dtype=int)
        rs = np.array([model.rindex[r] for x, r in assignments], dtype=int)
        retained, have, trace = literal_raf(model, xs, rs)
        return {"assignments": [{"catalyst": x, "product": r[0], "cut": r[1]}
                                for x, r in assignments],
                "RAF_exists": bool(len(retained)), "history": trace,
                "retained_reactions": [list(model.reactions[r]) for r in retained],
                "reachable_words": [model.words[i] for i in np.flatnonzero(have)]}

    # All unlisted catalysis coordinates are absent; full n=4 reaction universe.
    lost = run([("1111", ("000", 1)), ("000", ("0000", 3))])
    mutual = run([("0000", ("000", 1)), ("000", ("0000", 3))])
    food_only = run([("0", ("00", 1))])
    assert not lost["RAF_exists"] and mutual["RAF_exists"] and food_only["RAF_exists"]
    assert [r["reachable_molecules"] for r in lost["history"]] == [8, 6, 6]
    assert [r["reachable_molecules"] for r in mutual["history"]] == [8, 8]

    # Lemma 8.2: make target 101 from an already generated word 000.
    repair_model = Polymer(6)
    repairs = [("0001", 3), ("00010", 4), ("000101", 5), ("000101", 3)]
    active = np.zeros(repair_model.R, dtype=bool)
    for r in [("000", 1)] + repairs:
        active[repair_model.rindex[r]] = True
    assert all(len(w) > 3 for w, cut in repairs)
    assert len(set(repairs)) == 4
    assert repair_model.closure(active)[repair_model.index["101"]]
    return {"lost_catalyst": lost, "mutual_support": mutual,
            "food_only": food_only,
            "append_then_split": {"starting_word": "000", "target": "101",
                                  "reactions": repairs, "target_verified": True,
                                  "all_open_probability": "a^4"}}


def exact_positive_bound():
    """Verify the numerical condition of Theorem 7.1 exactly at a=1/2, m=2.

    theta <= 1/7 implies theta/(1-theta) <= 1/6. Proposition 7.2 then
    gives S(1/2) >= 2^(-(R_L+1)). The contour theorem itself is in the paper.
    Keep this tiny positive number symbolic to avoid floating-point underflow.
    """
    A = 2*7**64*81
    k = 1
    while 7*A*3**k > 4**k:
        k += 1
    assert 7*A*3**k <= 4**k
    assert k == 1 or 7*A*3**(k-1) > 4**(k-1)
    L = 10*k
    return {"a": "1/2", "lambda": "ln(2)", "m": 2, "k": k, "L": L,
            "exact_condition_verified": "7 * (2*7^64*81) * 3^k <= 4^k",
            "positive_lower_bound": f"2^(-(({L}-2)*2^({L}+1)+5))",
            "upper_bound_from_36_gateways": "1 - 2^(-36)",
            "meaning": "Bounds use Theorem 7.1 and Proposition 7.2; not a numerical estimate."}


def simulate_finite(intensities, caps, trials, seed):
    rows = []
    for n in caps:
        model = Polymer(n)
        rng = np.random.default_rng(np.random.SeedSequence(seed, spawn_key=(1, n)))
        marks = rng.random((trials, model.R))
        previous = np.zeros(trials, dtype=bool)
        for intensity in intensities:
            p = catalysis_probability(model, intensity)
            ranks = first_catalyst_ranks(marks, p, model.U)
            outcomes = np.array([canonical_raf(model, r)[0] for r in ranks])
            assert np.all(~previous | outcomes)  # common marks preserve inclusion
            previous = outcomes
            log_failure = -math.inf if p == 1 else 36*model.U*math.log1p(-p)
            rows.append(probability_row(int(outcomes.sum()), trials, n=n,
                        intensity=float(intensity), p=p, molecule_count=model.U,
                        reaction_count=model.R, finite_gateway_ceiling=-math.expm1(log_failure)))
        print(f"Finished finite RAF simulations: n={n}", flush=True)
    return rows


def simulate_escape(intensities, boundaries, trials, seed):
    """H_B: produce a word longer than B, allowing intermediate length <=2B.

    Lemma 3.1 says this exactly captures escape beyond B in the infinite field.
    H_B decreases to infinite survival as B grows. Its probability is an UPPER
    bound on S(a); our measured frequency estimates that bound, with sampling error.
    """
    largest = Polymer(2*max(boundaries))
    rng = np.random.default_rng(np.random.SeedSequence(seed, spawn_key=(2,)))
    marks = rng.random((trials, largest.R))
    rows = []
    for intensity in intensities:
        a = -math.expm1(-intensity)
        previous = np.ones(trials, dtype=bool)
        for boundary in boundaries:
            model = Polymer(2*boundary)
            outcomes = np.array([np.any(model.closure(u[:model.R] < a)
                                       & (model.length > boundary)) for u in marks])
            assert np.all(~outcomes | previous)
            previous = outcomes
            rows.append(probability_row(int(outcomes.sum()), trials,
                        boundary=boundary, cap=2*boundary,
                        intensity=float(intensity), openness=a))
        print(f"Finished static escape simulations: lambda={intensity:g}", flush=True)
    return rows


def reference_comparison(trials, seed):
    model = Polymer(6)
    rng = np.random.default_rng(np.random.SeedSequence(seed, spawn_key=(3,)))
    successes = [0, 0]
    p = catalysis_probability(model, 0.2)
    for _ in range(trials):
        xs, rs = sample_literal(model, p, rng)
        successes[0] += bool(len(literal_raf(model, xs, rs)[0]))
        ranks = first_catalyst_ranks(rng.random(model.R), p, model.U)
        successes[1] += canonical_raf(model, ranks)[0]
    return [probability_row(s, trials, method=method, n=6, intensity=0.2)
            for method, s in zip(("literal catalysis", "cardinality replacement"), successes)]


def static_barrier_check(trials, seed):
    """Check Lemma 6.1 in its canonical coupling, using k=U/2 > six food words."""
    model = Polymer(6)
    k = model.U // 2
    assert k > 6
    p = catalysis_probability(model, .2)
    rng = np.random.default_rng(np.random.SeedSequence(seed, spawn_key=(4,)))
    barrier_count = raf_count = 0
    for _ in range(trials):
        ranks = first_catalyst_ranks(rng.random(model.R), p, model.U)
        barrier = int(model.closure(ranks <= k).sum()) >= k
        raf, _ = canonical_raf(model, ranks)
        assert not barrier or raf
        barrier_count += barrier
        raf_count += raf
    return {"n": 6, "intensity": .2, "retained_pool_size_k": k,
            "static_openness": -math.expm1(k*math.log1p(-p)), "trials": trials,
            "barrier_successes": barrier_count, "RAF_successes": raf_count,
            "violations": 0,
            "scope": "Samplewise implication in canonical coupling; law transfers by Lemma 5.1."}


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


def make_figures(out, finite, escape, examples):
    plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 11,
                         "axes.spines.top": False, "axes.spines.right": False,
                         "axes.titleweight": "bold", "axes.labelcolor": "#263444",
                         "text.color": "#263444", "axes.edgecolor": "#AAB4BF",
                         "figure.facecolor": "#FAFAF7", "axes.facecolor": "#FAFAF7",
                         "savefig.facecolor": "#FAFAF7"})
    colors = ["#305879", "#147D83", "#B2742B", "#845577"]
    def save(fig, name):
        fig.savefig(out / f"{name}.png", dpi=180, bbox_inches="tight")
        fig.savefig(out / f"{name}.svg", bbox_inches="tight")
        plt.close(fig)

    fig, axes = plt.subplots(1, 2, figsize=(13.8, 5.6))
    for color, n in zip(colors, sorted({r["n"] for r in finite})):
        data = [r for r in finite if r["n"] == n and r["intensity"] <= .35]
        axes[0].errorbar([r["intensity"] for r in data], [r["estimate"] for r in data],
                        yerr=[[r["estimate"]-r["ci_low"] for r in data],
                              [r["ci_high"]-r["estimate"] for r in data]],
                        label=f"Maximum length {n}", color=color, marker="o", ms=3, lw=1.5, capsize=2)
    for color, boundary in zip(colors, sorted({r["boundary"] for r in escape})):
        data = [r for r in escape if r["boundary"] == boundary and r["intensity"] <= .35]
        axes[1].plot([r["intensity"] for r in data], [r["estimate"] for r in data],
                     label=f"Reach length > {boundary}", color=color, marker="o", ms=3, lw=1.5)
        axes[1].fill_between([r["intensity"] for r in data], [r["ci_low"] for r in data],
                            [r["ci_high"] for r in data], color=color, alpha=.08)
    axes[0].set_title("Does a finite system contain a RAF?", loc="left", pad=14)
    axes[1].set_title("Does the open-reaction closure escape?", loc="left", pad=14)
    axes[0].set_ylabel("Fraction of samples containing a RAF")
    axes[1].set_ylabel("Fraction of samples reaching beyond the boundary")
    for ax in axes:
        ax.set_xlabel(r"Catalysis intensity $\lambda$ (average catalysis / maximum length)")
        ax.set_ylim(-.025, 1.025)
        ax.set_xlim(-.005, .36)
        ax.grid(axis="y", alpha=.2)
        ax.legend(loc="lower right", fontsize=9, frameon=False)
    fig.text(.07, .01, "Finite samples with pointwise 95% intervals; view enlarged to 0–0.35. Neither panel plots the unknown limiting curve.", fontsize=10)
    fig.subplots_adjust(bottom=.19, wspace=.28)
    save(fig, "probability_comparison")

    caps = sorted({r["n"] for r in finite})
    intensities = sorted({r["intensity"] for r in finite})
    grid = np.array([[next(r["estimate"] for r in finite if r["n"] == n
                          and r["intensity"] == x) for x in intensities] for n in caps])
    fig, ax = plt.subplots(figsize=(13, 4.7))
    im = ax.imshow(grid, vmin=0, vmax=1, cmap="cividis", aspect="auto", origin="lower")
    ax.set_xticks(range(len(intensities)), [f"{x:g}" for x in intensities])
    ax.set_yticks(range(len(caps)), caps)
    ax.set_xlabel(r"Catalysis intensity $\lambda$ (displayed columns are separate sampled values)")
    ax.set_ylabel("Maximum molecule length")
    ax.set_title("Finite RAF probability across size and catalysis", loc="left", pad=18)
    for i in range(len(caps)):
        for j in range(len(intensities)):
            value = grid[i, j]
            label = "<1%" if 0 < value < .01 else ">99%" if .99 < value < 1 else f"{100*value:.0f}%"
            ax.text(j, i, label, ha="center", va="center",
                    color="white" if grid[i,j] < .48 else "#172530", fontsize=9)
    fig.colorbar(im, ax=ax, label="Fraction of samples containing a RAF", pad=.025)
    fig.text(.1, .01, "A finite-size probability map, not a proved phase boundary. A displayed 0% or 100% is a sample result.", fontsize=10)
    fig.subplots_adjust(bottom=.2)
    save(fig, "finite_size_map")

    food = ["0", "1", "00", "01", "10", "11"]
    fig, axes = plt.subplots(1, 2, figsize=(12.5, 5.6), gridspec_kw={"width_ratios": [1, 1.1]})
    grid = np.ones((6, 6)); grid[:2, :2] = 0
    axes[0].imshow(grid, cmap=matplotlib.colors.ListedColormap(["#DDD8CB", "#C9E6E1"]), vmin=0, vmax=1)
    for i, u in enumerate(food):
        for j, v in enumerate(food):
            axes[0].text(j, i, f"{u}|{v}", ha="center", va="center", fontsize=10)
    axes[0].set_xticks(range(6), food); axes[0].set_yticks(range(6), food)
    axes[0].set_xlabel("Second food word"); axes[0].set_ylabel("First food word")
    axes[0].set_title("36 reactions usable directly from food", loc="left", pad=15)
    x = np.linspace(0, .2, 300)
    axes[1].plot(x, -np.expm1(-36*x), color=colors[0], lw=2.5)
    axes[1].fill_between(x, -np.expm1(-36*x), 1, color="#E9DACA", alpha=.7)
    axes[1].text(.003, .98, "Excluded\nprobabilities", fontsize=10, va="top")
    axes[1].text(.075, .32, "The actual limiting probability\ncan be below the ceiling.", fontsize=11)
    axes[1].set_title("A proved ceiling on the limiting probability", loc="left", pad=15)
    axes[1].set_xlabel(r"Catalysis intensity $\lambda$"); axes[1].set_ylabel("Probability")
    axes[1].set_ylim(0, 1.02); axes[1].grid(axis="y", alpha=.2)
    axes[1].text(.085, .75, r"$1-e^{-36\lambda}$", fontsize=16)
    fig.text(.075, .01, "The bar marks the split. Four beige cells produce food already supplied; the other 32 can enlarge the closure.", fontsize=10)
    fig.subplots_adjust(bottom=.17, wspace=.36)
    save(fig, "food_boundary")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("results"))
    parser.add_argument("--trials", type=int, default=400)
    parser.add_argument("--seed", type=int, default=20260908)
    args = parser.parse_args()
    if args.trials < 1 or args.seed < 0:
        parser.error("trials must be positive and seed must be nonnegative")
    args.output.mkdir(parents=True, exist_ok=True)
    intensities = [0, .01, .03, .06, .1, .12, .14, .16, .18, .2, .25, .35, .5, 1, 2, 3]
    model = Polymer(4)
    gateway = (model.length[model.u] <= 2) & (model.length[model.v] <= 2)
    assert int(gateway.sum()) == 36
    assert int((gateway & (model.length[model.w] > 2)).sum()) == 32
    assert (Polymer(5).U, Polymer(5).R) == (62, 196)
    assert not canonical_raf(model, np.full(model.R, model.U+1))[0]
    assert canonical_raf(model, np.ones(model.R))[0]
    history = exact_history_check()
    examples = worked_examples()
    positive = exact_positive_bound()
    finite = simulate_finite(intensities, [4, 6, 8, 10], args.trials, args.seed)
    escape = simulate_escape(intensities, [2, 3, 4, 5], args.trials, args.seed)
    comparison = reference_comparison(args.trials, args.seed)
    barrier = static_barrier_check(args.trials, args.seed)
    save_csv(args.output / "finite_raf.csv", finite)
    save_csv(args.output / "static_escape.csv", escape)
    save_csv(args.output / "literal_comparison.csv", comparison)
    (args.output / "worked_examples.json").write_text(json.dumps(examples, indent=2)+"\n")
    make_figures(args.output, finite, escape, examples)
    metadata = {"paper": "Hordijk_Steel_Critical_Window_arxiv.pdf, September 2026",
                "seed": args.seed, "trials_per_point": args.trials,
                "python": platform.python_version(), "numpy": np.__version__,
                "matplotlib": matplotlib.__version__,
                "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                "exact_history_test": history, "positive_bound": positive,
                "static_barrier_check": barrier,
                "intervals": "Pointwise 95% Wilson; neither simultaneous nor finite-size error bounds.",
                "randomness": "NumPy default_rng, SeedSequence with documented separate spawn keys.",
                "coupling": "Common marks across intensities and static escape boundaries.",
                "limits": "No infinite survival, continuity, convergence rate or Lean proof was numerically verified."}
    (args.output / "run_metadata.json").write_text(json.dumps(metadata, indent=2)+"\n")
    lines = ["CRITICAL WINDOW: ACTUAL SAVED OUTPUT", "",
             f"Seed {args.seed}; {args.trials} independent samples per displayed point.",
             "PASS: U_5=62, R_5=196; 36 food gateways, of which 32 enlarge closure.",
             "PASS: all 64 small-example matrices give equal history polynomials.",
             "PASS: exact literal witnesses, failed-support example and food-only RAF.",
             "PASS: append-then-split produces target 101 using four fresh reactions.",
             "PASS: common-mark RAF monotonicity and nested static escape events.",
             f"PASS: static barrier: {barrier['barrier_successes']} barrier successes, "
             f"{barrier['RAF_successes']} RAFs; zero implication violations in {args.trials} samples.",
             f"PASS: exact positive-bound condition at a=1/2: k={positive['k']}, L={positive['L']}.",
             f"Paper-backed lower bound: Theta(ln 2) >= {positive['positive_lower_bound']} > 0.",
             "Paper-backed upper bound: Theta(ln 2) <= 1 - 2^(-36) < 1.", "",
             "FINITE RAF FREQUENCIES (95% pointwise Wilson intervals)"]
    for r in finite:
        if r["intensity"] in (.1, .5, 1.0):
            lines.append(f"n={r['n']:2d}, lambda={r['intensity']:g}: {r['successes']}/{args.trials} "
                         f"= {r['estimate']:.4f} [{r['ci_low']:.4f}, {r['ci_high']:.4f}]")
    lines += ["", "INDEPENDENT FINITE IMPLEMENTATION COMPARISON (n=6, lambda=0.2)"]
    for r in comparison:
        lines.append(f"{r['method']}: {r['successes']}/{args.trials}, "
                     f"95% interval [{r['ci_low']:.4f}, {r['ci_high']:.4f}]")
    lines += ["", "Static escape frequencies estimate upper approximations to S(a), not S(a) itself.",
              "The paper proves the infinite-size statements. This run did not compile Lean."]
    transcript = "\n".join(lines)+"\n"
    (args.output / "diagnostics.txt").write_text(transcript)
    print(transcript)


if __name__ == "__main__":
    main()
