"""Scan lee longitudinal dataset for brain T1c (post-contrast T1) series.

Each MR timepoint dir: <sid>/<YYYYMMDD>_MR_<n>/ contains top-level per-slice
DICOM-dump txt files '<sid>_<date>_MR_<ser>_<slice>.txt' and one jpg folder
named like '<description>_<sid>_<date>' holding '<sid>_<date>_MR_<ser>_<slice>_<inst>.jpg'.

T1c series = contrasted (bolus tags) + T1-weighted 3D sequence.
Writes: data/manifests/lee_t1c_raw.jsonl and (after subset selection)
        data/manifests/lee_t1c_selected.jsonl
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import re
import json
import argparse
import random
from src.common import ROOT, save_jsonl, load_jsonl

BASE = "/mnt/t24/Public/lee"
T1_NAME_RE = re.compile(r"t1|tfl|spgr|mp2rage|tse3d|vfl|mpage", re.I)
EXCL_RE = re.compile(r"flair|\bt2\b|dwi|dti|mra|angi|swi|bold|\bpp2d|\bpp3d|perf|t2\*|t2star", re.I)


def parse_txt(p):
    """Return dict tagtuple -> value (value column)."""
    t = {}
    try:
        with open(p, errors="ignore") as f:
            for line in f:
                m = re.match(r'"\((\d{4}),(\d{4})\)","([A-Z]+)","(\d+)",(.*),"(.*)"\s*$', line.strip())
                if m:
                    t[(int(m.group(1), 16), int(m.group(2), 16))] = m.group(5).strip().strip('"')
    except OSError:
        pass
    return t


def is_t1c(t):
    bolus = any(k in t for k in ((24, 4161), (24, 4168), (8, 307)))
    if not bolus:
        return False, "no bolus"
    name = (t.get((24, 36), "") + " " + t.get((24, 33), "")).lower()
    if EXCL_RE.search(name):
        return False, "excluded name " + name[:30]
    if not T1_NAME_RE.search(name):
        seq = t.get((24, 32), "")
        et = t.get((24, 129), "")
        ok = seq in ("GR", "SE", "GR\\IR", "SE ", "GR ")
        try:
            et_ok = et.replace('"', "").split()[0].replace(" ", "")[:1] != "" and float(et.split("\\")[0].strip()) < 30
        except (ValueError, IndexError):
            et_ok = True
        if not (ok and et_ok):
            return False, "not T1W"
    return True, name[:40]


def scan_timepoint(sid, date, tpd):
    """Returns list of candidate T1c series for one MR timepoint dir."""
    txts = {}
    jpg_dir = None
    for f in os.listdir(tpd):
        full = os.path.join(tpd, f)
        if os.path.isdir(full):
            jpg_dir = full
            continue
        m = re.match(rf"{sid}_{date}_MR_(\d+)_(\d+)\.txt$", f)
        if m:
            txts.setdefault(m.group(1), []).append((int(m.group(2)), full))
    if jpg_dir is None:
        return []
    jpgs = set(os.listdir(jpg_dir))
    cand = []
    for ser, sl in txts.items():
        if not sl:
            continue
        first = min(sl)[1]
        t = parse_txt(first)
        ok, why = is_t1c(t)
        if not ok:
            continue
        rows = t.get((40, 16), ""); cols = t.get((40, 17), "")
        # count jpgs of this series
        pat = re.compile(rf"{sid}_{date}_MR_{ser}_(\d+)_(\d+)\.jpg$")
        nslices = 0
        for jf in jpgs:
            if pat.match(jf):
                nslices += 1
                break
        nslices_all = sum(1 for jf in jpgs if pat.match(jf))
        cand.append({"ser": ser, "txt_first": first, "n_slices": nslices_all,
                     "rows": rows, "cols": cols, "why": why,
                     "jpg_dir": jpg_dir})
    # prefer 3D (MRAcq) & most slices
    def score(c):
        t = parse_txt(c["txt_first"])
        acq = "3" if "3D" in t.get((24, 35), "") else "2"
        return (acq, c["n_slices"])
    cand.sort(key=score, reverse=True)
    return cand


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--limit-subj", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max-multi-subj", type=int, default=60)
    ap.add_argument("--max-single-subj", type=int, default=0)
    ap.add_argument("--max-subj-timepoints", type=int, default=6)
    ap.add_argument("--select-only", action="store_true")
    args = ap.parse_args()
    random.seed(args.seed)
    raw_path = os.path.join(ROOT, "data/manifests/lee_t1c_raw.jsonl")
    if args.select_only:
        rows = load_jsonl(raw_path)
        print(f"select-only: {len(rows)} raw rows")
    else:
        rows = []
        sids = [s for s in os.listdir(BASE) if os.path.isdir(os.path.join(BASE, s)) and not s.endswith(".complete")]
    if not args.select_only:
        if args.limit_subj:
            sids = sids[: args.limit_subj]
        for i, sid in enumerate(sorted(sids)):
            tpd_dir = os.path.join(BASE, sid)
            for tp in os.listdir(tpd_dir):
                if not re.match(r"\d{8}_MR_", tp) or tp.endswith(".complete"):
                    continue
                tpd = os.path.join(tpd_dir, tp)
                if not os.path.isdir(tpd):
                    continue
                date = tp[:8]
                for c in scan_timepoint(sid, date, tpd):
                    rows.append({"sid": sid, "date": date, "tp": tp,
                                 "jpg_dir": c["jpg_dir"], "ser": c["ser"], "n_slices": c["n_slices"],
                                 "txt_first": c["txt_first"], "why": c["why"],
                                 "key": f"lee_{sid}_{date}_s{c['ser']}"})
            if i % 100 == 0 and i:
                print(f"scanned {i}/{len(sids)} subjects, {len(rows)} T1c candidates", flush=True)
        save_jsonl(rows, raw_path)
        print(f"total T1c candidates: {len(rows)}")

    # subset selection: prefer subjects with more timepoints (longitudinal consistency)
    by_subj = {}
    for r in rows:
        by_subj.setdefault(r["sid"], []).append(r)
    for v in by_subj.values():
        v.sort(key=lambda x: x["date"])
    multi = {k: v for k, v in by_subj.items() if len(v) >= 2}
    single = {k: v for k, v in by_subj.items() if len(v) == 1}
    rnd = random.Random(args.seed)
    depth = {k: len(v) for k, v in multi.items()}
    cand = [k for k in multi if depth[k] >= 3]
    deep = sorted(cand, key=lambda k: -depth[k])
    n_deep = max(1, args.max_multi_subj * 2 // 3)
    pick = deep[:n_deep]
    rest = [k for k in cand if k not in pick]
    rnd.shuffle(rest)
    pick += rest[: args.max_multi_subj - len(pick)]
    sel = []
    for k in pick:
        v = multi[k]
        if len(v) > args.max_subj_timepoints and args.max_subj_timepoints >= 2:
            v = [v[int(j * (len(v) - 1) / (args.max_subj_timepoints - 1))] for j in range(args.max_subj_timepoints)]
        sel.extend(v)
    rnd.shuffle(list(single))
    for k in list(single)[: args.max_single_subj]:
        sel.extend(single[k])
    sel.sort(key=lambda x: (x["sid"], x["date"], int(x["ser"])))
    save_jsonl(sel, os.path.join(ROOT, "data/manifests/lee_t1c_selected.jsonl"))
    print(f"selected: {len(sel)} timepoints from {len(set(r['sid'] for r in sel))} subjects")


if __name__ == "__main__":
    main()