#!/usr/bin/env python3
"""
ARES Context Compression Layer - Standalone Plug-and-Play Script

This is a single-file version that can be dropped into any orchestration stack.
No installation required - just copy this file and import it.

Usage:
    from compression_layer_standalone import compress_context
    
    # Your context tensor [seq_len, hidden_dim] or [batch, seq_len, hidden_dim]
    result = compress_context(context_tensor)
    
    # Use compressed context
    print(f"Compressed {result['original_length']} -> {result['compressed_length']} tokens")
    print(f"Memory saved: {result['memory_saved_mb']:.2f} MB")
"""

import torch
import torch.nn.functional as F
from typing import Optional, Dict, Any, Literal
from dataclasses import dataclass
import time


# =============================================================================
# CONFIGURATION
# =============================================================================

CompressionMode = Literal["conservative", "balanced", "aggressive"]

DEFAULT_CONFIGS = {
    "conservative": {
        "target_ratio": 0.6,
        "entropy_threshold": 0.5,
        "use_precision_switching": False,
    },
    "balanced": {
        "target_ratio": 0.5,
        "entropy_threshold": 1.0,
        "use_precision_switching": True,
    },
    "aggressive": {
        "target_ratio": 0.3,
        "entropy_threshold": 1.5,
        "use_precision_switching": True,
    },
}


# =============================================================================
# COMPRESSION RESULT
# =============================================================================

@dataclass
class CompressionResult:
    compressed_context: torch.Tensor
    attention_mask: Optional[torch.Tensor]
    original_length: int
    compressed_length: int
    compression_ratio: float
    memory_saved_mb: float
    latency_ms: float
    metadata: Dict[str, Any]
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "compressed_context": self.compressed_context,
            "attention_mask": self.attention_mask,
            "original_length": self.original_length,
            "compressed_length": self.compressed_length,
            "compression_ratio": self.compression_ratio,
            "memory_saved_mb": self.memory_saved_mb,
            "latency_ms": self.latency_ms,
            "metadata": self.metadata,
        }


# =============================================================================
# COMPRESSION LAYERS
# =============================================================================

class SemanticTokenPruning:
    def __init__(self, target_ratio: float = 0.5, min_tokens: int = 64):
        self.target_ratio = target_ratio
        self.min_tokens = min_tokens
    
    def compute_importance(self, context: torch.Tensor) -> torch.Tensor:
        norms = torch.norm(context, dim=-1)
        if norms.max() > norms.min():
            norms = (norms - norms.min()) / (norms.max() - norms.min())
        return norms
    
    def prune(self, context: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        original_len = context.shape[0]
        target_len = max(int(original_len * self.target_ratio), self.min_tokens)
        
        importance = self.compute_importance(context)
        _, top_indices = torch.topk(importance, target_len, sorted=True)
        top_indices = torch.sort(top_indices).values
        
        pruned_context = context[top_indices]
        pruned_mask = None
        if attention_mask is not None:
            pruned_mask = attention_mask[top_indices]
        
        metadata = {
            "layer": "semantic_pruning",
            "original_length": original_len,
            "pruned_length": target_len,
            "kept_ratio": target_len / original_len,
        }
        
        return pruned_context, pruned_mask, metadata


class DynamicPrecisionSwitching:
    def __init__(self, enabled: bool = True):
        self.enabled = enabled
    
    def compute_entropy(self, context: torch.Tensor) -> torch.Tensor:
        variance = torch.var(context, dim=-1)
        if variance.max() > variance.min():
            variance = (variance - variance.min()) / (variance.max() - variance.min())
        return variance
    
    def switch_precision(self, context: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        if not self.enabled:
            return context, attention_mask, {"layer": "dynamic_precision", "enabled": False}
        
        entropy = self.compute_entropy(context)
        metadata = {
            "layer": "dynamic_precision",
            "enabled": True,
            "avg_entropy": entropy.mean().item(),
            "high_entropy_tokens": (entropy > 0.5).sum().item(),
        }
        
        return context, attention_mask, metadata


class EntropyGatedCompression:
    def __init__(self, entropy_threshold: float = 1.0):
        self.entropy_threshold = entropy_threshold
    
    def compress(self, context: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        entropy = torch.var(context, dim=-1)
        high_entropy_mask = entropy > self.entropy_threshold
        
        if high_entropy_mask.sum() == 0:
            keep_mask = torch.ones(context.shape[0], dtype=torch.bool, device=context.device)
        else:
            keep_mask = high_entropy_mask
        
        compressed_context = context[keep_mask]
        compressed_mask = None
        if attention_mask is not None:
            compressed_mask = attention_mask[keep_mask]
        
        metadata = {
            "layer": "entropy_gated",
            "threshold": self.entropy_threshold,
            "kept_tokens": keep_mask.sum().item(),
            "filtered_tokens": (~keep_mask).sum().item(),
        }
        
        return compressed_context, compressed_mask, metadata


# =============================================================================
# MAIN COMPRESSOR
# =============================================================================

class ContextCompressor:
    def __init__(self, mode: CompressionMode = "balanced", device: Optional[str] = None, **kwargs):
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        
        config = DEFAULT_CONFIGS.get(mode, DEFAULT_CONFIGS["balanced"]).copy()
        config.update(kwargs)
        
        self.mode = mode
        self.config = config
        
        self.pruning_layer = SemanticTokenPruning(
            target_ratio=config["target_ratio"],
            min_tokens=64
        )
        self.precision_layer = DynamicPrecisionSwitching(
            enabled=config.get("use_precision_switching", False)
        )
        self.entropy_layer = EntropyGatedCompression(
            entropy_threshold=config.get("entropy_threshold", 1.0)
        )
        
        self.stats = {
            "total_compressions": 0,
            "total_tokens_processed": 0,
            "total_memory_saved_mb": 0.0,
            "total_latency_ms": 0.0,
        }
    
    def compress(self, context: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> CompressionResult:
        start_time = time.time()
        
        had_batch = context.dim() == 3
        if had_batch:
            context = context[0]
            if attention_mask is not None:
                attention_mask = attention_mask[0]
        
        context = context.to(self.device)
        if attention_mask is not None:
            attention_mask = attention_mask.to(self.device)
        
        original_length = context.shape[0]
        original_memory_mb = context.element_size() * context.numel() / (1024 * 1024)
        
        current_context = context
        current_mask = attention_mask
        all_metadata = {}
        
        current_context, current_mask, metadata = self.pruning_layer.prune(current_context, current_mask)
        all_metadata["semantic_pruning"] = metadata
        
        current_context, current_mask, metadata = self.entropy_layer.compress(current_context, current_mask)
        all_metadata["entropy_gated"] = metadata
        
        current_context, current_mask, metadata = self.precision_layer.switch_precision(current_context, current_mask)
        all_metadata["dynamic_precision"] = metadata
        
        if had_batch:
            current_context = current_context.unsqueeze(0)
            if current_mask is not None:
                current_mask = current_mask.unsqueeze(0)
        
        compressed_length = current_context.shape[-2] if had_batch else current_context.shape[0]
        compressed_memory_mb = current_context.element_size() * current_context.numel() / (1024 * 1024)
        memory_saved_mb = original_memory_mb - compressed_memory_mb
        compression_ratio = 1.0 - (compressed_length / original_length)
        latency_ms = (time.time() - start_time) * 1000
        
        self.stats["total_compressions"] += 1
        self.stats["total_tokens_processed"] += original_length
        self.stats["total_memory_saved_mb"] += memory_saved_mb
        self.stats["total_latency_ms"] += latency_ms
        
        return CompressionResult(
            compressed_context=current_context,
            attention_mask=current_mask,
            original_length=original_length,
            compressed_length=compressed_length,
            compression_ratio=compression_ratio,
            memory_saved_mb=memory_saved_mb,
            latency_ms=latency_ms,
            metadata=all_metadata,
        )
    
    def get_stats(self) -> Dict[str, Any]:
        if self.stats["total_compressions"] == 0:
            return self.stats
        
        return {
            **self.stats,
            "avg_compression_ratio": self.stats["total_memory_saved_mb"] / max(1, self.stats["total_tokens_processed"]) * 100,
            "avg_latency_ms": self.stats["total_latency_ms"] / self.stats["total_compressions"],
        }


def compress_context(context: torch.Tensor, mode: CompressionMode = "balanced", device: Optional[str] = None, **kwargs) -> CompressionResult:
    compressor = ContextCompressor(mode=mode, device=device, **kwargs)
    return compressor.compress(context)


def create_sample_context(seq_len: int = 2048, hidden_dim: int = 768, device: str = "cpu") -> torch.Tensor:
    return torch.randn(seq_len, hidden_dim, device=device)


if __name__ == "__main__":
    print("=" * 70)
    print("ARES Context Compression Layer - Standalone Demo")
    print("=" * 70)
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"\nDevice: {device}")
    
    print("\n" + "-" * 70)
    print("Demo 1: Basic Usage")
    print("-" * 70)
    
    context = create_sample_context(2048, 768, device)
    result = compress_context(context, mode="balanced", device=device)
    
    print(f"Original tokens: {result.original_length}")
    print(f"Compressed tokens: {result.compressed_length}")
    print(f"Compression ratio: {result.compression_ratio:.1%}")
    print(f"Memory saved: {result.memory_saved_mb:.2f} MB")
    print(f"Latency: {result.latency_ms:.2f} ms")
    
    print("\n" + "-" * 70)
    print("Demo 2: Mode Comparison")
    print("-" * 70)
    
    modes = ["conservative", "balanced", "aggressive"]
    context = create_sample_context(4096, 768, device)
    
    for mode in modes:
        result = compress_context(context, mode=mode, device=device)
        print(f"\n{mode.upper()}:")
        print(f"  Tokens: {result.original_length} -> {result.compressed_length}")
        print(f"  Compression: {result.compression_ratio:.1%}")
        print(f"  Memory saved: {result.memory_saved_mb:.2f} MB")
        print(f"  Latency: {result.latency_ms:.2f} ms")
    
    print("\n" + "=" * 70)
    print("Demo complete! The compression layer is ready to use.")
    print("=" * 70)
