"""Preprocess source volumes (T1c + optional label) into a uniform 1mm cropped/normalized form.

Usage:
  python scripts/preprocess.py --manifest data/manifests/train_labeled.jsonl --workers 16
Each row: {key, subject, date, img, label?, source}
Outputs:  data/proc/<key>.nii.gz, data/proc/<key>_label.nii.gz, data/procmeta/<key>.json
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import json
import argparse
import numpy as np
import SimpleITK as sitk
from concurrent.futures import ProcessPoolExecutor, as_completed
from src.common import ROOT, d, read_nii_img, head_mask_from_image, crop_box, normalize_volume, write_arr, save_jsonl

SPACING = (1.0, 1.0, 1.0)


def resample_volume(itk, spacing, interp):
    size = [max(1, int(round(s * sp / nsp))) for s, sp, nsp in zip(itk.GetSize(), itk.GetSpacing(), spacing)]
    ptype = sitk.sitkFloat32 if interp != sitk.sitkNearestNeighbor else itk.GetPixelID()
    return sitk.Resample(itk, size, sitk.Transform(), interp,
                         itk.GetOrigin(),
                         tuple(float(s) for s in spacing),
                         itk.GetDirection(),
                         0.0, ptype)


def worker(args):
    key, img_path, label_path, out_dir, meta_dir = args
    try:
        img_itk = read_nii_img(img_path)
        img1 = resample_volume(img_itk, SPACING, sitk.sitkLinear)
        a = sitk.GetArrayFromImage(img1)
        label1 = None
        if label_path and os.path.exists(label_path):
            lbl = read_nii_img(label_path)
            label1 = sitk.GetArrayFromImage(resample_volume(lbl, SPACING, sitk.sitkNearestNeighbor)).astype(np.uint8)
            label1 = (label1 > 0).astype(np.uint8)
        # ensure label1 grid == a grid
        if label1 is not None and label1.shape != a.shape:
            ref = sitk.GetImageFromArray(np.zeros(a.shape, dtype=np.uint8))
            ref.CopyInformation(img1)
            lbl2 = resample_volume(read_nii_img(label_path), SPACING, sitk.sitkNearestNeighbor)
            label1 = sitk.GetArrayFromImage(sitk.Resample(lbl2, ref, sitk.Transform(), sitk.sitkNearestNeighbor)).astype(np.uint8)
            label1 = (label1 > 0).astype(np.uint8)
            if label1.shape != a.shape:
                raise RuntimeError(f"label/image grid mismatch {label1.shape} vs {a.shape}")
        if a.max() <= a.min() + 1e-6:
            raise RuntimeError("empty volume")
        box = crop_box(a, label=label1, margin=12, cap=216, spacing=1.0)
        (s0, s1), (s2, s3), (s4, s5) = box
        ac_arr = a[s0:s1, s2:s3, s4:s5]
        lab_arr = label1[s0:s1, s2:s3, s4:s5] if label1 is not None else None
        m = head_mask_from_image(ac_arr)
        if m.sum() < 5000:
            raise RuntimeError("head mask too small")
        norm, (lo, hi) = normalize_volume(ac_arr, m)
        norm = norm.astype(np.float32)
        out = os.path.join(out_dir, key + ".nii.gz")
        write_arr(norm, out)
        meta = {"key": key, "img_spacing_in": list(img_itk.GetSpacing()), "img_size_in": list(img_itk.GetSize()),
                "origin": list(img_itk.GetOrigin()), "direction": list(img_itk.GetDirection()),
                "crop_vox": [int(s0), int(s2), int(s4)],
                "norm_lo": float(lo), "norm_hi": float(hi),
                "shape": list(norm.shape)}
        with open(os.path.join(meta_dir, key + ".json"), "w") as f:
            json.dump(meta, f)
        if lab_arr is not None:
            outl = os.path.join(out_dir, key + "_label.nii.gz")
            write_arr(lab_arr.astype(np.uint8), outl)
            meta["n_tumor_vox"] = int(lab_arr.sum())
            with open(os.path.join(meta_dir, key + ".json"), "w") as f:
                json.dump(meta, f)
        return key, None
    except Exception as e:  # noqa
        return key, repr(e)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--subset", type=int, default=0, help="only first N rows (0=all)")
    args = ap.parse_args()
    rows = [json.loads(l) for l in open(args.manifest) if l.strip()]
    if args.subset:
        rows = rows[: args.subset]
    out_dir = d("data/proc")
    meta_dir = d("data/procmeta")
    todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows
            if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz"))]
    print(f"rows={len(rows)} todo={len(todo)}")
    errs = []
    with ProcessPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(worker, t): t[0] for t in todo}
        done = 0
        for fu in as_completed(futs):
            k, err = fu.result()
            done += 1
            if err:
                errs.append((k, err))
                print("ERR", k, err, flush=True)
            if done % 50 == 0:
                print(f"  {done}/{len(todo)} done, {len(errs)} errs", flush=True)
    print(f"finished. errors: {len(errs)}")
    for k, e in errs[:20]:
        print("  ", k, e)


if __name__ == "__main__":
    main()