"""MR->MR control for the synthmorph investigation.

The benchmark pairs are CT->T1c (out-of-domain for the MRI-trained
SynthMorph model).  To separate the cross-modality effect from the other
suspects (extent truncation of the CT FOV, min-max normalization), this
registers two T1c scans of the SAME subject (m6_24V62QQM, 2023-02-23 and
2023-10-17 -- in-domain, MRI->MRI, same tightly-cropped geometry as the
benchmark's fixed image) and scores the result with the benchmark scoring
(MI over the head, ImpactLoss, skull-ring gap) plus the transform
decomposition against a SimpleITK reference.

Usage:
    python scripts/investigate_synthmorph_mm.py [--out results/synthmorph_investigation]
"""
import argparse
import json
import os
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import numpy as np
import SimpleITK as sitk

from src.common import read_nii_img, head_mask_from_image
from src import registration_pipelines as RP
from scripts.investigate_synthmorph import (
    _to_affine, _angle_between_rotations, head_centroid_world, BR, M6)

KEY_A = "m6_24V62QQM_20230223"
KEY_B = "m6_24V62QQM_20231017"


def score_pair(img_moving, img_fixed, T):
    """Benchmark scoring of the warped moving (in the fixed grid) ~ fixed."""
    ha = sitk.GetArrayFromImage(img_fixed).astype("float32")
    wc = sitk.GetArrayFromImage(sitk.Resample(img_moving, img_fixed, T,
                                              sitk.sitkLinear)).astype("float32")
    gap = M6._skull_ring_gap(wc, ha)
    mi = BR._mi_masked(wc, ha, head_mask_from_image(ha))
    il = BR._impact_loss_value(wc, ha)
    return {"mi": mi, "impact_loss": il, "gap": gap}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.path.join(ROOT, "results",
                                                  "synthmorph_investigation"))
    ap.add_argument("--gpu", type=int, default=0)
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    labeled, _ = M6.discover()
    bykey = {r["key"]: r for r in labeled}
    a, b = bykey[KEY_A], bykey[KEY_B]
    img_a = read_nii_img(a["img"])
    img_b = read_nii_img(b["img"])
    out = {"pair": [KEY_A, KEY_B], "note": ("same-subject T1c->T1c, 8 months "
                                            "apart; in-domain control for the "
                                            "CT->T1c benchmark pairs")}

    import math
    c_fix = head_centroid_world(img_b)
    T_ref, meta_ref = RP.register_simpleitk(img_a, img_b, threads=16,
                                            mode="benchmark")
    sref = score_pair(img_a, img_b, T_ref)
    R_ref, t_ref = _to_affine(T_ref)
    th_ref = math.degrees(float(np.clip((np.trace(R_ref) - 1) / 2, -1, 1)))
    out["reference"] = {
        "tool": meta_ref["toolkit"], "R": R_ref.tolist(), "t": t_ref.tolist(),
        "mi": sref["mi"], "impact_loss": sref["impact_loss"], "gap": sref["gap"],
        "rotation_deg": round(th_ref, 2),
        "head_centre_fix": [round(float(x), 1)
                            for x in (R_ref @ c_fix + t_ref)],
    }

    t0 = time.time()
    T_sm, meta_sm = RP.register_synthmorph(img_a, img_b, model="rigid",
                                           gpu=args.gpu, extent=256)
    ssm = score_pair(img_a, img_b, T_sm)
    R_sm, t_sm = _to_affine(T_sm)
    out["synthmorph_rigid"] = {
        "seconds": round(time.time() - t0, 1),
        "auto_inverted": meta_sm["auto_inverted"],
        "mi": ssm["mi"], "impact_loss": ssm["impact_loss"], "gap": ssm["gap"],
        "rot_err_vs_ref_deg": round(_angle_between_rotations(R_ref, R_sm), 2),
        "head_centre_err_mm": round(
            float(np.linalg.norm((R_sm @ c_fix + t_sm) - (R_ref @ c_fix + t_ref))), 1),
    }
    # no-clip sanity: MR inputs are never clipped by the adapter
    fp = os.path.join(args.out, "partE_mr_mr.json")
    with open(fp, "w") as f:
        json.dump(out, f, indent=1)
    print(json.dumps(out, indent=1))
    print("wrote", fp)


if __name__ == "__main__":
    main()