"""Build NTUH2022G4 labeled manifest: brain native T1c + tumor seg pairs from register_inv.

Rules:
- series name must contain T1 and +C
- exclude MRA/FLAIR/TOF/T2/SWI/DWI, _ROI1/_ROI re-exports, and spinal levels (e.g. T4-T8, T11-L3)
- dedupe per (subject, case, acquisition timestamp): prefer _MPR_Tra, else lowest series number
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import re
from collections import defaultdict
from src.common import ROOT, DATA, save_jsonl, path

T1C_RE = re.compile(r"T1.*\+C|\+C.*T1")
EXCL_RE = re.compile(r"MRA|FLAIR|TOF|T2|SWI|DWI|_ROI|ROI1", re.I)
SPINE_RE = re.compile(r"[TLC]\s?\d+\s?[-–]\s?[TLC]\s?\d+|[TLC]\d{2}", re.I)


def main(limit=None):
    out = os.path.join(DATA, "manifests", "ntuh.jsonl")
    reg_inv = path("ntuh_register_inv")
    cands = {}  # (subj, case, ts) -> list of (pref, ser, fname)
    n_subj = 0
    for s in sorted(os.listdir(reg_inv)):
        fp = os.path.join(reg_inv, s)
        if not os.path.isdir(fp):
            continue
        n_subj += 1
        for c in os.listdir(fp):
            cp = os.path.join(fp, c)
            if not os.path.isdir(cp):
                continue
            fl = set(os.listdir(cp))
            for f in sorted(fl):
                if not f.endswith(".nii.gz") or f.endswith((".seg.nii.gz", ".label.nii.gz")):
                    continue
                if not T1C_RE.search(f) or EXCL_RE.search(f) or SPINE_RE.search(f):
                    continue
                if f[: -len(".nii.gz")] + ".seg.nii.gz" not in fl:
                    continue
                mt = re.search(r"_(\d{14})_(\d+)\.nii\.gz$", f)
                ts = mt.group(1) if mt else f
                ser = int(mt.group(2)) if mt else 9999
                pref = 0 if "MPR_Tra" in f else 1
                cands.setdefault((s, c, ts), []).append((pref, ser, f))
                if limit and n_subj >= limit:
                    break
    rows = []
    for (s, c, ts), lst in sorted(cands.items()):
        lst.sort(key=lambda x: (x[0], x[1]))
        f = lst[0][2]
        rows.append({"key": f"ntuh_{s}_{c}_{ts}", "subject": s, "case": c,
                     "date": c2date(c), "img": os.path.join(reg_inv, s, c, f),
                     "label": os.path.join(reg_inv, s, c, f[: -len(".nii.gz")] + ".seg.nii.gz"),
                     "source": "ntuh", "dedup": len(lst)})
    save_jsonl(rows, out)
    print(f"subjects={n_subj} rows={len(rows)} (deduped {sum(r['dedup'] for r in rows) - len(rows)})")


def c2date(c):
    m = re.match(r"case?(\d{4})\.(\d{2})\.(\d{2})\.", c)
    return f"{m.group(1)}-{m.group(2)}-{m.group(3)}" if m else c


if __name__ == "__main__":
    main()