"""
Utility functions for the ARES Context Compression Layer.
"""

from typing import Dict, Any
import torch


def calculate_memory_size(tensor: torch.Tensor) -> float:
    """Calculate memory size of tensor in MB."""
    return tensor.element_size() * tensor.numel() / (1024 * 1024)


def format_compression_report(result: Any) -> str:
    """Format compression result as readable report."""
    report = f"""
=== ARES Context Compression Report ===
Original Length:     {result.original_length:,} tokens
Compressed Length:   {result.compressed_length:,} tokens
Compression Ratio:   {result.compression_ratio:.1%}
Memory Saved:        {result.memory_saved:.2f} MB
Latency:             {result.latency_ms:.2f} ms

=== Metadata ===
"""
    for layer, meta in result.metadata.items():
        report += f"\n{layer.upper()}:\n"
        for key, value in meta.items():
            if isinstance(value, float):
                report += f"  {key}: {value:.4f}\n"
            else:
                report += f"  {key}: {value}\n"

    return report


def create_sample_context(
    seq_len: int = 2048,
    hidden_dim: int = 768,
    device: str | None = None
) -> torch.Tensor:
    """Create sample context tensor for testing."""
    resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    return torch.randn(seq_len, hidden_dim, device=resolved_device)


def estimate_compression_benefit(
    current_vram_mb: float,
    target_reduction: float = 0.5
) -> Dict[str, float]:
    """Estimate memory savings from compression."""
    saved = current_vram_mb * target_reduction
    remaining = current_vram_mb - saved
    return {
        "current_vram_mb": current_vram_mb,
        "estimated_saved_mb": saved,
        "estimated_remaining_mb": remaining,
        "target_reduction_pct": target_reduction * 100
    }


class CompressionMonitor:
    """Monitor compression performance over time."""

    def __init__(self, window_size: int = 100):
        self.window_size = window_size
        self.metrics = {
            "compression_ratios": [],
            "latencies_ms": [],
            "memory_savings_mb": []
        }

    def update(self, result: Any) -> None:
        """Update metrics with new result."""
        self.metrics["compression_ratios"].append(result.compression_ratio)
        self.metrics["latencies_ms"].append(result.latency_ms)
        self.metrics["memory_savings_mb"].append(result.memory_saved)

        # Keep only recent metrics
        if len(self.metrics["compression_ratios"]) > self.window_size:
            self.metrics["compression_ratios"].pop(0)
            self.metrics["latencies_ms"].pop(0)
            self.metrics["memory_savings_mb"].pop(0)

    def get_summary(self) -> Dict[str, float]:
        """Get summary statistics."""
        if not self.metrics["compression_ratios"]:
            return {}

        return {
            "avg_compression_ratio": sum(self.metrics["compression_ratios"]) / len(self.metrics["compression_ratios"]),
            "avg_latency_ms": sum(self.metrics["latencies_ms"]) / len(self.metrics["latencies_ms"]),
            "avg_memory_saved_mb": sum(self.metrics["memory_savings_mb"]) / len(self.metrics["memory_savings_mb"]),
            "total_compressions": len(self.metrics["compression_ratios"])
        }
