"""First-stage curation + TV-slice QA figures for labeled rows.

Runs on the labeled source rows (data/manifests/{ntuh,m6_labeled}.jsonl, 1mm
processed) — the first place rows exist after fetching from the sources,
before any split (05). For each labeled volume:

  1. classify with the exclusion rules in src/exclusions.py (total TV > 1e6
     mm3, single target < 100 mm3, curated non-brain list),
  2. render one PNG with the three orthogonal slices (axial | coronal |
     sagittal) that carry the largest TV (target-volume) area, with the TV
     contour overlaid; when the label has several connected targets
     (26-connectivity), only the largest one is shown,
  3. write the figure to data/tv_slices/inclusion/ (kept) or
     data/tv_slices/exclusion/ (excluded, reason in the title).

Slices are picked per body axis (S/P/L) from the procmeta direction, so the
panel order follows radiological conventions (same mapping as
scripts/make_screenshots.py); volumes that are not body-aligned fall back to
the image-axis slices with a note in the title. Filenames are
<total_vol_mm3:07d>_<key>.png, so each folder lists small TV -> large TV.

The split files (05) are the dataset of record; this stage only curates and
documents. 05 applies the same src/exclusions rules after split assignment,
so the two always agree on kept/excluded.

Usage:
  python scripts/make_tv_slices.py                  # all source rows, 32 workers
  python scripts/make_tv_slices.py --from-splits    # rows/tag from split files
  python scripts/make_tv_slices.py --keys k1 k2     # specific rows
  python scripts/make_tv_slices.py --limit 50       # first N rows
  python scripts/make_tv_slices.py --force          # overwrite existing PNGs
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import glob
import json
import argparse
from collections import Counter

import numpy as np
import SimpleITK as sitk
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import ndimage

from src.common import DATA, d, load_jsonl, proc_row, labeled_source_rows
from src.exclusions import label_stats, exclusion_reason
from scripts.make_screenshots import vol_to_body


def out_name(vol_total, key):
    """Flat-folder filename: zero-padded TOTAL TV volume (mm3, all targets)
    first so a plain listing sorts small -> large, then the row key."""
    return f"{int(vol_total):07d}_{key}.png"


def largest_cc(mask, structure=None):
    """Largest connected component of a boolean mask (26-conn by default)."""
    if structure is None:
        structure = np.ones((3, 3, 3), dtype=int)
    lab, n = ndimage.label(mask, structure=structure)
    if n == 0:
        return mask, 0
    sizes = ndimage.sum(mask, lab, range(1, n + 1))
    return lab == (int(np.argmax(sizes)) + 1), int(n)


def max_area_slices(cc):
    """Per-axis index of the slice containing the largest CC area."""
    return tuple(int(np.argmax(cc.sum(axis=tuple(j for j in range(cc.ndim) if j != i))))
                 for i in range(cc.ndim))


def render(key, tag, mr, cc, n_cc, vol_total, direction, reason, out, dpi=110):
    """1x3 figure: axial | coronal | sagittal, max-CC-area slice each, with the
    largest-CC contour in red. Writes to `out`. Returns the path."""
    body, ok, _ = vol_to_body(mr, direction)
    if ok:
        c, _, _ = vol_to_body(cc.astype(np.float32), direction)
        si, pi, li = max_area_slices(c)
        panels = [
            (body[si], c[si], "axial", "S", si),
            (body[:, pi, :][::-1], c[:, pi, :][::-1], "coronal", "P", pi),
            (body[:, :, li][::-1], c[:, :, li][::-1], "sagittal", "L", li),
        ]
    else:  # not body-aligned: image-axis fallback
        zi, yi, xi = max_area_slices(cc)
        panels = [
            (mr[zi], cc[zi], "axial*", "z", zi),
            (mr[:, yi, :][::-1], cc[:, yi, :][::-1], "coronal*", "y", yi),
            (mr[:, :, xi][::-1], cc[:, :, xi][::-1], "sagittal*", "x", xi),
        ]

    lo, hi = np.percentile(mr[mr > 0], [1, 99.5])
    panel_specs = []
    for im, c2, title, axn, idx in panels:
        nr, nc = im.shape
        a = 1.0  # processed volumes are 1mm isotropic
        panel_specs.append((im, c2, f"{title}  {axn}={idx}", a, nr, nc))

    max_in, gap_in, m_in, top_in, bot_in = 4.4, 0.25, 0.3, 1.7, 0.2
    w_in, h_in = [], []
    for _, _, _, a, nrow, ncol in panel_specs:
        pw, ph = ncol, nrow * a
        m = max(pw, ph)
        w_in.append(max_in * pw / m)
        h_in.append(max_in * ph / m)
    plot_h = max_in
    fig_w = sum(w_in) + 2 * gap_in + 2 * m_in
    fig_h = plot_h + top_in + bot_in
    fig = plt.figure(figsize=(fig_w, fig_h))
    left = m_in / fig_w
    for (im, c2, title, a, _, _), w, h in zip(panel_specs, w_in, h_in):
        axi = fig.add_axes([left, (bot_in + (plot_h - h) / 2) / fig_h, w / fig_w, h / fig_h])
        axi.imshow(im, cmap="gray", vmin=lo, vmax=hi, origin="upper", aspect=a)
        if c2.any():
            axi.contour(c2, levels=[0.5], colors="r", linewidths=1.2)
        axi.set_title(title, color="w", fontsize=11)
        axi.set_xticks([])
        axi.set_yticks([])
        for spine in axi.spines.values():
            spine.set_edgecolor("0.35")
        left += (w + gap_in) / fig_w
    tkey = key.split("_", 1)[1] if "_" in key else key
    align = "body-aligned" if ok else "NOT body-aligned (image axes)"
    status = f"EXCLUDED: {reason}" if reason else "kept"
    fig.suptitle(
        f"{tkey}   [{tag}]   TV total={int(vol_total)} mm3, "
        f"largest={int(cc.sum())} mm3 ({n_cc} targets)   {status}   {align}",
        color="w", fontsize=12)
    fig.patch.set_facecolor("k")
    fig.savefig(out, facecolor="k", dpi=dpi)
    plt.close(fig)
    return out


def worker(args):
    key, tag, pimg, plabel, out_base, force, dpi = args
    try:
        meta = json.load(open(os.path.join(d("data/procmeta"), key + ".json")))
        stats = label_stats(plabel)
    except Exception:
        return key, "error: unreadable label/procmeta"
    reason = exclusion_reason(key, stats)
    cat = reason.split(" (", 1)[0] if reason else None
    # reconcile a stale figure for this key left in the other folder by an
    # earlier run under a different exclusion status (e.g. a key newly added to
    # NON_BRAIN_KEYS would otherwise keep its old inclusion/ figure)
    folder = "exclusion" if reason else "inclusion"
    other = "inclusion" if reason else "exclusion"
    for stale in glob.glob(os.path.join(out_base, other, f"*_{key}.png")):
        os.remove(stale)
    if stats["total"] == 0:
        return key, "empty-label", cat
    status = "excluded" if reason else "included"
    out = os.path.join(out_base, folder, out_name(stats["total"], key))
    if os.path.exists(out) and not force:
        return key, f"skip-{status}", cat
    try:
        mr = np.asarray(sitk.GetArrayFromImage(sitk.ReadImage(pimg)), dtype=np.float32)
        lab = sitk.GetArrayFromImage(sitk.ReadImage(plabel))
        cc, n_cc = largest_cc(lab > 0.5)
        try:
            render(key, tag, mr, cc, n_cc, int(stats["total"]),
                   np.array(meta["direction"]), reason, out, dpi)
        finally:
            plt.close("all")
        return key, status, cat
    except Exception as e:  # noqa
        return key, f"error: {e!r}", cat


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--from-splits", action="store_true",
                    help="rows/tags from split_{train,val,test}.jsonl instead of "
                         "the source manifests (ntuh + m6_labeled)")
    ap.add_argument("--keys", nargs="*", default=[])
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--out", default=d("data/tv_slices"))
    ap.add_argument("--dpi", type=int, default=110)
    args = ap.parse_args()
    os.makedirs(os.path.join(args.out, "inclusion"), exist_ok=True)
    os.makedirs(os.path.join(args.out, "exclusion"), exist_ok=True)

    if args.from_splits:
        rows = []
        for s in ("train", "val", "test"):
            for r in load_jsonl(os.path.join(DATA, "manifests", f"split_{s}.jsonl")):
                rows.append((r["key"], s, r["pimg"], r.get("plabel")))
    else:
        rows = [(r["key"],
                 "ntuh" if r["key"].startswith("ntuh_") else "m6",
                 r["pimg"], r.get("plabel")) for r in labeled_source_rows()]
    if args.keys:
        keep = set(args.keys)
        rows = [r for r in rows if r[0] in keep]
    if args.limit:
        rows = rows[: args.limit]
    print(f"rows: {len(rows)}")

    from concurrent.futures import ProcessPoolExecutor, as_completed
    tasks = [(k, tag, pimg, plabel, args.out, args.force, args.dpi)
             for k, tag, pimg, plabel in rows]
    stats = {}
    reasons = Counter()
    done = 0
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(worker, t): t for t in tasks}
        for fu in as_completed(futs):
            k, status, cat = fu.result()
            stats[status] = stats.get(status, 0) + 1
            if status == "excluded" and cat:
                reasons[cat] += 1
            done += 1
            if done % 100 == 0 or done == len(tasks):
                print(f"  {done}/{len(tasks)}  {stats}", flush=True)
    print(f"done: {stats}")
    if reasons:
        print("excluded by rule:")
        for r, n in sorted(reasons.items()):
            print(f"  {n:4d}  {r}")


if __name__ == "__main__":
    main()