"""Patient-specific segmentation vs. atlas registration.

Runs a pretrained segmentation model (TotalSegmentator brain-structures, or
FreeSurfer SynthSeg) directly on a sample T1 MRI, labelling eloquent cortex,
ventricles, and deep structures in the patient's own space — no atlas
registration. It then resamples the Harvard-Oxford atlas onto the same grid and
reports the per-structure Dice overlap between the two label sources, and writes
a side-by-side overlay figure.

Choose the backend with SULCORA_SEG_BACKEND=totalsegmentator|synthseg|auto
(default: auto). The model weights download on first use.

NOTE: the work is wrapped in ``if __name__ == "__main__":`` because the
TotalSegmentator/nnU-Net backend uses spawn-based multiprocessing — its worker
processes re-import this module, so the heavy work must not run at import time.

RESEARCH SOFTWARE — NOT FOR CLINICAL USE. Demonstration on public MNI152 data.
"""
import os, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
from sulcora.io import load
from sulcora import segmentation as SEG


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    nii = os.path.join(HERE, "sample_data", "mni152_t1_2mm.nii.gz")
    if not os.path.exists(nii):
        import nibabel as nib
        from nilearn.datasets import load_mni152_template
        os.makedirs(os.path.dirname(nii), exist_ok=True)
        nib.save(load_mni152_template(resolution=2), nii)

    vol = load(nii)
    backend = os.environ.get("SULCORA_SEG_BACKEND", "auto")

    # ---- run the patient-specific segmentation ---------------------------- #
    print(f"Segmenting {os.path.basename(nii)} with backend={backend!r} ...")
    try:
        seg = SEG.segment_t1(nii, backend=backend)
    except SEG.BackendUnavailable as e:
        print("\n[segmentation backend unavailable]\n" + str(e))
        print('\nInstall the optional extra to run this demo:\n'
              '    pip install -e ".[seg]"\n'
              '  or, for the SynthSeg backend, install FreeSurfer (provides mri_synthseg).')
        sys.exit(0)

    by_group = {g: seg.present_structures(groups=g) for g in SEG.GROUPS}
    print(f"backend used: {seg.backend}")
    for g in SEG.GROUPS:
        print(f"  {g:9s}: {', '.join(by_group[g]) or '(none)'}")

    if not any(by_group.values()):
        print(
            "\n[no target structures segmented]\n"
            "The model ran and produced a valid label map, but none of the eloquent /\n"
            "ventricle / deep structures Sulcora scores were labelled on this input.\n"
            "TotalSegmentator's brain-structures model is trained on clinical ~1 mm T1\n"
            "MRI; on a smooth, downsampled public template (e.g. the 2 mm MNI152 sample)\n"
            "it commonly returns empty labels. For a meaningful comparison, run on a real\n"
            "~1 mm patient T1 (a GPU is strongly recommended — 1 mm CPU inference is slow).\n"
            "The segmentation pipeline, backend, mask/Dice logic, and atlas comparison are\n"
            "all exercised by the unit tests (tests/test_segmentation.py) independently of\n"
            "any specific scan.")
        sys.exit(0)

    # ---- compare against the atlas-based labels --------------------------- #
    print("\nPatient segmentation vs. Harvard-Oxford atlas (resampled to patient grid):")
    rows = SEG.compare_to_atlas(seg)
    shared = [r for r in rows if r["seg_voxels"] > 0 and r["ref_voxels"] > 0]
    seg_only = [r for r in rows if r["seg_voxels"] > 0 and r["ref_voxels"] == 0]
    atlas_only = [r for r in rows if r["seg_voxels"] == 0 and r["ref_voxels"] > 0]

    print("  -- structures labelled by BOTH (like-for-like Dice overlap) --")
    print(f"  {'structure':22s} {'group':9s} {'seg_vox':>8s} {'atlas_vox':>9s} {'Dice':>6s}")
    for r in sorted(shared, key=lambda r: (r["group"], r["structure"])):
        print(f"  {r['structure']:22s} {r['group']:9s} {r['seg_voxels']:8d} "
              f"{r['ref_voxels']:9d} {r['dice']:6.2f}")
    dvals = [r["dice"] for r in shared]
    if dvals:
        print(f"  mean Dice over {len(dvals)} shared structures: {np.mean(dvals):.2f}")
    if seg_only:
        print("  -- segmented by the model, no direct Harvard-Oxford equivalent: "
              + ", ".join(sorted(r["structure"] for r in seg_only)))
    if atlas_only:
        print("  -- in the atlas but not produced by this model/task: "
              + ", ".join(sorted(r["structure"] for r in atlas_only)))

    # ---- overlay figure: segmentation masks vs atlas masks ---------------- #
    from sulcora import atlas as A
    cort, clab = A.load_harvard_oxford("cort-maxprob-thr25-2mm"); cort = A.resample_atlas_to(cort, vol)
    sub, slab = A.load_harvard_oxford("sub-maxprob-thr25-2mm"); sub = A.resample_atlas_to(sub, vol)

    seg_vent = seg.structure_mask(groups="ventricle").data
    seg_deep = seg.structure_mask(groups="deep").data
    seg_elo = seg.structure_mask(groups="eloquent").data

    def atlas_group_mask(group):
        out = np.zeros(vol.shape, np.uint8)
        for spec in SEG._selected_specs(groups=group):
            out |= SEG.mask_from_labelmap(cort, clab, spec.patterns).data
            out |= SEG.mask_from_labelmap(sub, slab, spec.patterns).data
        return out

    atl_vent, atl_deep, atl_elo = (atlas_group_mask(g) for g in ("ventricle", "deep", "eloquent"))

    disp = vol.window().data
    z = vol.shape[2] // 2

    def panel(ax, title, masks):
        ax.imshow(disp[:, :, z].T, cmap="gray", origin="lower")
        for mask, cmap in masks:
            m = np.ma.masked_where(mask[:, :, z].T == 0, mask[:, :, z].T)
            ax.imshow(m, cmap=cmap, alpha=0.55, origin="lower")
        ax.set_title(title, color="#9fd0ff", fontsize=11); ax.axis("off")

    fig, ax = plt.subplots(1, 2, figsize=(10.5, 5.4), facecolor="#0b0f18")
    panel(ax[0], "Patient-specific segmentation\n(eloquent=green ventricle=cyan deep=magenta)",
          [(seg_elo, "summer"), (seg_vent, "cool"), (seg_deep, "spring")])
    panel(ax[1], "Atlas registration (Harvard-Oxford)\nsame labels, registered",
          [(atl_elo, "summer"), (atl_vent, "cool"), (atl_deep, "spring")])
    fig.suptitle("Registration-free patient segmentation vs. atlas labels", color="#dce6f5", fontsize=13)
    meanstr = f"mean Dice {np.mean(dvals):.2f} over {len(dvals)} structures" if dvals else "no shared structures"
    fig.text(0.5, 0.02, f"backend: {seg.backend}  ·  {meanstr}  ·  "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE",
             color="#c0392b", ha="center", fontsize=9)
    fig.tight_layout(rect=[0, 0.05, 1, 0.95])
    out = os.path.join(docs, "segmentation_vs_atlas.png")
    fig.savefig(out, dpi=130, facecolor=fig.get_facecolor()); plt.close(fig)
    print("\nwrote", out, "\nSEGMENTATION PROOF OK")


if __name__ == "__main__":
    import multiprocessing as mp
    mp.freeze_support()
    main()
