"""
Index Map Module for CyMeta Compression.

This module implements compact integer index maps that store which dictionary
atom to use and how much to cyclically shift it for each weight matrix element.
"""

import torch
import numpy as np
from typing import Optional, Tuple, Union
import pickle
import gzip


class IndexMap:
    """
    Stores integer index matrices for cyclic-shift addressing.
    
    Instead of storing full weight matrices, we store compact integer indices
    that indicate:
    1. Which dictionary atom to use
    2. How much to cyclically shift the atom (wrap-around shift)
    
    This enables extreme compression while maintaining reconstruction capability.
    
    Args:
        atom_indices: Tensor of shape (out_dim, in_dim) containing atom indices
        shift_indices: Tensor of shape (out_dim, in_dim) containing shift amounts
        max_shift: Maximum absolute shift value (for validation)
    
    Attributes:
        atom_indices: Integer tensor mapping each weight position to a dictionary atom
        shift_indices: Integer tensor mapping each weight position to a shift amount
        shape: Shape of the original weight matrix (out_dim, in_dim)
    """
    
    def __init__(
        self,
        atom_indices: torch.Tensor,
        shift_indices: torch.Tensor,
        max_shift: Optional[int] = None,
    ):
        # Validate inputs
        if atom_indices.shape != shift_indices.shape:
            raise ValueError(
                f"Atom indices shape {atom_indices.shape} must match "
                f"shift indices shape {shift_indices.shape}"
            )
        
        if atom_indices.dtype.is_floating_point:
            raise ValueError(f"Atom indices must be integer type, got {atom_indices.dtype}")
        
        if shift_indices.dtype.is_floating_point:
            raise ValueError(f"Shift indices must be integer type, got {shift_indices.dtype}")
        
        self.atom_indices = atom_indices
        self.shift_indices = shift_indices
        self.shape = atom_indices.shape
        self.max_shift = max_shift
        
        # Validate shift indices if max_shift is provided
        if max_shift is not None:
            if torch.any(torch.abs(shift_indices) > max_shift):
                raise ValueError(
                    f"Shift indices exceed max_shift={max_shift}. "
                    f"Found max absolute shift: {torch.abs(shift_indices).max().item()}"
                )
    
    @classmethod
    def create_empty(
        cls,
        out_dim: int,
        in_dim: int,
        num_atoms: int,
        max_shift: Optional[int] = None,
        dtype: torch.dtype = torch.int32,
    ) -> "IndexMap":
        """
        Create an empty index map with random initialization.
        
        Args:
            out_dim: Output dimension of the weight matrix
            in_dim: Input dimension of the weight matrix
            num_atoms: Number of available dictionary atoms
            max_shift: Maximum absolute shift value (default: in_dim // 2)
            dtype: Data type for indices (default: torch.int32)
        
        Returns:
            IndexMap with randomly initialized indices
        """
        if max_shift is None:
            max_shift = in_dim // 2
        
        atom_indices = torch.randint(0, num_atoms, (out_dim, in_dim), dtype=dtype)
        shift_indices = torch.randint(-max_shift, max_shift + 1, (out_dim, in_dim), dtype=dtype)
        
        return cls(atom_indices, shift_indices, max_shift)
    
    def get_indices(self, row: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get atom and shift indices for a specific row.
        
        Args:
            row: Row index (0 <= row < out_dim)
        
        Returns:
            Tuple of (atom_indices, shift_indices) for the specified row
        """
        if row < 0 or row >= self.shape[0]:
            raise IndexError(f"Row index {row} out of range [0, {self.shape[0]})")
        
        return self.atom_indices[row], self.shift_indices[row]
    
    def get_atom_index(self, row: int, col: int) -> int:
        """
        Get the atom index for a specific (row, col) position.
        
        Args:
            row: Row index
            col: Column index
        
        Returns:
            Atom index for the specified position
        """
        return self.atom_indices[row, col].item()
    
    def get_shift_index(self, row: int, col: int) -> int:
        """
        Get the shift index for a specific (row, col) position.
        
        Args:
            row: Row index
            col: Column index
        
        Returns:
            Shift amount for the specified position
        """
        return self.shift_indices[row, col].item()
    
    def apply_wrap_around_indexing(
        self,
        base_indices: torch.Tensor,
        shifts: torch.Tensor,
        dim_size: int
    ) -> torch.Tensor:
        """
        Apply wrap-around indexing for cyclic addressing.
        
        This ensures that indices wrap around when they exceed boundaries,
        enabling circular shift operations.
        
        Args:
            base_indices: Base indices tensor
            shifts: Shift amounts tensor
            dim_size: Size of the dimension for wrapping
        
        Returns:
            Wrapped indices tensor
        """
        shifted = base_indices + shifts
        # Apply modulo for wrap-around
        return shifted % dim_size
    
    def to(self, device: Union[str, torch.device]) -> "IndexMap":
        """
        Move index map to a specific device.
        
        Args:
            device: Target device (e.g., 'cuda', 'cpu')
        
        Returns:
            New IndexMap on the specified device
        """
        return IndexMap(
            self.atom_indices.to(device),
            self.shift_indices.to(device),
            self.max_shift,
        )
    
    def cpu(self) -> "IndexMap":
        """Move index map to CPU."""
        return self.to("cpu")
    
    def cuda(self, device: Optional[int] = None) -> "IndexMap":
        """Move index map to CUDA device."""
        if device is not None:
            return self.to(f"cuda:{device}")
        return self.to("cuda")
    
    def save(self, filepath: str, compress: bool = True) -> None:
        """
        Save index map to disk.
        
        Args:
            filepath: Path to save the index map
            compress: Whether to use gzip compression (default: True)
        """
        data = {
            "atom_indices": self.atom_indices.cpu().numpy(),
            "shift_indices": self.shift_indices.cpu().numpy(),
            "max_shift": self.max_shift,
            "shape": self.shape,
        }
        
        if compress:
            with gzip.open(filepath, "wb") as f:
                pickle.dump(data, f)
        else:
            with open(filepath, "wb") as f:
                pickle.dump(data, f)
    
    @classmethod
    def load(cls, filepath: str) -> "IndexMap":
        """
        Load index map from disk.
        
        Args:
            filepath: Path to load the index map from
        
        Returns:
            Loaded IndexMap instance
        """
        try:
            with gzip.open(filepath, "rb") as f:
                data = pickle.load(f)
        except (gzip.BadGzipFile, OSError):
            with open(filepath, "rb") as f:
                data = pickle.load(f)
        
        atom_indices = torch.from_numpy(data["atom_indices"])
        shift_indices = torch.from_numpy(data["shift_indices"])
        
        return cls(atom_indices, shift_indices, data.get("max_shift"))
    
    def get_compression_ratio(self, original_dtype: torch.dtype = torch.float32) -> float:
        """
        Calculate compression ratio compared to full weight matrix.
        
        Args:
            original_dtype: Data type of the original weight matrix
        
        Returns:
            Compression ratio (original_size / compressed_size)
        """
        original_size = np.prod(self.shape) * torch.finfo(original_dtype).bits // 8
        compressed_size = (
            self.atom_indices.numel() * self.atom_indices.element_size() +
            self.shift_indices.numel() * self.shift_indices.element_size()
        )
        return original_size / compressed_size if compressed_size > 0 else float("inf")
    
    def __repr__(self) -> str:
        return (
            f"IndexMap(shape={self.shape}, "
            f"max_shift={self.max_shift}, "
            f"compression_ratio={self.get_compression_ratio():.2f}x)"
        )

