from __future__ import annotations

"""PyTorch Profiler wrapper for training performance analysis.

Profiles the first N batches of training to identify bottlenecks
in CUDA/CPU execution, memory allocation, and data loading.

Exports:
  - Chrome trace (viewable in chrome://tracing)
  - TensorBoard plugin data (if TB available)
  - Summary table to console

Usage:
    profiler = TrainingProfiler(output_dir, n_warmup=3, n_active=5)
    for batch_idx, batch in enumerate(loader):
        with profiler.step_context(batch_idx):
            # ... training step ...
        if profiler.is_done(batch_idx):
            break
    profiler.export()
"""

import logging
from pathlib import Path

import torch
from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler

LOGGER = logging.getLogger("caption_test.profiler")


class TrainingProfiler:
    """PyTorch profiler for first N training batches.

    Args:
        output_dir: Directory for profiler output.
        n_warmup: Number of warmup steps (not profiled).
        n_active: Number of steps to actively profile.
        n_repeat: Number of profiling cycles.
        record_shapes: Record tensor shapes.
        profile_memory: Track memory allocation.
        with_stack: Record Python call stacks.
    """

    def __init__(
        self,
        output_dir: str | Path,
        n_warmup: int = 3,
        n_active: int = 5,
        n_repeat: int = 1,
        record_shapes: bool = True,
        profile_memory: bool = True,
        with_stack: bool = False,
    ) -> None:
        self.output_dir = Path(output_dir) / "profiler"
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.n_warmup = n_warmup
        self.n_active = n_active
        self.n_repeat = n_repeat
        self.total_steps = (n_warmup + n_active) * n_repeat

        self._profiler = profile(
            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
            schedule=schedule(
                wait=0,
                warmup=n_warmup,
                active=n_active,
                repeat=n_repeat,
            ),
            on_trace_ready=tensorboard_trace_handler(str(self.output_dir)),
            record_shapes=record_shapes,
            profile_memory=profile_memory,
            with_stack=with_stack,
        )
        self._started = False

    def start(self) -> None:
        """Start the profiler."""
        self._profiler.__enter__()
        self._started = True
        LOGGER.info(
            "Profiler started: %d warmup + %d active steps, output: %s",
            self.n_warmup, self.n_active, self.output_dir,
        )

    def step(self) -> None:
        """Signal end of a profiling step."""
        if self._started:
            self._profiler.step()

    def is_done(self, batch_idx: int) -> bool:
        """Check if profiling is complete."""
        return batch_idx >= self.total_steps

    def export(self) -> None:
        """Export profiling results and print summary."""
        if not self._started:
            return

        self._profiler.__exit__(None, None, None)
        self._started = False

        # Print key averages summary.
        summary = self._profiler.key_averages().table(
            sort_by="cuda_time_total", row_limit=20,
        )
        LOGGER.info("Profiler summary (top 20 by CUDA time):\n%s", summary)

        # Export Chrome trace.
        trace_path = self.output_dir / "chrome_trace.json"
        self._profiler.export_chrome_trace(str(trace_path))
        LOGGER.info("Chrome trace exported: %s", trace_path)

        # Memory summary if available.
        if torch.cuda.is_available():
            mem_summary = torch.cuda.memory_summary(abbreviated=True)
            summary_path = self.output_dir / "memory_summary.txt"
            summary_path.write_text(mem_summary)
            LOGGER.info("CUDA memory summary: %s", summary_path)


def print_model_summary(model: torch.nn.Module, device: str = "cuda") -> str:
    """Print model summary using torchinfo (if available).

    Falls back to a simple parameter count if torchinfo is not installed.

    Returns:
        Summary string.
    """
    try:
        from torchinfo import summary as torchinfo_summary
        info = torchinfo_summary(
            model,
            input_data={
                "drone_img": torch.randn(1, 3, 256, 256, device=device),
                "sat_img": torch.randn(1, 3, 256, 256, device=device),
            },
            col_names=["input_size", "output_size", "num_params", "trainable"],
            verbose=0,
            depth=3,
        )
        summary_str = str(info)
        LOGGER.info("Model summary (torchinfo):\n%s", summary_str)
        return summary_str
    except ImportError:
        LOGGER.info("torchinfo not installed, using basic parameter count")
    except Exception as e:
        LOGGER.warning("torchinfo failed (%s), using basic parameter count", e)

    # Fallback: simple param count.
    lines = []
    total = 0
    trainable = 0
    for name, param in model.named_parameters():
        total += param.numel()
        if param.requires_grad:
            trainable += param.numel()
            lines.append(f"  [trainable] {name}: {list(param.shape)} ({param.numel():,})")

    summary_str = (
        f"Total parameters: {total:,}\n"
        f"Trainable parameters: {trainable:,} ({100*trainable/max(total,1):.2f}%)\n"
        + "\n".join(lines[:30])
    )
    LOGGER.info("Model summary:\n%s", summary_str)
    return summary_str
