"""
CPU-Optimized Utilities for CyMeta.

This module provides utilities for ensuring CyMeta-compressed models
run efficiently on CPU without requiring GPU.
"""

import torch
import torch.nn as nn
from typing import Optional


def ensure_cpu_model(model: nn.Module) -> nn.Module:
    """
    Ensure a model is on CPU and optimized for CPU inference.
    
    Args:
        model: PyTorch model (preferably CyMeta-compressed)
    
    Returns:
        Model moved to CPU with eval mode enabled
    """
    model = model.to("cpu")
    model.eval()
    
    # Disable gradient computation for inference
    for param in model.parameters():
        param.requires_grad = False
    
    return model


def estimate_model_size_mb(model: nn.Module) -> float:
    """
    Estimate model size in megabytes.
    
    Args:
        model: PyTorch model
    
    Returns:
        Estimated size in MB
    """
    total_params = sum(p.numel() for p in model.parameters())
    # Assume float32 (4 bytes per parameter)
    size_bytes = total_params * 4
    size_mb = size_bytes / (1024 * 1024)
    return size_mb


def get_compression_summary(model: nn.Module) -> dict:
    """
    Get compression summary for a CyMeta-compressed model.
    
    Args:
        model: CyMeta-compressed model
    
    Returns:
        Dictionary with compression statistics
    """
    from cymeta.models.linear import CyMetaLinear
    
    total_layers = 0
    compressed_layers = 0
    total_compression_ratio = 0.0
    
    def _count_layers(module: nn.Module):
        nonlocal total_layers, compressed_layers, total_compression_ratio
        
        for child in module.children():
            if isinstance(child, nn.Linear):
                total_layers += 1
            elif isinstance(child, CyMetaLinear):
                total_layers += 1
                compressed_layers += 1
                total_compression_ratio += child.get_compression_ratio()
            else:
                _count_layers(child)
    
    _count_layers(model)
    
    avg_compression = (
        total_compression_ratio / compressed_layers 
        if compressed_layers > 0 
        else 0.0
    )
    
    return {
        "total_layers": total_layers,
        "compressed_layers": compressed_layers,
        "average_compression_ratio": avg_compression,
        "compression_coverage": compressed_layers / total_layers if total_layers > 0 else 0.0,
    }


__all__ = [
    "ensure_cpu_model",
    "estimate_model_size_mb",
    "get_compression_summary",
]

