#!/usr/bin/env python
"""
Example usage of JDLSegmentor for myocardium and scar segmentation.

This script demonstrates how to use the ONNX-based JDL segmentor
to segment myocardium and scar tissue from cardiac MRI images.
"""

import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import sys

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))

from cardiac_toolkit.segmentation import JDLSegmentor


def example_basic_usage():
    """Example 1: Basic 2D segmentation."""
    print("=" * 80)
    print("Example 1: Basic 2D Segmentation")
    print("=" * 80)

    # Initialize segmentor
    # Default model paths (update with your actual paths)
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        print(f"Myocardium: {myocardium_model_path}")
        print(f"Scar: {scar_model_path}")
        print("Please update the model paths in this script.")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy image (replace with your own image loading)
    # In practice: image = np.load('your_cardiac_mri.npy')
    image = np.random.rand(256, 256).astype(np.float32) * 1000

    print(f"\nInput image shape: {image.shape}")
    print(f"Input image dtype: {image.dtype}")
    print(f"Input image range: [{image.min():.2f}, {image.max():.2f}]")

    # Perform segmentation
    print("\nPerforming segmentation...")
    myo_mask, scar_mask = segmentor.segment(image)

    print(f"\nMyocardium mask shape: {myo_mask.shape}")
    print(f"Myocardium pixels: {np.sum(myo_mask)}")
    print(f"Scar mask shape: {scar_mask.shape}")
    print(f"Scar pixels: {np.sum(scar_mask)}")

    return image, myo_mask, scar_mask


def example_with_bbox():
    """Example 2: Segmentation with bounding box."""
    print("\n" + "=" * 80)
    print("Example 2: Segmentation with Bounding Box")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy image
    image = np.random.rand(512, 512).astype(np.float32) * 1000

    # Define a bounding box around the heart region
    bbox = {
        'x1': 150,
        'y1': 150,
        'x2': 350,
        'y2': 350
    }

    print(f"\nInput image shape: {image.shape}")
    print(f"Bounding box: {bbox}")

    # Perform segmentation with bbox
    print("\nPerforming segmentation with bounding box...")
    myo_mask, scar_mask = segmentor.segment(
        image,
        bbox=bbox,
        bbox_scale_factor=1.1  # Expand bbox by 10%
    )

    print(f"\nMyocardium mask shape: {myo_mask.shape}")
    print(f"Myocardium pixels: {np.sum(myo_mask)}")
    print(f"Scar mask shape: {scar_mask.shape}")
    print(f"Scar pixels: {np.sum(scar_mask)}")

    return image, myo_mask, scar_mask, bbox


def example_batch_processing():
    """Example 3: Batch processing of multiple slices."""
    print("\n" + "=" * 80)
    print("Example 3: Batch Processing")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy batch of images (e.g., multiple slices)
    num_slices = 5
    images = [np.random.rand(256, 256).astype(np.float32) * 1000 for _ in range(num_slices)]

    print(f"\nNumber of slices: {num_slices}")
    print(f"Each slice shape: {images[0].shape}")

    # Process batch
    print("\nProcessing batch...")
    myo_masks, scar_masks = segmentor.segment_batch(images)

    print(f"\nProcessed {len(myo_masks)} slices")
    for i, (myo_mask, scar_mask) in enumerate(zip(myo_masks, scar_masks)):
        print(f"  Slice {i}: Myo pixels = {np.sum(myo_mask)}, Scar pixels = {np.sum(scar_mask)}")

    return images, myo_masks, scar_masks


def example_3d_volume():
    """Example 4: 3D volume segmentation."""
    print("\n" + "=" * 80)
    print("Example 4: 3D Volume Segmentation")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return None, None, None

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy 3D volume
    volume = np.random.rand(256, 256, 8).astype(np.float32) * 1000

    print(f"\nInput volume shape: {volume.shape}")

    # Segment 3D volume
    print("\nSegmenting 3D volume...")
    myo_mask_3d, scar_mask_3d = segmentor.segment_3d(
        volume,
        keep_largest_component='per_slice'  # or 'whole_volume'
    )

    print(f"\nMyocardium mask shape: {myo_mask_3d.shape}")
    print(f"Myocardium voxels: {np.sum(myo_mask_3d)}")
    print(f"Scar mask shape: {scar_mask_3d.shape}")
    print(f"Scar voxels: {np.sum(scar_mask_3d)}")

    return volume, myo_mask_3d, scar_mask_3d


def example_visualization():
    """Example 5: Visualizing results."""
    print("\n" + "=" * 80)
    print("Example 5: Visualization")
    print("=" * 80)

    # Initialize segmentor
    myocardium_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_myocardium.onnx"
    scar_model_path = "/gpfs/gibbs/project/kwan/jx332/code/2025-02-LGE-multiview-seg/checkpoints/jdl_scar.onnx"

    if not Path(myocardium_model_path).exists() or not Path(scar_model_path).exists():
        print(f"\n⚠ Models not found")
        return

    segmentor = JDLSegmentor(
        myocardium_model_path=myocardium_model_path,
        scar_model_path=scar_model_path
    )

    # Create dummy image
    image = np.random.rand(256, 256).astype(np.float32) * 1000

    # Perform segmentation
    print("\nPerforming segmentation...")
    myo_mask, scar_mask = segmentor.segment(image)

    # Create visualization
    print("\nCreating visualization...")
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))

    # Original image
    axes[0].imshow(image, cmap='gray')
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    # Myocardium mask
    axes[1].imshow(myo_mask, cmap='gray')
    axes[1].set_title('Myocardium Mask')
    axes[1].axis('off')

    # Scar mask
    axes[2].imshow(scar_mask, cmap='gray')
    axes[2].set_title('Scar Mask')
    axes[2].axis('off')

    # Overlay
    axes[3].imshow(image, cmap='gray')
    axes[3].imshow(myo_mask, cmap='Reds', alpha=0.3)
    axes[3].imshow(scar_mask, cmap='Blues', alpha=0.3)
    axes[3].set_title('Overlay\n(Red=Myo, Blue=Scar)')
    axes[3].axis('off')

    plt.tight_layout()

    # Save figure
    output_path = Path(__file__).parent / 'jdl_segmentation_result.png'
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"\nVisualization saved to: {output_path}")

    plt.close()


def main():
    """Run all examples."""
    print("\n" + "*" * 80)
    print("JDLSegmentor Examples")
    print("*" * 80)

    try:
        # Example 1: Basic usage
        example_basic_usage()

        # Example 2: With bounding box
        example_with_bbox()

        # Example 3: Batch processing
        example_batch_processing()

        # Example 4: 3D volume
        example_3d_volume()

        # Example 5: Visualization
        example_visualization()

        print("\n" + "=" * 80)
        print("All examples completed successfully!")
        print("=" * 80)
        print()

    except Exception as e:
        print(f"\n\nError running examples: {e}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()