"""merge every paper figure into one browsable pdf.

section-title pages divide main / supplement / biology; each figure gets
a small header page with its filename so a reader can find the source.
"""
from __future__ import annotations
from pathlib import Path
from pypdf import PdfWriter
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
ROOT = Path(str(PANDA_ROOT))
FIG = ROOT / "figures"
OUT = FIG / "PANDA_all_figures.pdf"
TMP = FIG / "_tmp_titlepages.pdf"

# Labels are the numbers PAPER.tex actually resolves to (PAPER.aux), not the
# historical filenames: fig5_* is Figure 2 and fig6_* is Figure 5.
MAIN = [
    ("Figure 1", "fig1_perclass_f1.pdf",        "Per-class held-out 5-fold F1 across three systems"),
    ("Figure 2", "fig5_dingwall_umap.pdf",      "Dingwall UMAP (predicted class and En1 genotype)"),
    ("Figure 3", "fig3_dahlin_heatmap.pdf",     "Dahlin Kit-W41 vs WT within-class pathway module deltas"),
    ("Figure 4", "fig4_veres_stage_stack.pdf",  "Veres predicted class fraction per protocol stage"),
    ("Figure 5", "fig6_multi_umap.pdf",         "Discovery-target UMAPs across skin / HSC / pancreas"),
]

SUPP = [
    ("S1",  "supplement/01_cv_summary.pdf",                                "CV summary"),
    ("S2",  "supplement/02_per_class_f1.pdf",                              "Per-class F1"),
    ("S3",  "supplement/03_prototype_cosine.pdf",                          "Prototype intra-cosine"),
    # S4 removed -- contained stale class vocab
    ("S5",  "supplement/05_adversary_purification.pdf",                    "Adversary purification"),
    ("S6",  "supplement/06_cross_system_prototypes.pdf",                   "Cross-system prototype map"),
    # S7 removed -- contained stale class vocab
    # S8 removed -- contained stale class vocab
    # S9 removed -- contained stale class vocab
    # S10 removed -- contained stale class vocab
    ("S11", "supplement/11_novel_populations.pdf",                         "Novel populations (legacy)"),
    # S12 removed -- contained stale class vocab
    ("S13", "supplement/13_dingwall_umap.pdf",                             "Dingwall UMAP (legacy)"),
    ("S14", "supplement/14_dahlin_umap.pdf",                               "Dahlin UMAP (legacy)"),
    ("S15", "supplement/15_veres_umap.pdf",                                "Veres UMAP (legacy)"),
    ("S16", "supplement/16_dingwall_discovery.pdf",                        "Dingwall discovery (legacy)"),
    ("S17", "supplement/17_dahlin_discovery.pdf",                          "Dahlin discovery (legacy)"),
    ("S18", "supplement/18_veres_discovery.pdf",                           "Veres discovery (legacy)"),
    ("S19", "supplement/19_myeloid_network.pdf",                           "Myeloid gene-gene network"),
    ("S20", "supplement/20_placode_wnt_module.pdf",                        "Placode WNT module"),
    ("S23", "supplement/23_anchor_delta_recall.pdf",                       "Anchor delta recall"),
    ("S24", "supplement/24_pca_vs_marker_umaps_dingwall_by_genotype.pdf",  "PCA vs Marker: Dingwall by genotype"),
    ("S24b","supplement/24b_pca_vs_marker_umaps_dingwall_by_class.pdf",    "PCA vs Marker: Dingwall by class"),
    ("S25", "supplement/25_pca_vs_marker_umaps_dahlin_by_genotype.pdf",    "PCA vs Marker: Dahlin by genotype"),
    ("S25b","supplement/25b_pca_vs_marker_umaps_dahlin_by_class.pdf",      "PCA vs Marker: Dahlin by class"),
    ("S26", "supplement/26_pca_vs_marker_umaps_veres_by_stage.pdf",        "PCA vs Marker: Veres by stage"),
    ("S26b","supplement/26b_pca_vs_marker_umaps_veres_by_class.pdf",       "PCA vs Marker: Veres by class"),
    ("S27", "supplement/27_dingwall_en1_enrichment.pdf",                   "Dingwall En1 class enrichment"),
    ("S28", "supplement/28_melanocyte_pathway_modules.pdf",                "Melanocyte pathway modules"),
]

BIO = [
    ("B1", "biology/biology_01_dingwall_umap.pdf",        "Dingwall UMAP (biology)"),
    ("B2", "biology/biology_02_primary_eden.pdf",         "Primary EDEN candidate: Derm2"),
    ("B3", "biology/biology_03_melanoblast_mitf.pdf",     "Melanoblast MITF-axis quadrant"),
    ("B4", "biology/biology_04_dahlin_metabolism.pdf",    "Dahlin per-lineage metabolism"),
    ("B5", "biology/biology_05_dahlin_composition.pdf",   "Dahlin composition shift"),
    ("B6", "biology/biology_06_veres_beta_quadrant.pdf",  "Veres beta-lineage quadrant"),
    ("B7", "biology/biology_07_veres_polyhormonal.pdf",   "Veres polyhormonal SC-alpha"),
    ("B8", "biology/biology_08_prototype_geometry.pdf",   "Prototype geometry"),
]


def make_title_page(text_top, text_body):
    fig = plt.figure(figsize=(8.5, 11))
    fig.text(0.5, 0.55, text_top,  ha="center", va="center", fontsize=17, weight="bold")
    fig.text(0.5, 0.45, text_body, ha="center", va="center", fontsize=12)
    # NO bbox_inches="tight" here: it crops the page down to the text bbox,
    # so the heading fills the sheet no matter what fontsize is set.
    fig.savefig(TMP)
    plt.close(fig)


def append_group(w, title, items):
    make_title_page(title, "")
    w.append(str(TMP))
    for tag, path, desc in items:
        p = FIG / path
        if not p.exists():
            print(f"  [skip] {path} missing")
            continue
        make_title_page(tag, f"{desc}\n\n{path}")
        w.append(str(TMP))
        w.append(str(p))
        print(f"  + {tag} {p.name}")


def main():
    w = PdfWriter()
    make_title_page("PANDA — all figures", "main text · supplement · biology")
    w.append(str(TMP))
    append_group(w, "Main text", MAIN)
    append_group(w, "Supplement", SUPP)
    append_group(w, "Biology", BIO)
    with open(OUT, "wb") as f:
        w.write(f)
    TMP.unlink(missing_ok=True)
    print(f"\n[all] wrote {OUT} ({OUT.stat().st_size / 1024:.0f} KB, "
          f"{len(w.pages)} pages)")


if __name__ == "__main__":
    main()