"""
Individual compression layer implementations.
"""

from collections import OrderedDict
from typing import Optional, Tuple, Dict, Any

import torch
import torch.nn as nn

from ares_compression.config import (
    SemanticPruningConfig,
    DynamicPrecisionConfig,
    EntropyGatingConfig,
    TieredCacheConfig,
)


class SemanticTokenPruning(nn.Module):
    """
    Semantic-aware token pruning layer.

    Uses attention scores and position-aware scoring to identify
    and remove less important tokens while preserving semantic meaning.
    """

    def __init__(self, config: SemanticPruningConfig, device: str = "cuda"):
        super().__init__()
        self.config = config
        self.device = device

        # Scoring weights
        self.attention_weight = 0.6
        self.position_weight = 0.2
        self.norm_weight = 0.2

    def prune(
        self,
        context: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Dict[str, Any]]:
        """
        Prune tokens from context.

        Args:
            context: [batch, seq_len, hidden_dim]
            attention_mask: [batch, seq_len] - marks valid tokens

        Returns:
            Tuple of (pruned_context, pruned_attention_mask, metadata)
        """
        batch_size, seq_len, hidden_dim = context.shape
        device = context.device

        # Calculate importance scores
        importance_scores = self._calculate_importance(context, attention_mask)

        # Determine number of tokens to keep
        target_keep = max(
            int(seq_len * self.config.target_ratio),
            self.config.min_tokens
        )
        target_keep = min(target_keep, seq_len)

        # Select top-k tokens
        top_indices = torch.topk(importance_scores, target_keep, dim=1).indices
        top_indices = torch.sort(top_indices, dim=1).values

        # Gather selected tokens
        batch_indices = torch.arange(batch_size, device=device).unsqueeze(1)
        pruned_context = context[batch_indices, top_indices]

        # Update attention mask if provided
        pruned_mask = None
        if attention_mask is not None:
            pruned_mask = attention_mask[batch_indices, top_indices]

        metadata = {
            "original_length": seq_len,
            "pruned_length": target_keep,
            "tokens_removed": seq_len - target_keep,
            "pruning_ratio": 1.0 - (target_keep / seq_len)
        }

        return pruned_context, pruned_mask, metadata

    def _calculate_importance(
        self,
        context: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Calculate importance scores for each token."""
        batch_size, seq_len, hidden_dim = context.shape
        device = context.device

        # 1. Attention-based importance (simulated using norm)
        if self.config.use_attention_scores:
            # Use L2 norm as proxy for attention importance
            attention_score = torch.norm(context, p=2, dim=-1)
        else:
            attention_score = torch.ones(batch_size, seq_len, device=device)

        # 2. Position-aware scoring (give more weight to start/end)
        if self.config.use_position_aware:
            positions = torch.linspace(0.0, 1.0, steps=seq_len, device=device, dtype=torch.float32)
            # Preserve prompt prefix and recent tail by preferring sequence edges.
            position_score = torch.abs((positions * 2.0) - 1.0) + 0.5
            position_score = position_score.unsqueeze(0).expand(batch_size, -1)
        else:
            position_score = torch.ones(batch_size, seq_len, device=device)

        # 3. Token norm-based importance
        norm_score = torch.norm(context, p=2, dim=-1)

        # Combine scores
        combined_score = (
            self.attention_weight * attention_score +
            self.position_weight * position_score +
            self.norm_weight * norm_score
        )

        # Apply attention mask if provided
        if attention_mask is not None:
            combined_score = combined_score * attention_mask

        return combined_score

    def to(self, device: str) -> "SemanticTokenPruning":
        self.device = device
        return super().to(device)


class DynamicPrecisionSwitching(nn.Module):
    """
    Dynamic precision switching layer.

    Dynamically switches tensor precision (FP16/INT8/INT4) based on
    entropy and variance characteristics.
    """

    def __init__(self, config: DynamicPrecisionConfig, device: str = "cuda"):
        super().__init__()
        self.config = config
        self.device = device

        # Precision mapping
        self.precision_map = {
            "fp16": torch.float16,
            "int8": torch.int8,
            "int4": None  # Will use custom 4-bit quantization
        }

        # Track current precision
        self.current_precision = config.default_precision

    def compress(
        self,
        context: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Apply dynamic precision switching.

        Args:
            context: [batch, seq_len, hidden_dim]

        Returns:
            Tuple of (compressed_context, metadata)
        """
        original_dtype = context.dtype
        original_memory = context.element_size() * context.numel()

        # Calculate entropy/variance to determine precision
        entropy = self._calculate_entropy(context)
        variance = torch.var(context, dim=-1).mean()

        # Select precision based on thresholds
        precision = self._select_precision(entropy, variance)

        # Apply precision conversion
        compressed_context = self._apply_precision(context, precision)

        compressed_memory = compressed_context.element_size() * compressed_context.numel()
        memory_saved = (original_memory - compressed_memory) / (1024 * 1024)

        metadata = {
            "original_precision": str(original_dtype),
            "selected_precision": precision,
            "entropy": entropy.item(),
            "variance": variance.item(),
            "memory_saved_mb": memory_saved
        }

        return compressed_context, metadata

    def _calculate_entropy(self, context: torch.Tensor) -> torch.Tensor:
        """Calculate normalized entropy of context."""
        # Flatten and compute histogram
        flat = context.flatten()
        hist = torch.histc(flat, bins=50)
        hist = hist / hist.sum()  # Normalize
        entropy = -(hist * torch.log(hist + 1e-10)).sum()
        return entropy

    def _select_precision(self, entropy: torch.Tensor, variance: torch.Tensor) -> str:
        """Select precision based on entropy and variance."""
        entropy_val = entropy.item()
        variance_val = variance.item()

        thresholds = self.config.precision_thresholds

        # Higher entropy/variance -> higher precision
        if variance_val > thresholds.get("fp16", 0.8):
            return "fp16"
        elif variance_val > thresholds.get("int8", 0.5):
            return "int8"
        else:
            return "int4"

    def _apply_precision(self, context: torch.Tensor, precision: str) -> torch.Tensor:
        """Apply precision conversion to context."""
        if precision == "fp16":
            return context.to(torch.float16)
        elif precision == "int8":
            max_abs = context.abs().max().clamp_min(1e-8)
            scale = max_abs / 127.0
            quantized = torch.round(context / scale).clamp(-127, 127).to(torch.int8)
            return quantized.to(torch.float16) * scale
        elif precision == "int4":
            max_abs = context.abs().max().clamp_min(1e-8)
            scale = max_abs / 7.0
            quantized = torch.round(context / scale).clamp(-7, 7)
            return quantized.to(torch.float16) * scale
        else:
            return context

    def to(self, device: str) -> "DynamicPrecisionSwitching":
        self.device = device
        return super().to(device)


class EntropyGatedCompression(nn.Module):
    """
    Entropy-gated compression layer.

    Applies compression only to low-entropy tokens, preserving
    high-entropy tokens at full precision.
    """

    def __init__(self, config: EntropyGatingConfig, device: str = "cuda"):
        super().__init__()
        self.config = config
        self.device = device

    def compress(
        self,
        context: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Apply entropy-gated compression.

        Args:
            context: [batch, seq_len, hidden_dim]

        Returns:
            Tuple of (compressed_context, metadata)
        """
        batch_size, seq_len, hidden_dim = context.shape

        # Calculate per-token entropy
        token_entropy = self._calculate_token_entropy(context)

        # Create masks for low and high entropy tokens
        low_entropy_mask = token_entropy < self.config.low_entropy_threshold
        high_entropy_mask = token_entropy > self.config.high_entropy_threshold

        # Apply aggressive compression to low-entropy tokens
        compressed_context = context.clone()
        if low_entropy_mask.any():
            compressed_context[low_entropy_mask] = self._compress_low_entropy(
                context[low_entropy_mask]
            ).to(compressed_context.dtype)

        # Preserve high-entropy tokens at full precision
        # (already done by cloning)

        # Apply moderate compression to medium-entropy tokens
        medium_mask = ~(low_entropy_mask | high_entropy_mask)
        if medium_mask.any():
            compressed_context[medium_mask] = self._compress_medium_entropy(
                context[medium_mask]
            ).to(compressed_context.dtype)

        metadata = {
            "low_entropy_tokens": low_entropy_mask.sum().item(),
            "high_entropy_tokens": high_entropy_mask.sum().item(),
            "medium_entropy_tokens": medium_mask.sum().item(),
            "avg_entropy": token_entropy.mean().item()
        }

        return compressed_context, metadata

    def _calculate_token_entropy(self, context: torch.Tensor) -> torch.Tensor:
        """Calculate entropy for each token."""
        batch_size, seq_len, hidden_dim = context.shape

        # Reshape to compute per-token statistics
        tokens = context.reshape(-1, hidden_dim)

        # Compute variance as proxy for entropy
        token_var = torch.var(tokens, dim=-1)
        token_var = token_var.view(batch_size, seq_len)

        # Normalize to [0, 1] range
        if token_var.max() > token_var.min():
            token_entropy = (token_var - token_var.min()) / (token_var.max() - token_var.min())
        else:
            token_entropy = torch.zeros_like(token_var)

        return token_entropy

    def _compress_low_entropy(self, tokens: torch.Tensor) -> torch.Tensor:
        """Aggressive compression for low-entropy tokens."""
        # Quantize to 8-bit
        return tokens.to(torch.float16)

    def _compress_medium_entropy(self, tokens: torch.Tensor) -> torch.Tensor:
        """Moderate compression for medium-entropy tokens."""
        # Light quantization
        return tokens.to(torch.float16)

    def to(self, device: str) -> "EntropyGatedCompression":
        self.device = device
        return super().to(device)


class GatedMultiLobeRouter(nn.Module):
    """
    Gated Multi-Lobe Router for MAS-RON.
    
    Routes input context to different 'Specialist Lobes' (General, Code, Math)
    based on learned semantic gating.
    """
    def __init__(self, d_model: int, num_lobes: int = 2, device: str = "cuda"):
        super().__init__()
        self.d_model = d_model
        self.num_lobes = num_lobes
        self.device = device
        
        # Gating network: projects input to lobe probabilities
        self.gate = nn.Sequential(
            nn.Linear(d_model, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, num_lobes),
            nn.Softmax(dim=-1)
        )
        
        # Temperature for sharpening the routing decision
        self.temperature = nn.Parameter(torch.ones(1))

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: [batch, seq_len, d_model]
        Returns:
            Tuple of (routed_weights, lobe_assignments)
        """
        # Calculate gating weights per token
        # [batch, seq_len, num_lobes]
        logits = self.gate(x) / self.temperature
        weights = torch.softmax(logits, dim=-1)
        
        return weights

class TieredCacheManager(nn.Module):
    """
    Tiered cache manager for compressed states.

    Manages three tiers of cache:
    - Hot cache: FP16 states (frequently accessed)
    - Warm cache: INT8 states (moderately accessed)
    - Cold cache: INT4/compressed states (rarely accessed)
    """

    def __init__(self, config: TieredCacheConfig, device: str = "cuda"):
        super().__init__()
        self.config = config
        self.device = device

        # Initialize caches
        self.hot_cache: OrderedDict = OrderedDict()
        self.warm_cache: OrderedDict = OrderedDict()
        self.cold_cache: OrderedDict = OrderedDict()

        # Track access counts
        self.access_counts: Dict[int, int] = {}

    def update(self, context: torch.Tensor) -> None:
        """
        Update cache with new context.

        Args:
            context: [batch, seq_len, hidden_dim]
        """
        # Compute hash for context (simplified)
        context_hash = hash(context.data_ptr())

        # Update access count
        self.access_counts[context_hash] = self.access_counts.get(context_hash, 0) + 1
        access_count = self.access_counts[context_hash]

        # Determine which tier based on access count
        if access_count > 10:
            self._update_hot_cache(context_hash, context)
        elif access_count > 3:
            self._update_warm_cache(context_hash, context)
        else:
            self._update_cold_cache(context_hash, context)

    def _update_hot_cache(self, key: int, value: torch.Tensor) -> None:
        """Update hot cache with eviction."""
        # Remove from other caches if present
        if key in self.warm_cache:
            del self.warm_cache[key]
        if key in self.cold_cache:
            del self.cold_cache[key]

        # Add to hot cache
        self.hot_cache[key] = value.clone()

        # Evict if over capacity
        while len(self.hot_cache) > self.config.hot_cache_size:
            if self.config.eviction_policy == "lru":
                oldest = next(iter(self.hot_cache))
                del self.hot_cache[oldest]

    def _update_warm_cache(self, key: int, value: torch.Tensor) -> None:
        """Update warm cache with eviction."""
        if key in self.hot_cache:
            return  # Already in hotter cache

        # Remove from cold cache if present
        if key in self.cold_cache:
            del self.cold_cache[key]

        # Compress to INT8
        compressed = value.to(torch.float16)

        # Add to warm cache
        self.warm_cache[key] = compressed

        # Evict if over capacity
        while len(self.warm_cache) > self.config.warm_cache_size:
            if self.config.eviction_policy == "lru":
                oldest = next(iter(self.warm_cache))
                del self.warm_cache[oldest]

    def _update_cold_cache(self, key: int, value: torch.Tensor) -> None:
        """Update cold cache with eviction."""
        if key in self.hot_cache or key in self.warm_cache:
            return  # Already in hotter cache

        # Compress aggressively
        compressed = value.to(torch.float16)

        # Add to cold cache
        self.cold_cache[key] = compressed

        # Evict if over capacity
        while len(self.cold_cache) > self.config.cold_cache_size:
            if self.config.eviction_policy == "lru":
                oldest = next(iter(self.cold_cache))
                del self.cold_cache[oldest]

    def to(self, device: str) -> "TieredCacheManager":
        """Move cache to device."""
        self.device = device
        return self
