"""
Serialization Utilities for CyMeta Models.

This module provides functions for saving and loading CyMeta compressed models,
including dictionaries, index maps, and gating networks.
"""

import torch
import json
import os
from typing import Dict, Optional, Any
from pathlib import Path
from cymeta.dictionary import MetaDictionary
from cymeta.index_map import IndexMap
from cymeta.gating import GatingNetwork


def save_cymeta_model(
    model: torch.nn.Module,
    save_dir: str,
    model_name: str = "cymeta_model",
    metadata: Optional[Dict[str, Any]] = None,
) -> None:
    """
    Save a CyMeta model to disk.
    
    Saves:
    - Model state dict
    - Dictionaries (if any)
    - Index maps (if any)
    - Gating networks (if any)
    - Metadata
    
    Args:
        model: CyMeta model to save
        save_dir: Directory to save the model
        model_name: Base name for saved files (default: "cymeta_model")
        metadata: Optional metadata dictionary to save
    """
    save_path = Path(save_dir)
    save_path.mkdir(parents=True, exist_ok=True)
    
    # Save model state dict
    model_path = save_path / f"{model_name}.pth"
    torch.save(model.state_dict(), model_path)
    
    # Extract and save CyMeta components
    dictionaries = {}
    index_maps = {}
    gating_networks = {}
    
    for name, module in model.named_modules():
        if hasattr(module, "dictionary"):
            dictionaries[name] = module.dictionary.state_dict()
        if hasattr(module, "index_map"):
            index_map_path = save_path / f"{model_name}_{name}_index_map.pkl"
            module.index_map.save(str(index_map_path))
            index_maps[name] = str(index_map_path)
        if hasattr(module, "gating_network"):
            gating_networks[name] = module.gating_network.state_dict()
    
    # Save dictionaries
    if dictionaries:
        dict_path = save_path / f"{model_name}_dictionaries.pth"
        torch.save(dictionaries, dict_path)
    
    # Save gating networks
    if gating_networks:
        gating_path = save_path / f"{model_name}_gating_networks.pth"
        torch.save(gating_networks, gating_path)
    
    # Save metadata
    save_metadata = {
        "model_name": model_name,
        "index_maps": index_maps,
        **(metadata or {}),
    }
    
    metadata_path = save_path / f"{model_name}_metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(save_metadata, f, indent=2)
    
    print(f"Model saved to {save_dir}")


def load_cymeta_model(
    model: torch.nn.Module,
    load_dir: str,
    model_name: str = "cymeta_model",
    strict: bool = True,
) -> torch.nn.Module:
    """
    Load a CyMeta model from disk.
    
    Args:
        model: Model instance to load weights into
        load_dir: Directory containing saved model files
        model_name: Base name for saved files (default: "cymeta_model")
        strict: Whether to strictly enforce state dict matching (default: True)
    
    Returns:
        Model with loaded weights
    """
    load_path = Path(load_dir)
    
    # Load model state dict
    model_path = load_path / f"{model_name}.pth"
    if model_path.exists():
        state_dict = torch.load(model_path, map_location="cpu")
        model.load_state_dict(state_dict, strict=strict)
    
    # Load metadata
    metadata_path = load_path / f"{model_name}_metadata.json"
    if metadata_path.exists():
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
        
        # Load index maps
        if "index_maps" in metadata:
            for name, index_map_path in metadata["index_maps"].items():
                full_path = load_path / Path(index_map_path).name
                if full_path.exists():
                    index_map = IndexMap.load(str(full_path))
                    # Set index map in model
                    parts = name.split(".")
                    module = model
                    for part in parts[:-1]:
                        module = getattr(module, part)
                    setattr(module, parts[-1], index_map)
    
    # Load dictionaries
    dict_path = load_path / f"{model_name}_dictionaries.pth"
    if dict_path.exists():
        dictionaries = torch.load(dict_path, map_location="cpu")
        for name, dict_state in dictionaries.items():
            parts = name.split(".")
            module = model
            for part in parts[:-1]:
                module = getattr(module, part)
            getattr(module, parts[-1]).load_state_dict(dict_state)
    
    # Load gating networks
    gating_path = load_path / f"{model_name}_gating_networks.pth"
    if gating_path.exists():
        gating_networks = torch.load(gating_path, map_location="cpu")
        for name, gating_state in gating_networks.items():
            parts = name.split(".")
            module = model
            for part in parts[:-1]:
                module = getattr(module, part)
            getattr(module, parts[-1]).load_state_dict(gating_state)
    
    print(f"Model loaded from {load_dir}")
    return model


def save_component(
    component: Any,
    filepath: str,
    component_type: str = "auto",
) -> None:
    """
    Save a single CyMeta component (dictionary, index_map, or gating_network).
    
    Args:
        component: Component to save
        filepath: Path to save the component
        component_type: Type of component ("dictionary", "index_map", "gating", or "auto")
    """
    if component_type == "auto":
        if isinstance(component, MetaDictionary):
            component_type = "dictionary"
        elif isinstance(component, IndexMap):
            component_type = "index_map"
        elif isinstance(component, GatingNetwork):
            component_type = "gating"
        else:
            raise ValueError(f"Cannot auto-detect component type for {type(component)}")
    
    if component_type == "dictionary":
        torch.save(component.state_dict(), filepath)
    elif component_type == "index_map":
        component.save(filepath)
    elif component_type == "gating":
        torch.save(component.state_dict(), filepath)
    else:
        raise ValueError(f"Unknown component type: {component_type}")


def load_component(
    filepath: str,
    component_type: str,
    **kwargs,
) -> Any:
    """
    Load a single CyMeta component.
    
    Args:
        filepath: Path to load the component from
        component_type: Type of component ("dictionary", "index_map", or "gating")
        **kwargs: Additional arguments for component initialization
    
    Returns:
        Loaded component
    """
    if component_type == "dictionary":
        state_dict = torch.load(filepath, map_location="cpu")
        # Need to know dimensions to create dictionary
        if "num_atoms" not in kwargs or "atom_dim" not in kwargs:
            raise ValueError("num_atoms and atom_dim required for dictionary")
        dictionary = MetaDictionary(**kwargs)
        dictionary.load_state_dict(state_dict)
        return dictionary
    elif component_type == "index_map":
        return IndexMap.load(filepath)
    elif component_type == "gating":
        state_dict = torch.load(filepath, map_location="cpu")
        # Need to know dimensions to create gating network
        if "input_dim" not in kwargs or "num_atoms" not in kwargs:
            raise ValueError("input_dim and num_atoms required for gating network")
        gating = GatingNetwork(**kwargs)
        gating.load_state_dict(state_dict)
        return gating
    else:
        raise ValueError(f"Unknown component type: {component_type}")

