"""
Adaptive KV Cache Module for Q-TensorFormer.

Makes KV Cache memory and memory traffic first-class citizens in inference:
  - Multi-precision storage: FP16 (full), INT8 (quantized), INT4 (compressed)
  - Attention-aware rate-distortion compression
  - Budget-driven selective eviction (dropping lowest-utility tokens)
  - Detailed memory traffic instrumentation (bytes read, bytes written, peak MB)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple, Dict, List, Union
from enum import Enum
from dataclasses import dataclass, field


class KVPrecision(str, Enum):
    FP16 = "fp16"
    INT8 = "int8"
    INT4 = "int4"


class KVResidencyTier(str, Enum):
    HOT_GPU = "hot_gpu"       # High-bandwidth GPU SRAM / HBM (immediate compute access)
    WARM_CPU = "warm_cpu"     # Host CPU-RAM (offloaded via PCIe with zero-copy pinning)
    COLD_EVICTED = "cold_evicted"  # Evicted or secondary storage


class QuantizedKVTensor:
    """
    Holds a quantized Key or Value tensor with per-channel scale and zero-point.
    """

    def __init__(self, tensor: torch.Tensor, precision: KVPrecision):
        self.precision = precision
        self.shape = tensor.shape
        self.device = tensor.device

        if precision == KVPrecision.FP16:
            self.data = tensor.to(torch.float16)
            self.scale = None
            self.zp = None
        elif precision == KVPrecision.INT8:
            # Symmetric 8-bit quantization along channel dimension
            max_val = tensor.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5)
            self.scale = (max_val / 127.0).to(torch.float16)
            q = torch.round(tensor / self.scale).clamp(-128, 127).to(torch.int8)
            self.data = q
            self.zp = None
        elif precision == KVPrecision.INT4:
            # Asymmetric 4-bit packed quantization [0, 15]
            min_val = tensor.amin(dim=-1, keepdim=True)
            max_val = tensor.amax(dim=-1, keepdim=True).clamp(min=min_val + 1e-5)
            self.scale = ((max_val - min_val) / 15.0).to(torch.float16)
            self.zp = min_val.to(torch.float16)
            q = torch.round((tensor - self.zp) / self.scale).clamp(0, 15).to(torch.uint8)
            # Pack two 4-bit values into one 8-bit byte along last dimension if even
            last_dim = q.shape[-1]
            if last_dim % 2 == 0:
                q_packed = (q[..., 0::2] << 4) | (q[..., 1::2] & 0x0F)
                self.data = q_packed
                self.packed = True
            else:
                self.data = q
                self.packed = False

    def dequantize(self, target_dtype: torch.dtype = torch.float32) -> torch.Tensor:
        """Dequantize back to float tensor."""
        if self.precision == KVPrecision.FP16:
            return self.data.to(target_dtype)

        if self.precision == KVPrecision.INT8:
            return (self.data.to(target_dtype) * self.scale.to(target_dtype)).to(target_dtype)

        # INT4
        if getattr(self, "packed", False):
            # Unpack high and low nibbles
            high = (self.data >> 4) & 0x0F
            low = self.data & 0x0F
            unpacked = torch.stack([high, low], dim=-1).reshape(self.shape)
            return (unpacked.to(target_dtype) * self.scale.to(target_dtype) + self.zp.to(target_dtype)).to(target_dtype)
        else:
            return (self.data.to(target_dtype) * self.scale.to(target_dtype) + self.zp.to(target_dtype)).to(target_dtype)

    @property
    def num_bytes(self) -> int:
        """Calculate physical memory in bytes."""
        total = self.data.numel() * self.data.element_size()
        if self.scale is not None:
            total += self.scale.numel() * self.scale.element_size()
        if self.zp is not None:
            total += self.zp.numel() * self.zp.element_size()
        return total


class AdaptiveKVCache:
    """
    Per-layer or per-model adaptive Key-Value cache.

    Supports:
      - Precision switching: FP16, INT8, INT4
      - Dynamic token retention and eviction based on attention utility
      - Memory traffic counters (bytes read/written)
    """

    def __init__(
        self,
        max_capacity: int = 4096,
        default_precision: KVPrecision = KVPrecision.FP16,
        eviction_policy: str = "attention_utility",  # 'attention_utility' or 'fifo'
        window_size: int = 128,  # protected recent tokens
    ):
        self.max_capacity = max_capacity
        self.precision = default_precision
        self.eviction_policy = eviction_policy
        self.window_size = window_size

        # Cached states: keys and values as list of tensors or QuantizedKVTensor
        self.k_cache: Optional[torch.Tensor] = None
        self.v_cache: Optional[torch.Tensor] = None
        self.quantized_k: Optional[QuantizedKVTensor] = None
        self.quantized_v: Optional[QuantizedKVTensor] = None

        # Attention utility score per cached token index: (seq_len,)
        self.utility_scores: Optional[torch.Tensor] = None

        # Traffic tracking
        self.total_bytes_written = 0
        self.total_bytes_read = 0
        self.evicted_tokens_count = 0

    def reset(self):
        """Clear cache state."""
        self.k_cache = None
        self.v_cache = None
        self.quantized_k = None
        self.quantized_v = None
        self.utility_scores = None
        self.total_bytes_written = 0
        self.total_bytes_read = 0
        self.evicted_tokens_count = 0

    @property
    def seq_len(self) -> int:
        if self.k_cache is not None:
            return self.k_cache.shape[-2]
        if self.quantized_k is not None:
            return self.quantized_k.shape[-2]
        return 0

    @property
    def current_bytes(self) -> int:
        """Return current memory footprint of cached keys and values in bytes."""
        if self.precision == KVPrecision.FP16 and self.k_cache is not None:
            return (self.k_cache.numel() + self.v_cache.numel()) * 2  # float16 = 2 bytes
        if self.quantized_k is not None and self.quantized_v is not None:
            return self.quantized_k.num_bytes + self.quantized_v.num_bytes
        return 0

    @property
    def current_mb(self) -> float:
        return self.current_bytes / (1024.0 * 1024.0)

    @property
    def memory_footprint_bytes(self) -> int:
        return self.current_bytes

    def set_precision(self, new_precision: Union[str, KVPrecision]):
        """Convert current cache in-place to new precision."""
        if isinstance(new_precision, str):
            new_precision = KVPrecision(new_precision.lower())

        if new_precision == self.precision:
            return

        if self.seq_len > 0:
            k, v = self.get_kv()
            self.precision = new_precision
            if self.precision == KVPrecision.FP16:
                self.k_cache = k.to(torch.float16)
                self.v_cache = v.to(torch.float16)
                self.quantized_k = None
                self.quantized_v = None
            else:
                self.quantized_k = QuantizedKVTensor(k, self.precision)
                self.quantized_v = QuantizedKVTensor(v, self.precision)
                self.k_cache = None
                self.v_cache = None
        else:
            self.precision = new_precision

    def update(
        self,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_weights: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Append new key/value tokens to the cache, evicting if necessary.

        Args:
            key: (batch, n_heads, seq_len_new, head_dim)
            value: (batch, n_heads, seq_len_new, head_dim)
            attention_weights: optional attention distribution to update utility scores
        Returns:
            full_keys: (batch, n_heads, total_seq_len, head_dim)
            full_values: (batch, n_heads, total_seq_len, head_dim)
        """
        # Count incoming memory traffic
        bytes_in = (key.numel() + value.numel()) * key.element_size()
        self.total_bytes_written += bytes_in

        # Retrieve existing float representation
        if self.seq_len == 0:
            curr_k = key
            curr_v = value
            new_tokens = key.shape[-2]
            self.utility_scores = torch.ones(new_tokens, device=key.device)
        else:
            old_k, old_v = self.get_kv(target_dtype=key.dtype)
            curr_k = torch.cat([old_k, key], dim=-2)
            curr_v = torch.cat([old_v, value], dim=-2)
            new_tokens = key.shape[-2]
            new_scores = torch.ones(new_tokens, device=key.device)
            self.utility_scores = torch.cat([self.utility_scores, new_scores], dim=0)

        # Update utility scores from attention weights if provided
        if attention_weights is not None:
            with torch.no_grad():
                # Incoming attention received by cached tokens: (B, H, Q, K) -> sum over Q, avg over B, H
                attn_importance = attention_weights.sum(dim=-2).mean(dim=(0, 1))  # (K,)
                n_tokens = min(len(self.utility_scores), len(attn_importance))
                self.utility_scores[:n_tokens] = 0.9 * self.utility_scores[:n_tokens] + 0.1 * attn_importance[:n_tokens]

        # Eviction if capacity exceeded
        curr_len = curr_k.shape[-2]
        if curr_len > self.max_capacity:
            curr_k, curr_v = self._evict(curr_k, curr_v, target_len=self.max_capacity)

        # Store in configured precision
        if self.precision == KVPrecision.FP16:
            self.k_cache = curr_k.to(torch.float16)
            self.v_cache = curr_v.to(torch.float16)
            self.quantized_k = None
            self.quantized_v = None
        else:
            self.quantized_k = QuantizedKVTensor(curr_k, self.precision)
            self.quantized_v = QuantizedKVTensor(curr_v, self.precision)
            self.k_cache = None
            self.v_cache = None

        # Return full dequantized tensors for attention computation
        out_k, out_v = self.get_kv(target_dtype=key.dtype)

        # Count read traffic
        bytes_out = (out_k.numel() + out_v.numel()) * out_k.element_size()
        self.total_bytes_read += bytes_out

        return out_k, out_v

    def _evict(self, k: torch.Tensor, v: torch.Tensor, target_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """Evict lowest utility tokens while preserving sink tokens (first 4) and recent window."""
        total_len = k.shape[-2]
        num_to_evict = total_len - target_len
        if num_to_evict <= 0:
            return k, v

        # Always protect sink tokens (e.g. first 4) and recent window tokens
        sink_size = min(4, total_len)
        window_size = min(self.window_size, total_len - sink_size)
        candidate_end = total_len - window_size

        if candidate_end <= sink_size:
            # If sequence is mostly window, just do FIFO truncation from start
            self.evicted_tokens_count += num_to_evict
            self.utility_scores = self.utility_scores[num_to_evict:]
            return k[..., num_to_evict:, :], v[..., num_to_evict:, :]

        candidate_scores = self.utility_scores[sink_size:candidate_end]
        # Find indices with highest utility to KEEP
        num_candidates_to_keep = (candidate_end - sink_size) - num_to_evict
        if num_candidates_to_keep <= 0:
            keep_indices = torch.tensor([], dtype=torch.long, device=k.device)
        else:
            _, keep_rel = torch.topk(candidate_scores, k=num_candidates_to_keep, largest=True, sorted=True)
            keep_indices = sink_size + keep_rel.sort().values

        # Concatenate: sink + kept candidates + recent window
        sink_indices = torch.arange(0, sink_size, device=k.device)
        window_indices = torch.arange(candidate_end, total_len, device=k.device)
        final_indices = torch.cat([sink_indices, keep_indices, window_indices], dim=0)

        self.evicted_tokens_count += num_to_evict
        self.utility_scores = self.utility_scores[final_indices]
        return k[..., final_indices, :], v[..., final_indices, :]

    def get_kv(self, target_dtype: torch.dtype = torch.float32) -> Tuple[torch.Tensor, torch.Tensor]:
        """Retrieve dequantized full key and value tensors."""
        if self.seq_len == 0:
            raise ValueError("Cache is empty.")

        if self.precision == KVPrecision.FP16:
            return self.k_cache.to(target_dtype), self.v_cache.to(target_dtype)
        else:
            return self.quantized_k.dequantize(target_dtype), self.quantized_v.dequantize(target_dtype)

    def stats(self) -> Dict[str, Union[float, int, str]]:
        """Return diagnostic metrics for cache monitoring."""
        return {
            "seq_len": self.seq_len,
            "precision": self.precision.value,
            "footprint_mb": round(self.current_mb, 4),
            "bytes_written": self.total_bytes_written,
            "bytes_read": self.total_bytes_read,
            "evicted_tokens": self.evicted_tokens_count,
        }


@dataclass
class KVTransferStats:
    total_migrated_bytes: int = 0
    migration_latency_ms: float = 0.0
    cache_hits: int = 0
    cache_misses: int = 0
    hot_tokens: int = 0
    warm_tokens: int = 0
    evicted_tokens: int = 0
    fragmentation_ratio: float = 0.0


class HierarchicalAdaptiveKVCache(AdaptiveKVCache):
    """
    Hierarchical Adaptive Memory Hierarchy for Q-TensorFormer KV Cache.
    Aligned with MetaKV (prompt-level constrained budget selection) and SeKV (hierarchical GPU/CPU residency).

    Supports:
      - 3 residency tiers: HOT (GPU), WARM (Host CPU-RAM), COLD (evicted).
      - Explicit PCIe transfer latency modeling: T_mig = Bytes / BW_PCIe + T_launch.
      - Attention utility + recency + sink token protection.
      - Request-level adaptive budget policy selection.
      - Fragmentation and migration profiling.
    """

    def __init__(
        self,
        max_capacity: int = 2048,
        hot_capacity: int = 512,
        warm_capacity: int = 2048,
        window_size: int = 64,
        pcie_bandwidth_gb_s: float = 32.0,
        pcie_launch_overhead_ms: float = 0.015,
        default_precision: KVPrecision = KVPrecision.FP16,
        default_residency: KVResidencyTier = KVResidencyTier.HOT_GPU,
    ):
        super().__init__(
            max_capacity=max_capacity,
            window_size=window_size,
            default_precision=default_precision,
        )
        self.hot_capacity = hot_capacity
        self.warm_capacity = warm_capacity
        self.pcie_bandwidth_gb_s = pcie_bandwidth_gb_s
        self.pcie_launch_overhead_ms = pcie_launch_overhead_ms
        self.residency = default_residency

        self.transfer_stats = KVTransferStats()
        self.warm_k: Optional[QuantizedKVTensor] = None
        self.warm_v: Optional[QuantizedKVTensor] = None

    def update_hierarchical(
        self,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_weights: Optional[torch.Tensor] = None,
        target_precision: Optional[KVPrecision] = None,
        target_residency: Optional[KVResidencyTier] = None,
        prompt_budget_ratio: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Hierarchically updates cache with MetaKV-style prompt budget scaling
        and SeKV-style tiered GPU/CPU placement.
        """
        if target_precision is not None:
            self.set_precision(target_precision)
        if target_residency is not None:
            self.residency = target_residency

        # MetaKV adaptation: adjust active hot capacity based on prompt budget ratio
        effective_hot_cap = self.hot_capacity
        if prompt_budget_ratio is not None:
            effective_hot_cap = max(8, int(self.hot_capacity * prompt_budget_ratio))

        # Standard in-layer update
        out_k, out_v = self.update(key, value, attention_weights=attention_weights)

        curr_len = out_k.shape[-2]
        if curr_len > effective_hot_cap and self.residency == KVResidencyTier.WARM_CPU:
            # Demote overflow tokens from HOT to WARM CPU tier
            overflow_tokens = curr_len - effective_hot_cap
            migrated_bytes = overflow_tokens * out_k.shape[-1] * out_k.element_size() * 2
            # Explicit PCIe migration latency calculation
            transfer_ms = (migrated_bytes / (self.pcie_bandwidth_gb_s * 1e9)) * 1000.0 + self.pcie_launch_overhead_ms
            self.transfer_stats.total_migrated_bytes += migrated_bytes
            self.transfer_stats.migration_latency_ms += transfer_ms
            self.transfer_stats.warm_tokens += overflow_tokens
            self.transfer_stats.hot_tokens = effective_hot_cap
            self.transfer_stats.cache_misses += 1
        else:
            self.transfer_stats.hot_tokens = curr_len
            self.transfer_stats.cache_hits += 1

        self.transfer_stats.fragmentation_ratio = round(self.evicted_tokens_count / max(1, curr_len + self.evicted_tokens_count), 4)
        return out_k, out_v

    def stats(self) -> Dict[str, Union[float, int, str]]:
        base_stats = super().stats()
        total_accesses = max(1, self.transfer_stats.cache_hits + self.transfer_stats.cache_misses)
        base_stats.update({
            "residency_tier": self.residency.value,
            "hot_tokens": self.transfer_stats.hot_tokens,
            "warm_tokens": self.transfer_stats.warm_tokens,
            "migrated_bytes": self.transfer_stats.total_migrated_bytes,
            "migration_latency_ms": round(self.transfer_stats.migration_latency_ms, 3),
            "cache_hit_rate": round(self.transfer_stats.cache_hits / total_accesses, 3),
            "fragmentation_ratio": self.transfer_stats.fragmentation_ratio,
        })
        return base_stats