r"""Build the PNAS-draft PDF from main.md (Markdown stays the source of truth).

Adapted from paper/arxiv/md2tex.py (same Markdown subset + pipe tables), with one addition: standalone
`*(FIG:name)*` markers compose multi-panel figures by stacking existing per-experiment vector PDFs
(LaTeX-level consolidation; bespoke unified figures are a submission-time polish, tracked in the work
order). Captions define the panel letters positionally (A = top, ...) because the sub-figures carry
their own internal panel labels.

Usage:  python paper/pnas/build.py && (cd paper/pnas && tectonic main.tex)
"""

from __future__ import annotations

import re
import shutil
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
HERE = Path(__file__).resolve().parent
SRC = HERE / "main.md"
OUT = HERE / "body.tex"

# figure name -> (list of source PDFs (stacked top->bottom), caption)
FIGURES: dict[str, tuple[list[str], str]] = {
    "fig1": (["results/E2/E2.pdf", "results/mnist_collapse/mnist_montage.pdf"],
        "Collapse is drift; grounding is immigration. (A, top) The grounding phase response in the "
        "minimal model: a critical real-data fraction $g^*\\!\\approx\\!0.05$ retains most diversity "
        "indefinitely, while tail survival obeys the per-item floor $m\\,p \\gtrsim 1$. (B, bottom) "
        "The same signs on real images: a convolutional VAE retrained each generation on its own "
        "output collapses to a single blurred mode (rows: generations), while $\\sim$10\\% grounding "
        "holds all thirty class$\\times$style modes."),
    "fig2": (["results/E4/E4.pdf", "results/E8/E8.pdf"],
        "Recombination: the conservation law and the Fisher--Muller effect. (A, top) Refitting a child "
        "to the mean of its parents' output distributions conserves rare-item mass at single-parent "
        "level regardless of parent count (blending inheritance); a strongest-source (union) operator "
        "realises the multi-parent gain. (B, bottom) Multi-locus recombination of decorrelated "
        "specialists assembles a genotype fitter than any parent, climbing to the optimum as parents "
        "are added, while the best single parent and the blended average plateau below."),
    "fig3": (["results/E9/E9.pdf", "results/E10/E10.pdf", "results/E14/E14.pdf"],
        "Rugged (epistatic) landscapes: risk, remedy, and structure. (A, top) Outbreeding depression: "
        "blind recombination of specialists drops offspring below their parents, worsening with "
        "ruggedness; the optimal recombination rate shrinks as skills entangle. (B, middle) Directed "
        "sex --- unbounded parents, chosen mates, verifier-screened offspring --- converts the "
        "catastrophe into a reliable gain at every ruggedness. (C, bottom) Mating structure: wide "
        "(promiscuous) mixing maximises the population mean but monotonically destroys diversity; the "
        "champion-optimal mate-pool breadth narrows as the landscape roughens."),
    "fig4": (["results/E11/E11.pdf"],
        "The society: grounding, sex, and diversity are jointly necessary. A finite agent population "
        "on a rugged NK landscape under a grounded selection score. Four-arm ablation: the full system "
        "climbs to near the global optimum; removing grounding collapses the population onto a "
        "confident, unfit consensus (self-consumption); removing recombination strands it on local "
        "optima; removing diversity converges it prematurely. Each ablation fails differently."),
    "fig5": (["results/E12/E12.pdf", "results/speciation_real/speciation_real.pdf",
              "results/llm_speciation/llm_speciation.pdf"],
        "Model speciation across three tiers. (A, top) Analytic: hybrid fitness traces compatible "
        "$\\rightarrow$ outbreeding depression $\\rightarrow$ inviability; the cliff arrives earlier "
        "the denser the incompatibilities; incompatibility count snowballs with divergence. (B, "
        "middle) Trained MLPs: the merge barrier decomposed under the complete unit symmetry group --- "
        "same-task/different-init barriers are coordinate artefacts (removed by alignment); "
        "conflicting-task barriers survive in full, with hybrid fitness falling 0.97 $\\rightarrow$ "
        "0.03; divergence without conflict produced no isolation, the merge instead rescuing the "
        "forgetting specialists. (C, bottom) Language models: conflicting conventions produce "
        "function-specific hybrid breakdown; over-training disjoint specialists produces none --- at "
        "every tier tested, isolation had to be provoked by functional conflict."),
    "fig6": (["results/llm_merge_seeds/llm_seeds.pdf", "results/llm_moe_hard_hpc/llm_moe.pdf",
              "results/llm_epistasis/llm_epistasis.pdf"],
        "The language-model tier. (A, top) Seed-replicated recombination claims (fixed test sets, "
        "training seed varied, 95\\% CI): merges beat every specialist; union-preserving routing and "
        "directed offspring selection beat the blend in every seed on headroom tasks, including one "
        "catastrophic blend failure they avoided. (B, middle) The headroom rule at 7B on hard "
        "(unsaturated) tasks: the weight-average dilutes a fragile specialist below the best single "
        "parent; routing preserves it. (C, bottom) The controlled predictive test: across a task grid "
        "with conflict, compatible-overlap, and duration axes decorrelated by construction, pre-merge "
        "functional disagreement predicts merge penalty (held-out $\\rho \\approx 0.4$) while "
        "weight-geometry baselines show no detectable association; paired predictor differences are "
        "not individually significant."),
}

UNICODE = {"—": "---", "–": "--", "→": r"\(\rightarrow\)", "≈": r"\(\approx\)", "≥": r"\(\geq\)",
           "≳": r"\(\gtrsim\)", "×": r"\(\times\)", "·": r"\(\cdot\)", "μ": r"\(\mu\)",
           "ρ": r"\(\rho\)", "≤": r"\(\leq\)", "≪": r"\(\ll\)", "∝": r"\(\propto\)"}
SPECIALS = {"&": r"\&", "%": r"\%", "#": r"\#", "_": r"\_", "$": r"\$",
            "~": r"\textasciitilde{}", "^": r"\textasciicircum{}"}


def esc(s: str) -> str:
    s = s.replace("\\", r"\textbackslash{}")
    for k, v in SPECIALS.items():
        s = s.replace(k, v)
    for k, v in UNICODE.items():
        s = s.replace(k, v)
    return s


def inline(s: str) -> str:
    parts = re.split(r"(`[^`]*`)", s)
    out = []
    for p in parts:
        if p.startswith("`") and p.endswith("`") and len(p) >= 2:
            out.append(r"\texttt{" + esc(p[1:-1]) + "}")
        else:
            p = esc(p)
            p = re.sub(r"\[([^\]]+)\]\((https?://[^)]+)\)", r"\\href{\2}{\1}", p)
            p = re.sub(r"\*\*([^*]+)\*\*", r"\\textbf{\1}", p)
            p = re.sub(r"\*([^*]+)\*", r"\\emph{\1}", p)
            p = re.sub(r'"([^"]+)"', r"``\1''", p)
            out.append(p)
    return "".join(out)


def figure_env(name: str) -> str:
    pdfs, caption = FIGURES[name]
    (HERE / "figs").mkdir(exist_ok=True)
    lines = [f"\\begin{{figure*}}[p]\\centering  % {name}"]
    for src in pdfs:
        dst = HERE / "figs" / (name + "_" + Path(src).name)
        shutil.copyfile(ROOT / src, dst)
        frac = min(0.98, 3.0 / len(pdfs) * 0.42)
        lines.append(f"\\includegraphics[width=\\textwidth,height={frac:.2f}\\textheight,"
                     f"keepaspectratio]{{figs/{dst.name}}}\\par\\smallskip")
    lines.append(f"\\caption{{{caption}}}\\label{{{name}}}")
    lines.append("\\end{figure*}")
    return "\n".join(lines)


def convert(text: str) -> str:
    lines = text.split("\n")
    i = 0
    while i < len(lines) and lines[i].strip() != "---":
        i += 1
    i += 1

    blocks: list[list[str]] = []
    cur: list[str] = []
    for line in lines[i:]:
        if line.strip() == "":
            if cur:
                blocks.append(cur); cur = []
        else:
            cur.append(line)
    if cur:
        blocks.append(cur)

    def emit_table(block, out):
        rows = [[c.strip() for c in line.strip().strip("|").split("|")] for line in block]
        header, body = rows[0], rows[2:]
        n = len(header)
        widths = " ".join([f"p{{{0.92 / n:.3f}\\textwidth}}"] * n)
        out += ["\\medskip\\noindent\\begin{center}\\footnotesize",
                f"\\begin{{tabular}}{{{widths}}}", "\\hline",
                " & ".join(inline(c) for c in header) + " \\\\ \\hline"]
        for r in body:
            r = (r + [""] * n)[:n]
            out.append(" & ".join(inline(c) for c in r) + " \\\\[3pt]")
        out += ["\\hline\\end{tabular}\\end{center}\\medskip", ""]

    out: list[str] = []
    for block in blocks:
        first = block[0].strip()
        m = re.match(r"^\*?\(FIG:(\w+)\)\*?$", first)
        if m:
            out.append(figure_env(m.group(1))); out.append("")
        elif first.startswith("|") and len(block) >= 2 and set(block[1].strip()) <= set("|-: "):
            emit_table(block, out)
        elif first == "---" and len(block) == 1:
            out.append("\\medskip\\hrule\\medskip"); out.append("")
        elif first.startswith("## "):
            out.append(f"\\section*{{{inline(first[3:])}}}"); out.append("")
        elif first.startswith("### "):
            out.append(f"\\subsection*{{{inline(first[4:])}}}"); out.append("")
        elif re.match(r"^(- |\d+\. )", first):
            env = "itemize" if first.startswith("- ") else "enumerate"
            out.append(f"\\begin{{{env}}}")
            items: list[str] = []
            for l in block:
                s = l.strip()
                if re.match(r"^(- |\d+\. )", s):
                    items.append(re.sub(r"^(- |\d+\. )", "", s))
                else:
                    items[-1] += " " + s
            for it in items:
                out.append("\\item " + inline(it.strip()))
            out.append(f"\\end{{{env}}}"); out.append("")
        else:
            joined = re.sub(r"\s{2,}", " ", " ".join(l.strip() for l in block)).strip()
            out.append(inline(joined)); out.append("")
    return "\n".join(out) + "\n"


if __name__ == "__main__":
    OUT.write_text(convert(SRC.read_text()))
    print(f"wrote {OUT}")