#!/usr/bin/env python3
"""
Verification script for MP-PCA denoising (3D/4D).
Generates a 4D phantom, adds Rician noise, runs Python MPPCA,
calls Octave MPPCA, and compares.
"""

import subprocess
from pathlib import Path
import os

import numpy as np
import scipy.io as sio

from qmrpy.models.noise.denoising_mppca import MPPCA

ROOT_DIR = Path(__file__).parent.parent
DATA_DIR = ROOT_DIR / "data/parity"
OCTAVE_SCRIPT_PATH = ROOT_DIR / "scripts/octave/verify_mppca.m"


def generate_phantom(sx=20, sy=20, sz=5, n_vol=30, snr=20):
    """Generate a simple 4D diffusion-like phantom."""
    np.random.seed(42)
    
    # Structure: A box in the middle
    ground_truth = np.zeros((sx, sy, sz, n_vol))
    
    # 3 regions: Background (0), Tissue A (decay fast), Tissue B (decay slow)
    # Signal = S0 * exp(-b * D)
    # Let's just simulate random exponential decays
    b_values = np.linspace(0, 1000, n_vol)
    
    # Tissue A (Center box)
    cx, cy = sx // 2, sy // 2
    r = sx // 4
    mask_a = np.zeros((sx, sy, sz), dtype=bool)
    mask_a[cx-r:cx+r, cy-r:cy+r, 1:sz-1] = True
    
    adc_a = 0.001 # mm2/s
    s0_a = 1000.0
    sig_a = s0_a * np.exp(-b_values * adc_a)
    
    # Assign
    ground_truth[mask_a] = sig_a
    
    # Add Rician noise
    sigma = s0_a / snr
    
    # Rician = sqrt( (S + n1)^2 + n2^2 )
    n1 = np.random.normal(0, sigma, size=ground_truth.shape)
    n2 = np.random.normal(0, sigma, size=ground_truth.shape)
    
    noisy_data = np.sqrt((ground_truth + n1)**2 + n2**2)
    
    return noisy_data, ground_truth, sigma


def run_octave(mat_path_in, mat_path_out):
    """Call Octave script to run MP-PCA."""
    qMRLab_path = os.environ.get("QMRLAB_PATH")
    if not qMRLab_path:
        raise FileNotFoundError("qMRLab path is required. Set QMRLAB_PATH.")
    qMRLab_path = Path(qMRLab_path)
    if not qMRLab_path.exists():
        raise FileNotFoundError(f"qMRLab not found at: {qMRLab_path}")
    
    cmd = [
        "octave",
        "--no-gui",
        "--eval",
        f"qMRLab_path='{qMRLab_path}'; input_mat='{mat_path_in}'; output_mat='{mat_path_out}'; source('{OCTAVE_SCRIPT_PATH}');"
    ]
    print("Running Octave...")
    subprocess.check_call(cmd)


def main():
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    
    # 1. Generate Data
    print("Generating 4D phantom...")
    noisy, true_sig, true_sigma = generate_phantom(sx=20, sy=20, sz=10, n_vol=30, snr=20)
    
    # Save for Octave
    input_mat = DATA_DIR / "mppca_input.mat"
    output_mat = DATA_DIR / "mppca_octave.mat"
    
    sio.savemat(input_mat, {"Data4D": noisy})
    
    # 2. Run Python MPPCA
    print("Running Python MPPCA...")
    model = MPPCA(kernel=(5, 5, 5))
    res_py = model.fit(noisy)
    denoised_py = res_py["denoised"]
    sigma_py = res_py["sigma"]
    
    # 3. Run Octave MPPCA
    run_octave(input_mat, output_mat)
    
    # Load Octave results
    if not output_mat.exists():
        print("Octave output not found!")
        return
        
    res_oct = sio.loadmat(output_mat)
    denoised_oct = res_oct["denoised"]
    sigma_oct = res_oct["sigma"]
    
    # 4. Compare
    print("\n--- Verification: MP-PCA ---")
    
    # Compare center values (avoid boundaries)
    mid_x, mid_y, mid_z = 10, 10, 5
    print(f"Center Pixel ({mid_x},{mid_y},{mid_z}):")
    print(f"  Py Sigma: {sigma_py[mid_x, mid_y, mid_z]:.4f}")
    print(f"  Oct Sigma: {sigma_oct[mid_x, mid_y, mid_z]:.4f}")
    print(f"  Py Signal[0]: {denoised_py[mid_x, mid_y, mid_z, 0]:.4f}")
    print(f"  Oct Signal[0]: {denoised_oct[mid_x, mid_y, mid_z, 0]:.4f}")

    # ROI Comparison (exclude boundaries and qMRLab weirdness)
    # qMRLab zeros out boundaries widely.
    # We crop 3 pixels from each side (kernel half-width is 2, plus margin)
    sl = slice(3, -3)
    # Z-dimension is small (10). 3:-3 leaves 4 slices (3,4,5,6).
    roi = (sl, sl, sl)
    
    d_py_roi = denoised_py[roi]
    d_oct_roi = denoised_oct[roi]
    s_py_roi = sigma_py[roi]
    s_oct_roi = sigma_oct[roi]
    
    diff_sig = np.abs(d_py_roi - d_oct_roi)
    max_diff_sig = np.max(diff_sig)
    print(f"\nROI Max Diff Signal: {max_diff_sig:.6e}")
    
    diff_sigma = np.abs(s_py_roi - s_oct_roi)
    max_diff_sigma = np.max(diff_sigma)
    print(f"ROI Max Diff Sigma: {max_diff_sigma:.6e}")

    # Metrics
    if max_diff_sig < 1e-4:
        print("✅ Signal Parity OK")
    else:
        print("❌ Signal Parity FAIL")
        
    if max_diff_sigma < 1e-4:
        print("✅ Sigma Parity OK")
    else:
        print("❌ Sigma Parity FAIL")
    print(f"Center Pixel ({mid_x},{mid_y},{mid_z}):")
    print(f"  Py Sigma: {sigma_py[mid_x, mid_y, mid_z]:.4f}")
    print(f"  Oct Sigma: {sigma_oct[mid_x, mid_y, mid_z]:.4f}")
    print(f"  Py n_pars: {res_py['n_pars'][mid_x, mid_y, mid_z]:.1f}")
    print(f"  Py Signal[0]: {denoised_py[mid_x, mid_y, mid_z, 0]:.4f}")
    print(f"  Oct Signal[0]: {denoised_oct[mid_x, mid_y, mid_z, 0]:.4f}")

    # Metrics
    if max_diff_sig < 1e-4:
        print("✅ Signal Parity OK")
    else:
        print("❌ Signal Parity FAIL")
        
    if max_diff_sigma < 1e-4:
        print("✅ Sigma Parity OK")
    else:
        print("❌ Sigma Parity FAIL")


if __name__ == "__main__":
    main()
