"""
Benchmarking Utilities for CyMeta

This script provides benchmarking utilities to compare:
- Model size (parameters)
- Memory usage
- Inference speed
- Compression ratios
"""

import torch
import time
import numpy as np
from typing import Dict, List
from cymeta.models.linear import CyMetaLinear
import torch.nn as nn


def benchmark_model_size(model: nn.Module) -> Dict[str, int]:
    """
    Benchmark model size.
    
    Args:
        model: Model to benchmark
    
    Returns:
        Dictionary with size metrics
    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    # Estimate size in bytes (assuming float32)
    size_bytes = total_params * 4
    
    return {
        "total_parameters": total_params,
        "trainable_parameters": trainable_params,
        "size_bytes": size_bytes,
        "size_mb": size_bytes / (1024 * 1024),
    }


def benchmark_inference_speed(
    model: nn.Module,
    input_shape: tuple,
    num_runs: int = 100,
    warmup_runs: int = 10,
    device: str = "cpu",
) -> Dict[str, float]:
    """
    Benchmark inference speed.
    
    Args:
        model: Model to benchmark
        input_shape: Shape of input tensor (batch_size, ...)
        num_runs: Number of inference runs
        warmup_runs: Number of warmup runs
        device: Device to run on
    
    Returns:
        Dictionary with timing metrics
    """
    model = model.to(device)
    model.eval()
    
    # Create dummy input
    dummy_input = torch.randn(*input_shape).to(device)
    
    # Warmup
    with torch.no_grad():
        for _ in range(warmup_runs):
            _ = model(dummy_input)
    
    # Synchronize if CUDA
    if device.startswith("cuda"):
        torch.cuda.synchronize()
    
    # Benchmark
    times = []
    with torch.no_grad():
        for _ in range(num_runs):
            if device.startswith("cuda"):
                torch.cuda.synchronize()
            
            start_time = time.time()
            _ = model(dummy_input)
            
            if device.startswith("cuda"):
                torch.cuda.synchronize()
            
            end_time = time.time()
            times.append(end_time - start_time)
    
    times = np.array(times)
    
    return {
        "mean_time_ms": np.mean(times) * 1000,
        "std_time_ms": np.std(times) * 1000,
        "min_time_ms": np.min(times) * 1000,
        "max_time_ms": np.max(times) * 1000,
        "throughput_samples_per_sec": 1.0 / np.mean(times),
    }


def compare_models(
    standard_model: nn.Module,
    cymeta_model: nn.Module,
    input_shape: tuple,
    device: str = "cpu",
) -> Dict:
    """
    Compare standard and CyMeta models.
    
    Args:
        standard_model: Standard (uncompressed) model
        cymeta_model: CyMeta compressed model
        input_shape: Input tensor shape
        device: Device to run on
    
    Returns:
        Dictionary with comparison metrics
    """
    print("Benchmarking models...")
    print(f"Device: {device}")
    print(f"Input shape: {input_shape}")
    print()
    
    # Model size
    print("1. Model Size Comparison")
    print("-" * 40)
    standard_size = benchmark_model_size(standard_model)
    cymeta_size = benchmark_model_size(cymeta_model)
    
    print(f"Standard Model:")
    print(f"  Parameters: {standard_size['total_parameters']:,}")
    print(f"  Size: {standard_size['size_mb']:.2f} MB")
    
    print(f"\nCyMeta Model:")
    print(f"  Parameters: {cymeta_size['total_parameters']:,}")
    print(f"  Size: {cymeta_size['size_mb']:.2f} MB")
    
    compression_ratio = standard_size['total_parameters'] / cymeta_size['total_parameters']
    size_reduction = (1 - cymeta_size['size_mb'] / standard_size['size_mb']) * 100
    
    print(f"\nCompression:")
    print(f"  Ratio: {compression_ratio:.2f}x")
    print(f"  Size Reduction: {size_reduction:.1f}%")
    print()
    
    # Inference speed
    print("2. Inference Speed Comparison")
    print("-" * 40)
    standard_speed = benchmark_inference_speed(standard_model, input_shape, device=device)
    cymeta_speed = benchmark_inference_speed(cymeta_model, input_shape, device=device)
    
    print(f"Standard Model:")
    print(f"  Mean time: {standard_speed['mean_time_ms']:.3f} ms")
    print(f"  Throughput: {standard_speed['throughput_samples_per_sec']:.1f} samples/sec")
    
    print(f"\nCyMeta Model:")
    print(f"  Mean time: {cymeta_speed['mean_time_ms']:.3f} ms")
    print(f"  Throughput: {cymeta_speed['throughput_samples_per_sec']:.1f} samples/sec")
    
    speed_ratio = standard_speed['mean_time_ms'] / cymeta_speed['mean_time_ms']
    print(f"\nSpeed:")
    print(f"  Ratio: {speed_ratio:.2f}x ({'faster' if speed_ratio > 1 else 'slower'})")
    print()
    
    return {
        "standard_size": standard_size,
        "cymeta_size": cymeta_size,
        "compression_ratio": compression_ratio,
        "size_reduction_percent": size_reduction,
        "standard_speed": standard_speed,
        "cymeta_speed": cymeta_speed,
        "speed_ratio": speed_ratio,
    }


def main():
    """Example benchmark."""
    print("=" * 60)
    print("CyMeta Benchmarking Example")
    print("=" * 60)
    print()
    
    # Create models
    in_features = 128
    out_features = 64
    
    standard_layer = nn.Linear(in_features, out_features)
    cymeta_layer = CyMetaLinear(in_features, out_features, dict_size=32)
    
    # Compare
    input_shape = (4, in_features)  # (batch_size, in_features)
    
    results = compare_models(
        standard_layer,
        cymeta_layer,
        input_shape,
        device="cpu",
    )
    
    print("=" * 60)
    print("Benchmark complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()

