"""
Expert offloading — keep inactive experts on CPU, only load selected
ones to GPU on demand for real RAM savings.

Instead of holding all experts in GPU memory, this manager:
1. Stores all experts on CPU (or disk)
2. When the gate selects top-K, loads only those K to GPU
3. After computing, moves them back to CPU
4. Maintains a small cache to avoid redundant transfers
"""

from __future__ import annotations

import copy
from collections import OrderedDict
from typing import Any, Dict, List, Optional, Set, Tuple

import torch
import torch.nn as nn


class ExpertOffloader:
    """Manages expert device placement for memory efficiency.

    Keeps experts on CPU by default and moves only the active ones
    to the compute device (GPU) on demand.

    Args:
        device: The compute device to move active experts to (e.g., 'cuda').
        cache_size: Number of experts to keep on GPU in cache.
            Avoids repeated CPU→GPU transfers if the same expert is
            selected repeatedly. Set to 0 for minimal GPU RAM usage.
        verbose: Print offloading activity.
    """

    def __init__(
        self,
        device: str = "cuda",
        cache_size: int = 2,
        verbose: bool = False,
    ) -> None:
        self.device = torch.device(device) if isinstance(device, str) else device
        self.cache_size = cache_size
        self.verbose = verbose

        # LRU cache: expert_idx → True (most recently used at end)
        self._cache: OrderedDict[int, bool] = OrderedDict()

    def offload_all_to_cpu(self, experts: nn.ModuleList) -> None:
        """Move all experts to CPU for storage."""
        for i, expert in enumerate(experts):
            expert.to("cpu")
        self._cache.clear()
        if self.verbose:
            print(f"[Offloader] Moved all {len(experts)} experts to CPU")

    def load_experts(
        self,
        experts: nn.ModuleList,
        indices: List[int],
    ) -> None:
        """Load specific experts to the compute device.

        Only transfers experts not already in the GPU cache.

        Args:
            experts: The ModuleList of all experts.
            indices: Indices of experts to load.
        """
        for idx in indices:
            if idx in self._cache:
                # Already on GPU — just mark as recently used
                self._cache.move_to_end(idx)
            else:
                # Move to GPU
                experts[idx].to(self.device)
                self._cache[idx] = True
                self._cache.move_to_end(idx)

                if self.verbose:
                    print(f"[Offloader] Loaded expert {idx} to {self.device}")

                # Evict oldest if cache is full
                while len(self._cache) > self.cache_size and len(self._cache) > len(indices):
                    evict_idx, _ = self._cache.popitem(last=False)
                    # Don't evict if it's one we just loaded
                    if evict_idx in indices:
                        self._cache[evict_idx] = True
                        continue
                    experts[evict_idx].to("cpu")
                    if self.verbose:
                        print(f"[Offloader] Evicted expert {evict_idx} to CPU")

    def get_device(self) -> torch.device:
        """Get the compute device."""
        return self.device

    def get_cached_experts(self) -> List[int]:
        """Get indices of experts currently in GPU cache."""
        return list(self._cache.keys())

    def memory_stats(self, experts: nn.ModuleList) -> Dict[str, Any]:
        """Get memory usage statistics."""
        total_params = sum(
            sum(p.numel() for p in e.parameters()) for e in experts
        )
        cached_params = sum(
            sum(p.numel() for p in experts[i].parameters())
            for i in self._cache.keys()
            if i < len(experts)
        )

        bytes_per_param = 4  # fp32
        return {
            "total_experts": len(experts),
            "cached_experts": len(self._cache),
            "total_params": total_params,
            "cached_params": cached_params,
            "total_memory_mb": total_params * bytes_per_param / (1024 * 1024),
            "gpu_memory_mb": cached_params * bytes_per_param / (1024 * 1024),
            "ram_savings_pct": (
                (1 - cached_params / max(total_params, 1)) * 100
            ),
        }
