#!/usr/bin/env python3
"""Profiling script for SDF loss functions.

This script profiles memory usage, timing, and device transfers for the
SDF-weighted loss functions. It can be run locally or on a GPU cluster.

Usage:
    # Local CPU profiling
    python scripts/profile_loss.py --device cpu

    # Local GPU profiling
    python scripts/profile_loss.py --device cuda

    # GPU cluster (with specific GPU)
    CUDA_VISIBLE_DEVICES=0 python scripts/profile_loss.py --device cuda

    # Full profiling with PyTorch profiler
    python scripts/profile_loss.py --device cuda --full-profile

    # Custom batch size and image size
    python scripts/profile_loss.py --device cuda --batch-size 16 --image-size 512
"""

import argparse
import time
import gc
from contextlib import contextmanager
from dataclasses import dataclass

import torch
from torch.profiler import profile, ProfilerActivity, record_function

# Add parent directory to path for imports
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from sdf_loss import DiceLoss, SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss


@dataclass
class ProfilingResult:
    """Container for profiling results."""

    loss_name: str
    device: str
    batch_size: int
    image_size: int
    forward_time_ms: float
    backward_time_ms: float
    total_time_ms: float
    peak_memory_mb: float
    memory_allocated_mb: float
    memory_reserved_mb: float


@contextmanager
def cuda_memory_tracking():
    """Context manager for tracking CUDA memory."""
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()
    yield
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def create_test_data(batch_size: int, height: int, width: int, device: str):
    """Create test prediction and target tensors."""
    # Create realistic oil spill-like target
    target = torch.zeros(batch_size, 1, height, width, device=device)
    for b in range(batch_size):
        cx, cy = height // 2, width // 2
        rx, ry = height // 4, width // 4
        y, x = torch.meshgrid(
            torch.arange(height, device=device),
            torch.arange(width, device=device),
            indexing="ij",
        )
        mask = ((x - cx) ** 2 / (rx**2) + (y - cy) ** 2 / (ry**2)) < 1
        target[b, 0] = mask.float()

    # Create prediction (logits)
    pred = torch.randn(batch_size, 1, height, width, device=device, requires_grad=True)

    return pred, target


def profile_loss(
    loss_fn,
    loss_name: str,
    batch_size: int,
    image_size: int,
    device: str,
    num_warmup: int = 3,
    num_iterations: int = 10,
) -> ProfilingResult:
    """Profile a single loss function."""
    height = width = image_size

    # Warmup
    for _ in range(num_warmup):
        pred, target = create_test_data(batch_size, height, width, device)
        loss = loss_fn(pred, target)
        loss.backward()
        del pred, target, loss
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()

    # Reset memory stats after warmup
    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()

    forward_times = []
    backward_times = []

    for _ in range(num_iterations):
        pred, target = create_test_data(batch_size, height, width, device)

        # Time forward pass
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()

        loss = loss_fn(pred, target)

        if device == "cuda":
            torch.cuda.synchronize()
        forward_time = time.perf_counter() - start
        forward_times.append(forward_time)

        # Time backward pass
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()

        loss.backward()

        if device == "cuda":
            torch.cuda.synchronize()
        backward_time = time.perf_counter() - start
        backward_times.append(backward_time)

        del pred, target, loss

    # Collect memory stats
    if device == "cuda":
        torch.cuda.synchronize()
        peak_memory = torch.cuda.max_memory_allocated() / 1024 / 1024
        memory_allocated = torch.cuda.memory_allocated() / 1024 / 1024
        memory_reserved = torch.cuda.memory_reserved() / 1024 / 1024
    else:
        peak_memory = 0
        memory_allocated = 0
        memory_reserved = 0

    avg_forward = sum(forward_times) / len(forward_times) * 1000
    avg_backward = sum(backward_times) / len(backward_times) * 1000

    return ProfilingResult(
        loss_name=loss_name,
        device=device,
        batch_size=batch_size,
        image_size=image_size,
        forward_time_ms=avg_forward,
        backward_time_ms=avg_backward,
        total_time_ms=avg_forward + avg_backward,
        peak_memory_mb=peak_memory,
        memory_allocated_mb=memory_allocated,
        memory_reserved_mb=memory_reserved,
    )


def run_pytorch_profiler(
    loss_fn, loss_name: str, batch_size: int, image_size: int, device: str
):
    """Run PyTorch profiler for detailed analysis."""
    height = width = image_size
    pred, target = create_test_data(batch_size, height, width, device)

    activities = [ProfilerActivity.CPU]
    if device == "cuda":
        activities.append(ProfilerActivity.CUDA)

    with profile(
        activities=activities,
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    ) as prof:
        with record_function(f"{loss_name}_forward"):
            loss = loss_fn(pred, target)

        with record_function(f"{loss_name}_backward"):
            loss.backward()

    return prof


def check_device_transfers(loss_fn, batch_size: int, image_size: int, device: str):
    """Check for unexpected device transfers during loss computation."""
    if device != "cuda":
        print("  Device transfer check only available for CUDA")
        return

    height = width = image_size
    pred, target = create_test_data(batch_size, height, width, device)

    # Use profiler to detect CPU-GPU transfers
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
    ) as prof:
        loss = loss_fn(pred, target)
        loss.backward()

    # Look for cudaMemcpy events (indicates transfers)
    events = prof.key_averages()
    transfer_events = [
        e for e in events if "cudaMemcpy" in e.key or "to" in e.key.lower()
    ]

    if transfer_events:
        print("  WARNING: Potential device transfers detected:")
        for event in transfer_events[:5]:  # Show first 5
            print(f"    - {event.key}: {event.cpu_time_total / 1000:.2f}ms")
    else:
        print("  No unexpected device transfers detected")


def print_results(results: list[ProfilingResult]):
    """Print profiling results in a formatted table."""
    print("\n" + "=" * 100)
    print("PROFILING RESULTS")
    print("=" * 100)

    # Header
    print(
        f"{'Loss':<25} {'Device':<8} {'Batch':<6} {'Size':<6} "
        f"{'Forward':<12} {'Backward':<12} {'Total':<12} {'Peak Mem':<12}"
    )
    print("-" * 100)

    for r in results:
        print(
            f"{r.loss_name:<25} {r.device:<8} {r.batch_size:<6} {r.image_size:<6} "
            f"{r.forward_time_ms:>8.2f} ms  {r.backward_time_ms:>8.2f} ms  "
            f"{r.total_time_ms:>8.2f} ms  {r.peak_memory_mb:>8.1f} MB"
        )

    print("=" * 100)


def main():
    parser = argparse.ArgumentParser(description="Profile SDF loss functions")
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        choices=["cpu", "cuda"],
        help="Device to run profiling on",
    )
    parser.add_argument(
        "--batch-size", type=int, default=8, help="Batch size for profiling"
    )
    parser.add_argument(
        "--image-size", type=int, default=256, help="Image size (height=width)"
    )
    parser.add_argument(
        "--num-warmup", type=int, default=3, help="Number of warmup iterations"
    )
    parser.add_argument(
        "--num-iterations", type=int, default=10, help="Number of profiling iterations"
    )
    parser.add_argument(
        "--full-profile",
        action="store_true",
        help="Run full PyTorch profiler (generates detailed trace)",
    )
    parser.add_argument(
        "--check-transfers",
        action="store_true",
        help="Check for device transfers",
    )
    parser.add_argument(
        "--output-trace",
        type=str,
        default=None,
        help="Output path for Chrome trace (only with --full-profile)",
    )

    args = parser.parse_args()

    # Check device availability
    if args.device == "cuda" and not torch.cuda.is_available():
        print("CUDA not available, falling back to CPU")
        args.device = "cpu"

    print("Profiling SDF Loss Functions")
    print(f"Device: {args.device}")
    if args.device == "cuda":
        print(f"GPU: {torch.cuda.get_device_name()}")
    print(f"Batch size: {args.batch_size}")
    print(f"Image size: {args.image_size}x{args.image_size}")
    print(f"Warmup iterations: {args.num_warmup}")
    print(f"Profiling iterations: {args.num_iterations}")
    print()

    # Loss functions to profile
    losses = [
        ("DiceLoss (baseline)", DiceLoss(from_logits=True)),
        ("SDFWeightedBCELoss", SDFWeightedBCELoss()),
        ("SDFWeightedDiceLoss", SDFWeightedDiceLoss()),
        ("DiSCoLoss (full)", DiSCoLoss(baseloss_weight=1, sdfweighted_weight=1)),
        ("DiSCoLoss (SDF only)", DiSCoLoss(baseloss_weight=0, sdfweighted_weight=1)),
    ]

    results = []

    for name, loss_fn in losses:
        print(f"Profiling {name}...")

        # Check for device transfers
        if args.check_transfers:
            print("  Checking device transfers...")
            check_device_transfers(
                loss_fn, args.batch_size, args.image_size, args.device
            )

        # Run profiling
        result = profile_loss(
            loss_fn,
            name,
            args.batch_size,
            args.image_size,
            args.device,
            args.num_warmup,
            args.num_iterations,
        )
        results.append(result)
        print(
            f"  Forward: {result.forward_time_ms:.2f}ms, Backward: {result.backward_time_ms:.2f}ms"
        )

        # Full PyTorch profiler
        if args.full_profile:
            print("  Running full profiler...")
            prof = run_pytorch_profiler(
                loss_fn, name, args.batch_size, args.image_size, args.device
            )

            # Print top operations
            print("  Top operations by CPU time:")
            print(
                prof.key_averages().table(
                    sort_by="cpu_time_total", row_limit=5, top_level_events_only=True
                )
            )

            # Export trace if requested
            if args.output_trace:
                trace_path = f"{args.output_trace}_{name.replace(' ', '_')}.json"
                prof.export_chrome_trace(trace_path)
                print(f"  Trace exported to: {trace_path}")

        gc.collect()
        if args.device == "cuda":
            torch.cuda.empty_cache()

    # Print summary
    print_results(results)

    # Memory summary for CUDA
    if args.device == "cuda":
        print("\nCUDA Memory Summary:")
        print(torch.cuda.memory_summary())


if __name__ == "__main__":
    main()
