"""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, DATA, save_jsonl, load_jsonl, is_head_series, path

BASE = path("lee")
T1_NAME_RE = re.compile(r"t1|tfl|spgr|mp2rage|tse3d|vfl|mpage", re.I)
# fiesta/ciss/ssfp: balanced-SSFP (fluid-bright, T2-dominant) — their short TE
# fools the seq/TE fallback below, so they are excluded by name, not T1c
# tof: time-of-flight MRA (angiographic; "spgr" in the name matches T1_NAME_RE)
# sub/subtraction: post/pre subtraction images (KVP carries POST CONTRAST);
# "sub" = the token prefix followed by any non-letter (sub_s16, sub:3d,
# _SUB, sub. — but not "subject"/"subtraction", handled separately)
EXCL_RE = re.compile(r"\bt2\b|dwi|dti|mra|mrv|angi|swi|bold|\bpp2d|\bpp3d|perf|t2\*|t2star"
                     r"|fiesta|ciss|ssfp"
                     r"|(?<![a-z])tof(?![a-z])"
                     r"|sub(?![a-z])|subtraction", 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'"\(([0-9A-Fa-f]{4}),([0-9A-Fa-f]{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):
    # newer exports tag contrast only via (0018,0010) Agent ("Y GD", "POST
    # CONTRAST", ...); pre-contrast series leave it empty in all exports
    bolus = any(k in t for k in ((24, 4161), (24, 4168), (8, 307))) or bool(t.get((24, 16), "").strip())
    if not bolus:
        return False, "no bolus"
    # newer exports leave Sequence Name empty and put the protocol in the
    # series description, so match/exclude on all three
    name = (t.get((24, 36), "") + " " + t.get((24, 33), "") + " "
            + t.get((8, 0x103E), "")).lower()
    if EXCL_RE.search(name):
        return False, "excluded name " + name[:30]
    # FLAIR is usually T2-weighted (excluded); T1-FLAIR is a valid T1c
    if re.search(r"flair", name) and "t1" not in name:
        return False, "excluded name (FLAIR, not T1) " + name[:30]
    if not T1_NAME_RE.search(name):
        seq = t.get((24, 32), "").strip()
        et = t.get((24, 129), "")
        ok = seq in ("GR", "SE", "GR\\IR", "RM\\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]


MAX_SPACING = 4.0  # mm; reject series with any voxel spacing above this


def max_spacing(t):
    """Max voxel spacing (mm): (0028,0030) pixel spacing + the larger of
    (0018,0050) slice thickness and (0018,0088) spacing between slices
    (interleaved slices can be thinner than their center-to-center gap).
    Returns 0.0 if no spacing info is available."""
    vals = []
    ps = t.get((40, 48), "")
    if ps:
        try:
            vals += [float(x) for x in ps.split("\\")[:2]]
        except ValueError:
            pass
    for v in (t.get((24, 80), ""), t.get((24, 136), "")):
        if v:
            try:
                vals.append(float(v.split("\\")[0].strip()))
            except ValueError:
                pass
    return max(vals) if vals else 0.0


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
        if not is_head_series(t.get((8, 0x1030)), t.get((8, 0x103E))):
            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,
                     "max_sp": max_spacing(t),
                     "acq": "3" if "3D" in t.get((24, 35), "") else "2"})
    # series thicker than MAX_SPACING are the primary candidates only when this
    # timepoint has no thinner valid T1c; the others are flagged
    # thick_dropped: out of the selected manifest, but kept in raw so 04 can
    # fall back to them when every other T1c of the timepoint is lost
    if not cand:
        return []
    thin = [c for c in cand if c["max_sp"] <= MAX_SPACING]
    keep = thin if thin else [max(cand, key=lambda c: (c["acq"], c["n_slices"]))]
    dropped = [c for c in cand if c not in keep]
    for c in dropped:
        c["thick_dropped"] = True
    keep.sort(key=lambda c: (c["acq"], c["n_slices"]), reverse=True)
    return keep + dropped


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")
    ap.add_argument("--all", action="store_true",
                    help="select every head T1c candidate (no subject/timepoint caps)")
    args = ap.parse_args()
    random.seed(args.seed)
    raw_path = os.path.join(DATA, "manifests/lee_t1c_raw.jsonl")
    if args.select_only:
        rows = load_jsonl(raw_path)
        print(f"select-only: {len(rows)} raw rows")
        from concurrent.futures import ThreadPoolExecutor
        def is_head_row(r):
            t = parse_txt(r["txt_first"])
            return is_head_series(t.get((8, 0x1030)), t.get((8, 0x103E)))
        with ThreadPoolExecutor(max_workers=32) as ex:
            head = list(ex.map(is_head_row, rows))
        rows = [r for r, ok in zip(rows, head) if ok]
        rows = [r for r in rows if not r.get("thick_dropped")]
        print(f"head/brain rows: {len(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):
                    row = {"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 c.get("thick_dropped"):
                        row["thick_dropped"] = True
                    rows.append(row)
            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)}")

    # thick_dropped rows stay out of the selection; they remain in raw for 04's
    # last-resort pass on timepoints that otherwise yield no T1c volume
    base = [r for r in rows if not r.get("thick_dropped")]
    if args.all:
        # no caps: every head T1c candidate enters the dataset
        sel = list(base)
    else:
        # subset selection: prefer subjects with more timepoints (longitudinal consistency)
        by_subj = {}
        for r in base:
            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(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()