#!/usr/bin/env python3
"""
Probe script: find the exact Euler angle mapping SV-DRR uses.

For each test view (0, 404, 1239), we try multiple permutations of how the
two coordinate values [c0, c1] from camera_views.json map to DiffDRR's
3-component Euler angles rot=[rz, rx, ry] (convention="ZXY"), with xyz=[0, 1800, 0].

NO look_at_rotation is used — raw Euler angles are injected directly.

Output: probe_output/{view_id}_{perm_name}.png   (+ SV-DRR ground truth copy)
"""

import json
import math
import sys
from pathlib import Path

import numpy as np
from PIL import Image

SCRIPT_DIR = Path(__file__).resolve().parent
DIFFDRR_DIR = SCRIPT_DIR / "DiffDRR"
if str(DIFFDRR_DIR) not in sys.path:
    sys.path.insert(0, str(DIFFDRR_DIR))

import torch

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
CT_CLEAN = SCRIPT_DIR / "data" / "lidc_TotalSeg_test" / "LIDC-IDRI-0001" / "01_body" / "CT_clean.nii.gz"
MASK_NII = SCRIPT_DIR / "data" / "lidc_TotalSeg_test" / "LIDC-IDRI-0001" / "02_totalseg" / "mask.nii.gz"
CAMERA_JSON = SCRIPT_DIR.parent / "wt_dataset" / "LIDC_IDRI" / "img_complex_fb_256" / "camera_views.json"
SVDRR_DIR = SCRIPT_DIR.parent / "wt_dataset" / "LIDC_IDRI" / "img_complex_fb_256" / "LIDC-IDRI-0001"
OUT_DIR = SCRIPT_DIR / "probe_output"

SDD = 2000.0
HEIGHT = 256
DELX = 2.0
RADIUS = 1800.0
EULER_CONV = "ZXY"

VIEW_INDICES = [0, 404, 1239]


def save_uint8_png(arr_2d: np.ndarray, path: Path, do_flip: bool = True):
    """Percentile-clip, optional horizontal flip, save as 8-bit PNG."""
    if do_flip:
        arr_2d = np.ascontiguousarray(np.flip(arr_2d, axis=-1))
    lo = float(np.percentile(arr_2d, 0.5))
    hi = float(np.percentile(arr_2d, 99.5))
    if hi - lo > 0:
        normed = np.clip(arr_2d, lo, hi)
        normed = (normed - lo) / (hi - lo)
    else:
        normed = np.zeros_like(arr_2d)
    uint8 = (normed * 255).astype(np.uint8)
    Image.fromarray(uint8, mode="L").save(str(path))


def main():
    from diffdrr.data import read
    from diffdrr.drr import DRR
    from diffdrr.pose import convert

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    # Load camera views
    with open(CAMERA_JSON) as f:
        cam_views = json.load(f)

    # Load volume (no labelmap needed for plain DRR)
    print("Loading CT volume...")
    subject = read(
        volume=str(CT_CLEAN),
        labelmap=str(MASK_NII),
        orientation="AP",
        center_volume=True,
    )

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    drr_module = DRR(
        subject,
        sdd=SDD,
        height=HEIGHT,
        delx=DELX,
    ).to(device)

    # ---------------------------------------------------------------------------
    # Define permutations: how [c0, c1] map to rot=[rz, rx, ry] for ZXY conv.
    # Convention="ZXY" means: rot is applied as Rz(rot[0]) @ Rx(rot[1]) @ Ry(rot[2])
    # ---------------------------------------------------------------------------
    def make_permutations(c0, c1):
        """Return dict of perm_name -> [rz, rx, ry] for ZXY convention."""
        return {
            # Basic placements of two angles into three slots
            "A_c0_0_c1":     [c0, 0.0, c1],
            "B_c0_c1_0":     [c0, c1, 0.0],
            "C_c1_c0_0":     [c1, c0, 0.0],
            "D_0_c0_c1":     [0.0, c0, c1],
            "E_0_c1_c0":     [0.0, c1, c0],
            "F_c1_0_c0":     [c1, 0.0, c0],
            # Negative sign variants (common in C-arm: pitch inversion)
            "G_nc0_0_c1":    [-c0, 0.0, c1],
            "H_c0_0_nc1":    [c0, 0.0, -c1],
            "I_nc0_0_nc1":   [-c0, 0.0, -c1],
            "J_0_nc0_c1":    [0.0, -c0, c1],
            "K_0_c0_nc1":    [0.0, c0, -c1],
            "L_0_nc0_nc1":   [0.0, -c0, -c1],
            # Sign flips on the other placements
            "M_nc1_c0_0":    [-c1, c0, 0.0],
            "N_c1_nc0_0":    [c1, -c0, 0.0],
            "O_nc1_nc0_0":   [-c1, -c0, 0.0],
            "P_nc0_c1_0":    [-c0, c1, 0.0],
            "Q_c0_nc1_0":    [c0, -c1, 0.0],
            "R_nc0_nc1_0":   [-c0, -c1, 0.0],
        }

    xyz_fixed = torch.tensor([[0.0, RADIUS, 0.0]], dtype=torch.float32, device=device)

    for view_idx in VIEW_INDICES:
        cv = cam_views[view_idx]
        c0, c1 = cv["coordinate"]
        svdrr_id = cv["id"]
        print(f"\n--- View {view_idx} (svdrr_id={svdrr_id}): c0={c0:.6f}, c1={c1:.6f} ---")

        # Copy SV-DRR ground truth for comparison
        gt_path = SVDRR_DIR / f"{svdrr_id}.png"
        if gt_path.exists():
            gt_img = Image.open(gt_path)
            gt_img.save(str(OUT_DIR / f"{view_idx:04d}_GT.png"))
            print(f"  Saved GT: {view_idx:04d}_GT.png")

        perms = make_permutations(c0, c1)
        for pname, rot_vec in perms.items():
            rot_t = torch.tensor([rot_vec], dtype=torch.float32, device=device)

            pose = convert(
                rot_t, xyz_fixed,
                parameterization="euler_angles",
                convention=EULER_CONV,
            )

            with torch.no_grad():
                img = drr_module(pose)  # (1, 1, H, W)

            img_np = img.squeeze().cpu().numpy()  # (H, W)

            out_name = f"{view_idx:04d}_{pname}.png"
            save_uint8_png(img_np, OUT_DIR / out_name, do_flip=True)

        print(f"  Saved {len(perms)} permutation images for view {view_idx}")

    # Also render view 0 with rot=[0,0,0] as sanity check
    print("\n--- View 0 sanity check: rot=[0, 0, 0] ---")
    rot_zero = torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float32, device=device)
    pose_zero = convert(rot_zero, xyz_fixed, parameterization="euler_angles", convention=EULER_CONV)
    with torch.no_grad():
        img_zero = drr_module(pose_zero)
    save_uint8_png(img_zero.squeeze().cpu().numpy(), OUT_DIR / "0000_rot000.png", do_flip=True)
    print("  Saved 0000_rot000.png")

    # Cleanup
    del drr_module, subject
    torch.cuda.empty_cache()

    print(f"\nAll probe images saved to: {OUT_DIR}/")
    print("Compare *_GT.png (ground truth) with permutation PNGs to find the match.")


if __name__ == "__main__":
    main()