"""
KitsuneOptimizer - High-level API for optimized PyTorch training.

Provides a drop-in replacement for PyTorch training loops with
automatic optimization through:
- Graph capture and scheduling
- CUDA stream parallelism
- Memory pooling
- Data prefetching
- Kernel fusion (Week 5)
- Automatic Mixed Precision (Week 6)
"""

from __future__ import annotations

import time
from typing import Optional, Dict, Any, Iterator, Callable, Tuple, List
from dataclasses import dataclass, field
from contextlib import contextmanager
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from ..core import (
    DataflowScheduler,
    ComputationGraph,
    ExecutionPlan,
    ModelExecutor,
    ParallelForwardExecutor,
)
from ..cuda import StreamPool, get_stream_pool
from ..memory import (
    MemoryPool,
    get_memory_pool,
    CUDAPrefetcher,
    create_prefetched_loader,
    LifetimeAnalyzer,
)
from ..fusion import FusionEngine, FusionDetector, get_fusion_engine
from ..amp import AMPConfig, PrecisionMode, AMPOptimizer, autocast_context, KitsuneGradScaler
from ..profiler import Profiler, get_logger

logger = get_logger(__name__)


@dataclass
class OptimizationConfig:
    """Configuration for Kitsune optimizations."""
    # Stream parallelism
    num_streams: int = 4
    enable_streams: bool = True

    # Memory optimization
    enable_memory_pool: bool = True
    max_cached_bytes: int = 1 << 30  # 1GB

    # Data prefetching
    enable_prefetch: bool = True
    prefetch_factor: int = 2

    # Graph optimization
    enable_graph_capture: bool = True
    scheduler_type: str = "wavefront"

    # Kernel fusion (Week 5)
    enable_fusion: bool = True

    # AMP (Week 6)
    enable_amp: bool = True
    amp_precision: Optional[str] = None  # "fp16", "bf16", "auto"

    # Profiling
    enable_profiling: bool = False
    warmup_iterations: int = 3


@dataclass
class OptimizationStats:
    """Statistics from optimization."""
    total_time_ms: float = 0.0
    compute_time_ms: float = 0.0
    transfer_time_ms: float = 0.0
    iterations: int = 0
    cache_hit_rate: float = 0.0
    peak_memory_mb: float = 0.0
    speedup: float = 1.0


class KitsuneOptimizer:
    """
    High-level optimizer for PyTorch models.

    Automatically applies Kitsune optimizations to training and inference.

    Usage:
        optimizer = KitsuneOptimizer(model, config=OptimizationConfig())

        # Optimized training loop
        for epoch in range(epochs):
            for batch in optimizer.prefetch(dataloader):
                with optimizer.optimize():
                    output = model(batch)
                    loss = criterion(output, target)
                    loss.backward()
                    opt.step()

        # Get optimization stats
        print(optimizer.stats)
    """

    def __init__(
        self,
        model: nn.Module,
        config: Optional[OptimizationConfig] = None,
        device: Optional[torch.device] = None,
    ):
        """
        Initialize Kitsune optimizer.

        Args:
            model: PyTorch model to optimize
            config: Optimization configuration
            device: Target device
        """
        self.model = model
        self.config = config or OptimizationConfig()
        self._device = device or self._get_device(model)

        # Initialize components based on config
        self._stream_pool: Optional[StreamPool] = None
        self._memory_pool: Optional[MemoryPool] = None
        self._scheduler: Optional[DataflowScheduler] = None
        self._profiler: Optional[Profiler] = None
        self._graph: Optional[ComputationGraph] = None
        self._plan: Optional[ExecutionPlan] = None

        # Fusion and AMP (Week 5-6)
        self._fusion_engine: Optional[FusionEngine] = None
        self._amp_config: Optional[AMPConfig] = None
        self._grad_scaler: Optional[KitsuneGradScaler] = None

        # Statistics
        self.stats = OptimizationStats()
        self._baseline_time: Optional[float] = None

        # Initialize
        self._initialize()

        logger.info(f"KitsuneOptimizer initialized for {type(model).__name__}")

    def _get_device(self, model: nn.Module) -> torch.device:
        """Get device from model parameters."""
        try:
            return next(model.parameters()).device
        except StopIteration:
            return torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def _initialize(self) -> None:
        """Initialize optimization components."""
        if not torch.cuda.is_available():
            logger.warning("CUDA not available, optimizations will be limited")
            return

        # Stream pool
        if self.config.enable_streams:
            self._stream_pool = StreamPool(num_streams=self.config.num_streams)

        # Memory pool
        if self.config.enable_memory_pool:
            self._memory_pool = MemoryPool(
                device=self._device,
                max_cached_bytes=self.config.max_cached_bytes,
            )

        # Scheduler
        if self.config.enable_graph_capture:
            self._scheduler = DataflowScheduler(
                num_streams=self.config.num_streams,
                scheduler_type=self.config.scheduler_type,
            )

        # Fusion engine (Week 5)
        if self.config.enable_fusion:
            self._fusion_engine = get_fusion_engine()

        # AMP configuration (Week 6)
        if self.config.enable_amp:
            precision_mode = PrecisionMode.AUTO
            if self.config.amp_precision:
                precision_map = {
                    "fp16": PrecisionMode.FP16,
                    "bf16": PrecisionMode.BF16,
                    "fp32": PrecisionMode.FP32,
                    "auto": PrecisionMode.AUTO,
                }
                precision_mode = precision_map.get(
                    self.config.amp_precision.lower(),
                    PrecisionMode.AUTO
                )
            self._amp_config = AMPConfig(precision_mode=precision_mode)
            self._grad_scaler = KitsuneGradScaler(config=self._amp_config)

        # Profiler
        if self.config.enable_profiling:
            self._profiler = Profiler()

    def capture_graph(self, sample_input: torch.Tensor) -> ComputationGraph:
        """
        Capture computation graph from model.

        Args:
            sample_input: Sample input for graph capture

        Returns:
            Captured computation graph
        """
        if self._scheduler is None:
            raise RuntimeError("Graph capture not enabled in config")

        self._plan = self._scheduler.capture_and_schedule(self.model, sample_input)
        self._graph = self._scheduler.graph

        logger.info(f"Captured graph: {self._graph.num_tasks} tasks")
        return self._graph

    def prefetch(self, dataloader: DataLoader) -> Iterator:
        """
        Wrap DataLoader with prefetching.

        Args:
            dataloader: PyTorch DataLoader

        Returns:
            Prefetching iterator
        """
        if not self.config.enable_prefetch or not torch.cuda.is_available():
            return iter(dataloader)

        return create_prefetched_loader(
            dataloader,
            device=self._device,
            prefetch_factor=self.config.prefetch_factor,
        )

    @contextmanager
    def optimize(self, use_amp: bool = True):
        """
        Context manager for optimized execution.

        Args:
            use_amp: Whether to use AMP autocast (default True if enabled)

        Usage:
            with optimizer.optimize():
                output = model(input)
        """
        start_time = time.perf_counter()

        try:
            # Apply AMP autocast if enabled
            if use_amp and self._amp_config is not None:
                with autocast_context(config=self._amp_config):
                    yield
            else:
                yield
        finally:
            elapsed = (time.perf_counter() - start_time) * 1000
            self.stats.total_time_ms += elapsed
            self.stats.iterations += 1

    @property
    def grad_scaler(self) -> Optional[KitsuneGradScaler]:
        """Get the gradient scaler for AMP training."""
        return self._grad_scaler

    @property
    def amp_config(self) -> Optional[AMPConfig]:
        """Get the AMP configuration."""
        return self._amp_config

    @property
    def fusion_engine(self) -> Optional[FusionEngine]:
        """Get the fusion engine."""
        return self._fusion_engine

    def synchronize(self) -> None:
        """Synchronize all CUDA operations."""
        if torch.cuda.is_available():
            torch.cuda.synchronize()

        if self._stream_pool is not None:
            self._stream_pool.synchronize_all()

    def benchmark(
        self,
        input_fn: Callable[[], torch.Tensor],
        num_iterations: int = 100,
        warmup: int = 10,
    ) -> Dict[str, float]:
        """
        Benchmark model with and without optimizations.

        Args:
            input_fn: Function that returns input tensor
            num_iterations: Number of timed iterations
            warmup: Warmup iterations

        Returns:
            Benchmark results
        """
        if not torch.cuda.is_available():
            return {"error": "CUDA required for benchmarking"}

        # Warmup
        for _ in range(warmup):
            with torch.no_grad():
                _ = self.model(input_fn())
        torch.cuda.synchronize()

        # Baseline (no optimization)
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)

        start.record()
        for _ in range(num_iterations):
            with torch.no_grad():
                _ = self.model(input_fn())
        end.record()
        end.synchronize()

        baseline_time = start.elapsed_time(end) / num_iterations
        self._baseline_time = baseline_time

        # With optimization context
        start.record()
        for _ in range(num_iterations):
            with self.optimize():
                with torch.no_grad():
                    _ = self.model(input_fn())
        end.record()
        end.synchronize()

        optimized_time = start.elapsed_time(end) / num_iterations

        speedup = baseline_time / optimized_time if optimized_time > 0 else 1.0
        self.stats.speedup = speedup

        return {
            "baseline_ms": baseline_time,
            "optimized_ms": optimized_time,
            "speedup": speedup,
            "num_iterations": num_iterations,
        }

    def get_memory_stats(self) -> Dict[str, Any]:
        """Get memory pool statistics."""
        if self._memory_pool is None:
            return {}
        return self._memory_pool.get_stats()

    def summary(self) -> str:
        """Get optimization summary."""
        lines = [
            "=" * 50,
            "KitsuneOptimizer Summary",
            "=" * 50,
            f"Model: {type(self.model).__name__}",
            f"Device: {self._device}",
            "",
            "Configuration:",
            f"  Streams: {self.config.num_streams if self.config.enable_streams else 'disabled'}",
            f"  Memory pool: {'enabled' if self.config.enable_memory_pool else 'disabled'}",
            f"  Prefetching: {'enabled' if self.config.enable_prefetch else 'disabled'}",
            f"  Graph capture: {'enabled' if self.config.enable_graph_capture else 'disabled'}",
            f"  Fusion: {'enabled' if self.config.enable_fusion else 'disabled'}",
        ]

        # AMP info
        if self._amp_config is not None:
            lines.append(f"  AMP: {self._amp_config.precision_mode.name}")
        else:
            lines.append("  AMP: disabled")

        lines.extend([
            "",
            "Statistics:",
            f"  Total iterations: {self.stats.iterations}",
            f"  Total time: {self.stats.total_time_ms:.2f} ms",
        ])

        if self.stats.iterations > 0:
            avg_time = self.stats.total_time_ms / self.stats.iterations
            lines.append(f"  Avg time/iter: {avg_time:.3f} ms")

        if self.stats.speedup != 1.0:
            lines.append(f"  Speedup: {self.stats.speedup:.2f}x")

        if self._graph is not None:
            lines.extend([
                "",
                "Graph:",
                f"  Tasks: {self._graph.num_tasks}",
            ])

        if self._fusion_engine is not None:
            lines.extend([
                "",
                "Fusion:",
                f"  Backend: {self._fusion_engine.backend}",
            ])

        if self._memory_pool is not None:
            mem_stats = self._memory_pool.get_stats()
            lines.extend([
                "",
                "Memory Pool:",
                f"  Hit rate: {mem_stats['hit_rate']:.1%}",
                f"  Cached: {mem_stats['bytes_cached'] / 1e6:.2f} MB",
            ])

        if self._grad_scaler is not None:
            lines.extend([
                "",
                "Gradient Scaler:",
                f"  Current scale: {self._grad_scaler.scale:.0f}",
            ])

        lines.append("=" * 50)
        return "\n".join(lines)


def optimize_model(
    model: nn.Module,
    sample_input: torch.Tensor,
    **config_kwargs,
) -> KitsuneOptimizer:
    """
    Quick setup for model optimization.

    Args:
        model: Model to optimize
        sample_input: Sample input for graph capture
        **config_kwargs: Configuration options

    Returns:
        Configured KitsuneOptimizer
    """
    config = OptimizationConfig(**config_kwargs)
    optimizer = KitsuneOptimizer(model, config)

    if config.enable_graph_capture:
        optimizer.capture_graph(sample_input)

    return optimizer
