"""Reconstruct lee T1c volumes from JPG slices + DICOM txt metadata.

The timepoint dir holds only sampled per-slice txts (typically slices 1, 2, N);
slice positions are reconstructed by fitting a linear IPP(s) model, validated
against all available samples.
Usage: python scripts/04_reconstruct_lee.py [--manifest ...selected.jsonl] [--workers 48]
Output: data/lee_nii/<key>.nii.gz (uint8, native grid)
"""
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 numpy as np
import SimpleITK as sitk
from concurrent.futures import ProcessPoolExecutor, as_completed
from PIL import Image
from src.common import ROOT, d, load_jsonl


LINE_RE = re.compile(
    r'^\s*"\((\d{4}),(\d{4})\)","([A-Z]+)","(\d+)",(?:\"([^\"]*)\"|([^,]+)),"')


def parse_txt(p):
    ipp = spacing = iop = rows = cols = None
    try:
        with open(p, errors="ignore") as f:
            for raw in f:
                m = LINE_RE.match(raw)
                if not m:
                    continue
                tag = (int(m.group(1), 16), int(m.group(2), 16))
                if tag not in ((32, 50), (40, 48), (32, 55), (40, 16), (40, 17)):
                    continue
                val = (m.group(5) if m.group(5) is not None else m.group(6)).strip()
                if not val:
                    continue
                try:
                    if tag == (32, 50):
                        ipp = np.array([float(x) for x in val.split("\\")[:3]])
                    elif tag == (40, 48):
                        spacing = np.array([float(x) for x in val.split("\\")[:2]])
                    elif tag == (32, 55):
                        iop = np.array([float(x) for x in val.split("\\")[:6]])
                    elif tag == (40, 16):
                        rows = int(float(val))
                    elif tag == (40, 17):
                        cols = int(float(val))
                except (ValueError, IndexError):
                    continue
        if ipp is None or spacing is None or iop is None or rows is None or cols is None:
            return None
        return ipp, spacing, iop, rows, cols
    except OSError:
        return None


def reconstruct(row):
    sid, date, ser, jpg_dir = row["sid"], row["date"], row["ser"], row["jpg_dir"]
    out = os.path.join(d("data/lee_nii"), row["key"] + ".nii.gz")
    if os.path.exists(out):
        return row["key"], True, "cached"
    pat = re.compile(rf"^{re.escape(sid)}_{date}_MR_{ser}_(\d+)_(\d+)\.jpg$")
    tpat = re.compile(rf"^{re.escape(sid)}_{date}_MR_{ser}_(\d+)\.txt$")
    tp_dir = os.path.dirname(jpg_dir)
    jpg_map, txt_samples = {}, []
    try:
        for f in os.listdir(jpg_dir):
            m = pat.match(f)
            if m:
                sl, inst = int(m.group(1)), int(m.group(2))
                if sl not in jpg_map or inst < jpg_map[sl][1]:
                    jpg_map[sl] = (f, inst)
        for f in os.listdir(tp_dir):
            m = tpat.match(f)
            if m:
                p = parse_txt(os.path.join(tp_dir, f))
                if p is not None:
                    txt_samples.append((int(m.group(1)), p))
    except OSError as e:
        return row["key"], False, f"listdir fail {e!r}"
    if not jpg_map or not txt_samples:
        return row["key"], False, f"no jpg({len(jpg_map)}) or txt({len(txt_samples)})"
    txt_samples.sort()
    # majority (rows, cols, pixel-spacing) as reference geometry
    from collections import Counter
    geo = Counter((p[3], p[4], round(p[1][0], 4), round(p[1][1], 4)) for _, p in txt_samples)
    (rows, cols, ps0, ps1) = geo.most_common(1)[0][0]
    ref = [p for _, p in txt_samples if (p[3], p[4]) == (rows, cols) and np.allclose(p[1], (ps0, ps1), atol=1e-3)]
    good_idx = [sl for sl, p in txt_samples if (p[3], p[4]) == (rows, cols)]
    if len(good_idx) < 2 or not ref:
        return row["key"], False, f"only {len(good_idx)} consistent geometry samples"
    ipp0, ps, iop0, _, _ = ref[len(ref) // 2]
    # linear IPP model (fit on geometry-consistent samples only)
    good = [(sl, p) for sl, p in txt_samples if p[3] == rows and p[4] == cols and np.allclose(p[1], ps, atol=1e-3)]
    if len(good) < 2:
        return row["key"], False, "fewer than 2 consistent geometry samples"
    (s_lo, p_lo), (s_hi, p_hi) = good[0], good[-1]
    step = (p_hi[0] - p_lo[0]) / max(s_hi - s_lo, 1)
    resid = max(np.linalg.norm(p[0] - p_lo[0] - (s - s_lo) * step) for s, p in good)
    if resid > 0.35 * np.linalg.norm(step):
        return row["key"], False, f"non-linear slice positions resid={resid:.3f}"
    u, v = iop0[:3], iop0[3:]
    n = np.cross(u, v)
    # the stack direction must follow the nominal slice normal (rejects multiplanar 2D exports)
    if abs(float(np.dot(step, n))) < 0.6 * float(np.linalg.norm(step)):
        return row["key"], False, "stack not along slice normal"
    sl_list = sorted(jpg_map)

    def ipp_of(s):
        return p_lo[0] + (s - s_lo) * step
    ipp_all = {sl: ipp_of(sl) for sl in sl_list}
    r0 = min(sl_list, key=lambda s: np.dot(ipp_all[s] - ipp_all[sl_list[0]], n))
    p_ref = ipp_all[r0]
    zs, pos = [], []
    for sl in sl_list:
        off = ipp_all[sl] - p_ref
        zs.append(float(np.dot(off, n)))
        pos.append((int(round(float(np.dot(off, u)) / ps[0])), int(round(float(np.dot(off, v)) / ps[1]))))
    zs = np.array(zs)
    dmed = np.median(np.abs(np.diff(np.sort(zs))))
    if dmed <= 0 or not np.isfinite(dmed):
        return row["key"], False, "bad slice spacing"
    nz = int(round((zs.max() - zs.min()) / dmed)) + 1
    ro_min = min(p[0] for p in pos)
    ro_max = max(p[0] for p in pos)
    co_min = min(p[1] for p in pos)
    co_max = max(p[1] for p in pos)
    vol = np.zeros((rows + (ro_max - ro_min), cols + (co_max - co_min), nz), dtype=np.uint8)
    nread = 0
    for idx, (sl, jf, (ro, co)) in enumerate(zip(sl_list, [jpg_map[s][0] for s in sl_list], pos)):
        zi = int(round(zs[idx] / dmed))
        try:
            im = np.asarray(Image.open(os.path.join(jpg_dir, jf)).convert("L"), dtype=np.uint8)
        except Exception:
            continue
        if im.shape[:2] != (rows, cols):
            if im.shape[:2] == (cols, rows):
                im = im.T.copy()
            else:
                continue
        vol[ro - ro_min: ro - ro_min + rows, co - co_min: co - co_min + cols, zi] = im
        nread += 1
    if nread < 0.95 * len(sl_list) or vol.max() == 0:
        return row["key"], False, f"only {nread}/{len(sl_list)} slices read"
    ext = (vol.shape[0] * ps[0], vol.shape[1] * ps[1], vol.shape[2] * dmed)
    if min(ext) < 60 or max(ext) > 350:
        return row["key"], False, f"plausible extent failed {tuple(round(e, 1) for e in ext)}mm"
    origin = p_ref + (ro_min * ps[0]) * u + (co_min * ps[1]) * v
    direction = tuple(float(x) for x in np.concatenate([u, v, n]))
    img = sitk.GetImageFromArray(vol)
    img.SetOrigin(tuple(float(x) for x in origin))
    img.SetSpacing((float(ps[0]), float(ps[1]), float(dmed)))
    img.SetDirection(direction)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    sitk.WriteImage(img, out, True)
    return row["key"], True, f"{vol.shape} n={nread}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", default=os.path.join(ROOT, "data/manifests/lee_t1c_selected.jsonl"))
    ap.add_argument("--workers", type=int, default=48)
    args = ap.parse_args()
    rows = load_jsonl(args.manifest)
    rows = [r for r in rows if not os.path.exists(os.path.join(d("data/lee_nii"), r["key"] + ".nii.gz"))]
    print(f"todo={len(rows)} workers={args.workers}")
    ok = err = 0
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(reconstruct, r): r for r in rows}
        for i, fu in enumerate(as_completed(futs), 1):
            try:
                k, good, msg = fu.result()
            except Exception as e:  # noqa
                k, good, msg = futs[fu]["key"], False, repr(e)
            if good:
                ok += 1
            else:
                err += 1
                if err <= 40:
                    print(" ERR", k, msg, flush=True)
            if i % 100 == 0:
                print(f"  {i}/{len(futs)} ok={ok} err={err}", flush=True)
    print(f"done ok={ok} err={err}")


if __name__ == "__main__":
    main()