#!/usr/bin/env python3
"""Generate publication figures only from audited manuscript tables."""

from __future__ import annotations

import csv
from pathlib import Path

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


ROOT = Path(__file__).resolve().parent.parent
TABLES = ROOT / "tables"
FIGURES = ROOT / "figures"

COLORS = {
    "COCO-full": "#0072B2",
    "ImageNet+COCO": "#E69F00",
    "SimCLR+COCO": "#009E73",
    "SupCon+COCO": "#CC79A7",
    "COCO": "#0072B2",
    "ImageNet": "#E69F00",
    "SimCLR": "#009E73",
    "SupCon": "#CC79A7",
}
MARKERS = {"COCO-full": "o", "ImageNet+COCO": "s", "SimCLR+COCO": "^", "SupCon+COCO": "D"}


def rows(name: str) -> list[dict[str, str]]:
    with (TABLES / name).open(encoding="utf-8", newline="") as handle:
        return list(csv.DictReader(handle))


def style() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8,
            "axes.labelsize": 8,
            "axes.titlesize": 9,
            "legend.fontsize": 7,
            "xtick.labelsize": 7,
            "ytick.labelsize": 7,
            "axes.linewidth": 0.7,
            "lines.linewidth": 1.3,
            "lines.markersize": 4.5,
            "savefig.bbox": "tight",
            "savefig.pad_inches": 0.03,
        }
    )


def finish(fig: plt.Figure, stem: str) -> None:
    FIGURES.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGURES / f"{stem}.pdf")
    fig.savefig(FIGURES / f"{stem}.eps")
    fig.savefig(FIGURES / f"{stem}.png", dpi=600)
    plt.close(fig)


def fig1_design() -> None:
    fig, ax = plt.subplots(figsize=(7.0, 3.0))
    ax.set_xlim(0, 12)
    ax.set_ylim(0, 4.5)
    ax.axis("off")
    boxes = [
        (0.15, 2.75, 2.70, 1.25, "Pretraining inputs", "COCO-full | ImageNet\nSimCLR | SupCon"),
        (3.10, 2.75, 2.70, 1.25, "Module-controlled\ntransfer", "Fixed Mask R-CNN\nbackbone-only replacement"),
        (6.05, 2.75, 2.70, 1.25, "Development controls", "Matched random components\nbackbone LR \u00d7 budget"),
        (9.00, 2.75, 2.70, 1.25, "Frozen primary test", "Official LIVECell test\n3 seeds | AP@3000"),
        (3.10, 0.45, 2.70, 1.25, "Dense protocol", "maxDets response\ncell-density strata"),
        (6.05, 0.45, 2.70, 1.25, "Conditional\nuncertainty", "Paired image bootstrap\nfixed models only"),
        (9.00, 0.45, 2.70, 1.25, "Supplementary diagnostic", "Capped Cellpose\nconstruction audit"),
    ]
    for i, (x, y, w, h, title, body) in enumerate(boxes):
        color = "#DCEEF8" if i < 4 else "#E8F3E8"
        ax.add_patch(plt.Rectangle((x, y), w, h, facecolor=color, edgecolor="#333333", linewidth=0.8))
        ax.text(x + 0.13, y + h - 0.20, title, weight="bold", va="top", fontsize=6.8, linespacing=1.0)
        ax.text(x + 0.13, y + h - 0.72, body, va="top", fontsize=5.9, linespacing=1.14)
    arrow = dict(arrowstyle="->", color="#444444", linewidth=1.0)
    for x1, x2 in ((2.85, 3.10), (5.80, 6.05), (8.75, 9.00)):
        ax.annotate("", xy=(x2, 3.38), xytext=(x1, 3.38), arrowprops=arrow)
    finish(fig, "Fig1_study_design")


def fig2_primary_density() -> None:
    primary = rows("table_primary_livecell_test_n3.csv")
    density = rows("table_density_strata_n3.csv")
    fig, axes = plt.subplots(1, 2, figsize=(7.0, 2.65), gridspec_kw={"wspace": 0.32})

    ax = axes[0]
    for x, row in enumerate(primary):
        method = row["method"]
        vals = [float(row[f"seed{s}"]) for s in (42, 43, 44)] if "seed42" in row else []
        # The primary summary stores only aggregate statistics, so retrieve seeds from density/all.
        if not vals:
            drow = next(r for r in density if r["method"] == method and r["density"] == "all")
            vals = [float(drow[f"seed{s}"]) for s in (42, 43, 44)]
        jitter = [-0.08, 0.0, 0.08]
        ax.scatter(np.array(jitter) + x, vals, color=COLORS[method], marker=MARKERS[method], zorder=3)
        ax.plot([x - 0.15, x + 0.15], [float(row["segm_AP_mean"])] * 2, color="black", lw=1.2)
    ax.set_xticks(range(4), ["COCO", "ImageNet", "SimCLR", "SupCon"], rotation=20, ha="right")
    ax.set_ylabel("Mask AP@3000")
    ax.set_title("(a) Official LIVECell test (three seeds)", loc="left")
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    ax.set_ylim(0.185, 0.255)

    ax = axes[1]
    strata = ["low", "medium", "high", "extreme"]
    labels = ["<=100", "101-300", "301-1000", ">1000"]
    for method in COLORS:
        if method not in MARKERS:
            continue
        vals = [float(next(r for r in density if r["method"] == method and r["density"] == d)["segm_AP_mean"]) for d in strata]
        ax.plot(labels, vals, label=method.replace("+COCO", ""), color=COLORS[method], marker=MARKERS[method])
    ax.set_ylabel("Mask AP@3000")
    ax.set_xlabel("Ground-truth instances per image")
    ax.set_title("(b) Density-stratified performance", loc="left")
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    ax.legend(frameon=False, ncol=2, loc="upper right")
    finish(fig, "Fig2_primary_and_density")


def fig3_controls() -> None:
    random_rows = rows("matched_random_raw.csv")
    factorial = rows("factorial_100epoch_raw.csv")
    effects = rows("factorial_100epoch_effects_by_seed.csv")
    fig, axes = plt.subplots(1, 3, figsize=(7.0, 2.55), gridspec_kw={"wspace": 0.58})

    ax = axes[0]
    backbones = ["COCO", "ImageNet", "SimCLR", "SupCon"]
    for x, name in enumerate(backbones):
        vals = [float(r["AP"]) for r in random_rows if r["backbone"] == name]
        ax.scatter([x - 0.04, x + 0.04], vals, color=COLORS[name], marker=["o", "s", "^", "D"][x])
        ax.plot([x - 0.14, x + 0.14], [np.mean(vals)] * 2, color="black", lw=1.1)
    ax.set_xticks(range(4), backbones, rotation=28, ha="right")
    ax.set_ylabel("Validation mask AP@3000")
    ax.set_title("(a) Random components", loc="left", fontsize=8)
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)

    ax = axes[1]
    factorial_seeds = sorted({int(row["seed"]) for row in factorial})
    line_styles = ("-", "--", ":")
    for backbone, color, marker in (("COCO-full", COLORS["COCO-full"], "o"), ("SimCLR+COCO", COLORS["SimCLR+COCO"], "^")):
        for seed, ls in zip(factorial_seeds, line_styles):
            vals = [float(next(r for r in factorial if r["backbone"] == backbone and int(r["seed"]) == seed and int(r["backbone_lr_multiplier"]) == lr)["AP"]) for lr in (1, 5)]
            ax.plot([1, 5], vals, color=color, marker=marker, linestyle=ls, label=f"{backbone.split('+')[0]}, s{seed}")
    ax.set_xticks([1, 5], ["1\u00d7", "5\u00d7"])
    ax.set_xlabel("Backbone LR multiplier")
    ax.set_ylabel("Validation mask AP@3000")
    ax.set_title("(b) COCO \u00d7 SimCLR, 100 epochs", loc="left", fontsize=8)
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    ax.legend(frameon=False, fontsize=6, loc="best")

    ax = axes[2]
    did = [float(r["difference_in_differences"]) for r in effects]
    bar_x = list(range(len(did)))
    ax.bar(bar_x, did, color=["#56B4E9", "#999999", "#D55E00"][:len(did)], width=0.6)
    ax.axhline(0, color="black", lw=0.7)
    ax.set_xticks(bar_x, [f"seed {row['seed']}" for row in effects])
    ax.set_ylabel("Descriptive difference-in-differences\nin validation AP", fontsize=5.7, labelpad=1)
    ax.set_title("(c) Descriptive interaction", loc="left", fontsize=8)
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    finish(fig, "Fig3_component_and_optimization_controls")


def fig4_dense_evaluation_caps() -> None:
    cap_file = "table_maxdets_ablation_n3.csv" if (TABLES / "table_maxdets_ablation_n3.csv").exists() else "table_maxdets_ablation_n2.csv"
    caps = rows(cap_file)
    fig, axes = plt.subplots(1, 2, figsize=(7.0, 2.65), gridspec_kw={"wspace": 0.32})

    ax = axes[0]
    for method in MARKERS:
        sub = sorted((r for r in caps if r["method"] == method), key=lambda r: int(r["max_dets"]))
        ax.plot([int(r["max_dets"]) for r in sub], [float(r["segm_AP_mean"]) for r in sub], color=COLORS[method], marker=MARKERS[method], label=method.replace("+COCO", ""))
    ax.set_xscale("log")
    ax.set_xticks([100, 300, 1000, 3000], ["100", "300", "1000", "3000"])
    ax.set_xlabel("maxDets per image")
    ax.set_ylabel("LIVECell mask AP")
    ax.set_title(f"(a) AP vs cap ({caps[0]['n_seeds']} seeds)", loc="left", fontsize=8)
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    ax.legend(frameon=False, ncol=2, fontsize=5.8)

    ax = axes[1]
    for method in MARKERS:
        sub = sorted((r for r in caps if r["method"] == method), key=lambda r: int(r["max_dets"]))
        ax.plot(
            [int(r["max_dets"]) for r in sub],
            [float(r["segm_AR_mean"]) for r in sub],
            color=COLORS[method], marker=MARKERS[method], label=method.replace("+COCO", ""),
        )
    ax.set_xscale("log")
    ax.set_xticks([100, 300, 1000, 3000], ["100", "300", "1000", "3000"])
    ax.set_xlabel("maxDets per image")
    ax.set_ylabel("LIVECell mask AR")
    ax.set_title(f"(b) AR vs cap ({caps[0]['n_seeds']} seeds)", loc="left", fontsize=8)
    ax.grid(axis="y", color="#D0D0D0", lw=0.5)
    finish(fig, "Fig4_dense_evaluation_caps")


def main() -> None:
    style()
    fig1_design()
    fig2_primary_density()
    fig3_controls()
    fig4_dense_evaluation_caps()
    print(f"MVA_FIGURES_OK output={FIGURES}")


if __name__ == "__main__":
    main()
