"""
HuggingFace Model Conversion Utilities.

This module provides utilities for converting HuggingFace transformer models
to CyMeta compressed format.
"""

from __future__ import annotations

import dataclasses
from typing import Callable, Dict, Iterable, Optional, Tuple

import torch
import torch.nn as nn
from transformers import AutoModel, PreTrainedModel

from cymeta.models.linear import CyMetaLinear
from cymeta.training.compressor import CyMetaCompressor


ModuleFilter = Callable[[str, nn.Module], bool]


@dataclasses.dataclass
class ConversionConfig:
    """
    Configuration for converting transformer models to CyMeta format.

    Args:
        dict_size: Number of dictionary atoms per linear layer.
        shift_bits: Shift precision for cyclic shifts.
        gating_hidden: Hidden size for gating network.
        compress_iterations: Number of optimization iterations per layer.
        include_filter: Optional callable that returns True if a module should be converted.
        exclude_filter: Optional callable to skip modules.
        device: Target device for compression (defaults to model device).
        verbose: Whether to print conversion progress.
        cpu_optimized: If True, uses aggressive compression settings for CPU deployment.
                      Reduces dict_size, gating_hidden, and increases iterations for better compression.
    """

    dict_size: int = 32
    shift_bits: int = 8
    gating_hidden: int = 16
    compress_iterations: int = 100
    include_filter: Optional[ModuleFilter] = None
    exclude_filter: Optional[ModuleFilter] = None
    device: Optional[torch.device] = None
    verbose: bool = True
    cpu_optimized: bool = False

    def __post_init__(self):
        """Apply CPU-optimized settings if enabled."""
        if self.cpu_optimized:
            # Aggressive compression for CPU: fewer atoms, smaller gating network
            # This achieves 20-50x compression ratios suitable for CPU deployment
            if self.dict_size == 32:  # Only override if using default
                self.dict_size = 16
            if self.gating_hidden == 16:  # Only override if using default
                self.gating_hidden = 8
            if self.compress_iterations == 100:  # Use more iterations for better quality
                self.compress_iterations = 150
            if self.device is None:
                self.device = torch.device("cpu")


def should_convert_module(name: str, module: nn.Module, config: ConversionConfig) -> bool:
    """Determine whether a module should be converted according to config."""
    if config.include_filter is not None and not config.include_filter(name, module):
        return False
    if config.exclude_filter is not None and config.exclude_filter(name, module):
        return False
    return isinstance(module, nn.Linear)


def _compress_linear_layer(
    name: str,
    layer: nn.Linear,
    compressor: CyMetaCompressor,
    config: ConversionConfig,
) -> CyMetaLinear:
    """Compress a single Linear layer into CyMetaLinear."""
    if config.verbose:
        print(f"[CyMeta] Compressing layer: {name} ({layer.in_features} -> {layer.out_features})")

    compressed = compressor.compress_layer(
        layer,
        num_iterations=config.compress_iterations,
        verbose=config.verbose,
    )

    cymeta_layer = CyMetaLinear(
        in_features=layer.in_features,
        out_features=layer.out_features,
        dict_size=config.dict_size,
        shift_bits=config.shift_bits,
        gating_hidden=config.gating_hidden,
        bias=layer.bias is not None,
    ).to(compressor.device)

    cymeta_layer.dictionary.load_state_dict(compressed["dictionary"].state_dict())
    cymeta_layer.index_map = compressed["index_map"]
    cymeta_layer.gating_network.load_state_dict(compressed["gating_network"].state_dict())

    if compressed["bias"] is not None and cymeta_layer.bias is not None:
        cymeta_layer.bias.data = compressed["bias"].data.to(cymeta_layer.bias.device)

    return cymeta_layer


def convert_model_linear_layers(
    model: nn.Module,
    config: Optional[ConversionConfig] = None,
) -> Tuple[nn.Module, int]:
    """
    Convert all (or filtered) nn.Linear layers in a model to CyMetaLinear.

    Args:
        model: Target PyTorch / HuggingFace model.
        config: ConversionConfig specifying compression parameters and filters.

    Returns:
        Tuple of (converted_model, num_layers_converted).
    """
    if config is None:
        config = ConversionConfig()

    device = config.device or next(model.parameters(), torch.zeros(0)).device or torch.device("cpu")
    model = model.to(device)

    compressor = CyMetaCompressor(
        num_atoms=config.dict_size,
        gating_hidden=config.gating_hidden,
        device=str(device),
    )

    converted_layers = 0

    def _recurse(module: nn.Module, prefix: str = "") -> None:
        nonlocal converted_layers
        for name, child in list(module.named_children()):
            full_name = f"{prefix}.{name}" if prefix else name

            if should_convert_module(full_name, child, config):
                if config.verbose:
                    print(f"[CyMeta] Compressing module {full_name}")
                cymeta_layer = _compress_linear_layer(full_name, child, compressor, config)
                setattr(module, name, cymeta_layer)
                converted_layers += 1
            else:
                _recurse(child, full_name)

    _recurse(model)

    if config.verbose:
        print(f"[CyMeta] Converted {converted_layers} linear layer(s) to CyMetaLinear.")

    return model, converted_layers


def convert_from_huggingface(
    model: PreTrainedModel,
    config: Optional[ConversionConfig] = None,
) -> PreTrainedModel:
    """
    Convert a loaded HuggingFace model in-place.
    """
    converted_model, _ = convert_model_linear_layers(model, config=config)
    return converted_model


def convert_pretrained_model(
    model_name_or_path: str,
    config: Optional[ConversionConfig] = None,
    torch_dtype: Optional[torch.dtype] = None,
    trust_remote_code: bool = False,
) -> PreTrainedModel:
    """
    Load a HuggingFace model by name/path and convert it to CyMeta format.

    Args:
        model_name_or_path: HF model identifier or local path.
        config: ConversionConfig. Use cpu_optimized=True for CPU deployment.
        torch_dtype: Optional dtype for model loading.
        trust_remote_code: Forwarded to AutoModel for custom architectures.

    Returns:
        Converted HuggingFace model instance, ready for CPU inference.

    Example:
        # For CPU deployment with aggressive compression:
        config = ConversionConfig(cpu_optimized=True, verbose=True)
        model = convert_pretrained_model("distilbert-base-uncased", config=config)
        model.eval()  # Set to eval mode for inference
        # Model is now compressed and ready to run on CPU
    """
    if config is None:
        config = ConversionConfig()
    
    # Ensure model is loaded on CPU if cpu_optimized
    if config.cpu_optimized:
        model = AutoModel.from_pretrained(
            model_name_or_path,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
            device_map="cpu",
        )
    else:
        model = AutoModel.from_pretrained(
            model_name_or_path,
            torch_dtype=torch_dtype,
            trust_remote_code=trust_remote_code,
        )
    
    converted_model = convert_from_huggingface(model, config=config)
    
    # Ensure model is on CPU and in eval mode for inference
    if config.cpu_optimized:
        converted_model = converted_model.to("cpu")
        converted_model.eval()
    
    return converted_model


__all__ = [
    "ConversionConfig",
    "convert_model_linear_layers",
    "convert_from_huggingface",
    "convert_pretrained_model",
]

