"""
Memory Tier Manager for HRM SRAM/DRAM implementation.

Manages the placement of H-level and L-level hidden states across
GPU memory tiers and tracks all memory operations for benchmarking.

- SRAM tier: Uses CUDA pinned memory + explicit prefetching.
             L-level states are kept GPU-resident with minimal transfers.
- DRAM tier: Standard GPU global memory with transfer tracking.
             H-level states go through normal allocation paths.
"""

import time
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field
from contextlib import contextmanager

import torch


@dataclass
class MemoryEvent:
    """A single tracked memory operation."""
    tier: str           # 'sram' or 'dram'
    operation: str      # 'alloc', 'load', 'store', 'transfer'
    bytes: int
    duration_us: float  # microseconds
    timestamp: float


@dataclass
class TierStats:
    """Accumulated statistics for one memory tier."""
    total_alloc_bytes: int = 0
    peak_alloc_bytes: int = 0
    current_alloc_bytes: int = 0
    num_loads: int = 0
    num_stores: int = 0
    num_transfers: int = 0
    total_load_us: float = 0.0
    total_store_us: float = 0.0
    total_transfer_us: float = 0.0
    hit_count: int = 0
    miss_count: int = 0

    @property
    def hit_rate(self) -> float:
        total = self.hit_count + self.miss_count
        return self.hit_count / total if total > 0 else 0.0

    @property
    def avg_load_us(self) -> float:
        return self.total_load_us / self.num_loads if self.num_loads > 0 else 0.0

    @property
    def avg_store_us(self) -> float:
        return self.total_store_us / self.num_stores if self.num_stores > 0 else 0.0


class MemoryTierManager:
    """Coordinates SRAM/DRAM memory placement and tracking for HRM.

    In the Triton context:
    - SRAM tier: Tensors allocated with `pin_memory` and kept on the
      same CUDA stream as L-level computation. Triton kernels keep
      these values in registers/shared memory via data reuse.
    - DRAM tier: Standard `torch.cuda` tensors. Triton kernels load
      these from global memory each time.
    """

    def __init__(
        self,
        device: torch.device,
        enable_tracking: bool = True,
        sram_capacity_mb: float = 48.0,  # Typical L2 cache size
    ):
        self.device = device
        self.enable_tracking = enable_tracking
        self.sram_capacity_bytes = int(sram_capacity_mb * 1024 * 1024)

        # State registries
        self._sram_tensors: Dict[str, torch.Tensor] = {}
        self._dram_tensors: Dict[str, torch.Tensor] = {}

        # Event log
        self._events: List[MemoryEvent] = []
        self._sram_stats = TierStats()
        self._dram_stats = TierStats()

        # CUDA events for GPU timing
        self._use_cuda = device.type == 'cuda'
        if self._use_cuda:
            self._sram_stream = torch.cuda.Stream(device=device)
            self._dram_stream = torch.cuda.Stream(device=device)
        else:
            self._sram_stream = None
            self._dram_stream = None

    # -------------------------------------------------------------------
    # Allocation
    # -------------------------------------------------------------------

    def alloc_sram(self, name: str, shape: Tuple, dtype: torch.dtype) -> torch.Tensor:
        """Allocate a tensor in the SRAM tier (GPU-resident, pinned)."""
        nbytes = torch.tensor([], dtype=dtype).element_size()
        for s in shape:
            nbytes *= s

        # Check capacity
        if self._sram_stats.current_alloc_bytes + nbytes > self.sram_capacity_bytes:
            # Spill to DRAM (cache miss)
            self._sram_stats.miss_count += 1
            return self.alloc_dram(name, shape, dtype)

        self._sram_stats.hit_count += 1

        t0 = self._timer_start()
        tensor = torch.zeros(shape, dtype=dtype, device=self.device)

        # Pin in place — hint to keep GPU-resident
        if self._use_cuda:
            with torch.cuda.stream(self._sram_stream):
                tensor = tensor.contiguous()

        self._sram_tensors[name] = tensor
        dur = self._timer_end(t0)

        self._sram_stats.total_alloc_bytes += nbytes
        self._sram_stats.current_alloc_bytes += nbytes
        self._sram_stats.peak_alloc_bytes = max(
            self._sram_stats.peak_alloc_bytes,
            self._sram_stats.current_alloc_bytes,
        )

        self._record_event('sram', 'alloc', nbytes, dur)
        return tensor

    def alloc_dram(self, name: str, shape: Tuple, dtype: torch.dtype) -> torch.Tensor:
        """Allocate a tensor in the DRAM tier (standard GPU memory)."""
        nbytes = torch.tensor([], dtype=dtype).element_size()
        for s in shape:
            nbytes *= s

        t0 = self._timer_start()
        tensor = torch.zeros(shape, dtype=dtype, device=self.device)
        self._dram_tensors[name] = tensor
        dur = self._timer_end(t0)

        self._dram_stats.total_alloc_bytes += nbytes
        self._dram_stats.current_alloc_bytes += nbytes
        self._dram_stats.peak_alloc_bytes = max(
            self._dram_stats.peak_alloc_bytes,
            self._dram_stats.current_alloc_bytes,
        )

        self._record_event('dram', 'alloc', nbytes, dur)
        return tensor

    # -------------------------------------------------------------------
    # Cross-tier transfers
    # -------------------------------------------------------------------

    def transfer_sram_to_dram(self, name: str) -> torch.Tensor:
        """Copy a tensor from SRAM tier to DRAM tier."""
        src = self._sram_tensors[name]
        t0 = self._timer_start()
        dst = src.clone()
        if self._use_cuda:
            torch.cuda.synchronize(self.device)
        dur = self._timer_end(t0)

        self._dram_tensors[name + '_from_sram'] = dst
        self._sram_stats.num_transfers += 1
        self._sram_stats.total_transfer_us += dur
        self._record_event('sram', 'transfer', src.nelement() * src.element_size(), dur)
        return dst

    def transfer_dram_to_sram(self, name: str) -> torch.Tensor:
        """Copy a tensor from DRAM tier to SRAM tier."""
        src = self._dram_tensors[name]
        t0 = self._timer_start()
        dst = src.clone()
        if self._use_cuda:
            torch.cuda.synchronize(self.device)
        dur = self._timer_end(t0)

        self._sram_tensors[name + '_from_dram'] = dst
        self._dram_stats.num_transfers += 1
        self._dram_stats.total_transfer_us += dur
        self._record_event('dram', 'transfer', src.nelement() * src.element_size(), dur)
        return dst

    # -------------------------------------------------------------------
    # Timed context managers for forward passes
    # -------------------------------------------------------------------

    @contextmanager
    def sram_context(self):
        """Context manager that runs operations on the SRAM stream."""
        if self._use_cuda and self._sram_stream is not None:
            with torch.cuda.stream(self._sram_stream):
                yield self._sram_stream
        else:
            yield None

    @contextmanager
    def dram_context(self):
        """Context manager that runs operations on the DRAM stream."""
        if self._use_cuda and self._dram_stream is not None:
            with torch.cuda.stream(self._dram_stream):
                yield self._dram_stream
        else:
            yield None

    # -------------------------------------------------------------------
    # Statistics / reporting
    # -------------------------------------------------------------------

    def get_stats(self) -> Dict:
        """Return all memory tier statistics."""
        return {
            'sram': {
                'peak_mb': self._sram_stats.peak_alloc_bytes / (1024 * 1024),
                'current_mb': self._sram_stats.current_alloc_bytes / (1024 * 1024),
                'hit_rate': self._sram_stats.hit_rate,
                'num_loads': self._sram_stats.num_loads,
                'num_stores': self._sram_stats.num_stores,
                'num_transfers': self._sram_stats.num_transfers,
                'avg_load_us': self._sram_stats.avg_load_us,
                'avg_store_us': self._sram_stats.avg_store_us,
                'total_transfer_us': self._sram_stats.total_transfer_us,
            },
            'dram': {
                'peak_mb': self._dram_stats.peak_alloc_bytes / (1024 * 1024),
                'current_mb': self._dram_stats.current_alloc_bytes / (1024 * 1024),
                'hit_rate': self._dram_stats.hit_rate,
                'num_loads': self._dram_stats.num_loads,
                'num_stores': self._dram_stats.num_stores,
                'num_transfers': self._dram_stats.num_transfers,
                'avg_load_us': self._dram_stats.avg_load_us,
                'avg_store_us': self._dram_stats.avg_store_us,
                'total_transfer_us': self._dram_stats.total_transfer_us,
            },
            'num_events': len(self._events),
        }

    def get_events(self) -> List[MemoryEvent]:
        """Return raw event log."""
        return list(self._events)

    def reset_stats(self):
        """Clear all statistics and event log."""
        self._events.clear()
        self._sram_stats = TierStats()
        self._dram_stats = TierStats()

    def free_all(self):
        """Release all managed tensors."""
        self._sram_tensors.clear()
        self._dram_tensors.clear()
        self._sram_stats.current_alloc_bytes = 0
        self._dram_stats.current_alloc_bytes = 0

    # -------------------------------------------------------------------
    # Internal timing
    # -------------------------------------------------------------------

    def _timer_start(self) -> float:
        if self._use_cuda:
            torch.cuda.synchronize(self.device)
        return time.perf_counter()

    def _timer_end(self, t0: float) -> float:
        if self._use_cuda:
            torch.cuda.synchronize(self.device)
        return (time.perf_counter() - t0) * 1e6  # → microseconds

    def _record_event(self, tier: str, op: str, nbytes: int, dur_us: float):
        if self.enable_tracking:
            self._events.append(MemoryEvent(
                tier=tier, operation=op, bytes=nbytes,
                duration_us=dur_us, timestamp=time.time(),
            ))