"""
Meta-Dictionary Module for CyMeta Compression.

This module implements the core meta-dictionary data structure that stores
learned high-precision 1D vectors (dictionary atoms) and supports efficient
circular shift operations.
"""

import torch
import torch.nn as nn
from typing import Optional, Tuple
import numpy as np


class MetaDictionary(nn.Module):
    """
    A meta-dictionary that stores learned high-precision 1D dictionary vectors.
    
    Each dictionary atom is a learnable parameter vector that can be cyclically
    shifted to create variations. This enables compact representation of weight
    matrices through dictionary-based compression.
    
    Args:
        num_atoms: Number of dictionary atoms (typically 8-64)
        atom_dim: Dimension of each dictionary atom (should match weight dimensions)
        dtype: Data type for dictionary atoms (default: torch.float32)
        init_scale: Initialization scale for dictionary atoms (default: 0.02)
    
    Attributes:
        atoms: Learnable parameter tensor of shape (num_atoms, atom_dim)
    """
    
    def __init__(
        self,
        num_atoms: int,
        atom_dim: int,
        dtype: torch.dtype = torch.float32,
        init_scale: float = 0.02,
    ):
        super().__init__()
        self.num_atoms = num_atoms
        self.atom_dim = atom_dim
        self.dtype = dtype
        
        # Initialize dictionary atoms as learnable parameters
        self.atoms = nn.Parameter(
            torch.randn(num_atoms, atom_dim, dtype=dtype) * init_scale
        )
    
    def forward(self) -> torch.Tensor:
        """
        Returns the dictionary atoms.
        
        Returns:
            Tensor of shape (num_atoms, atom_dim) containing all dictionary atoms
        """
        return self.atoms
    
    def get_atom(self, index: int) -> torch.Tensor:
        """
        Get a specific dictionary atom by index.
        
        Args:
            index: Index of the atom to retrieve (0 <= index < num_atoms)
        
        Returns:
            Tensor of shape (atom_dim,) containing the requested atom
        """
        if index < 0 or index >= self.num_atoms:
            raise IndexError(f"Atom index {index} out of range [0, {self.num_atoms})")
        return self.atoms[index]
    
    def apply_cyclic_shift(
        self,
        atom_index: int,
        shift: int,
        dim: int = -1
    ) -> torch.Tensor:
        """
        Apply a cyclic shift (circular rotation) to a dictionary atom.
        
        This implements wrap-around shifting where elements that shift beyond
        the boundary wrap around to the beginning. This is more efficient than
        positional biases and enables expressive weight reconstruction.
        
        Args:
            atom_index: Index of the atom to shift
            shift: Number of positions to shift (can be negative for reverse shift)
            dim: Dimension along which to apply the shift (default: -1)
        
        Returns:
            Tensor of shape (atom_dim,) containing the shifted atom
        """
        atom = self.get_atom(atom_index)
        # Use torch.roll for efficient circular shift
        return torch.roll(atom, shifts=shift, dims=dim)
    
    def apply_cyclic_shifts_batch(
        self,
        atom_indices: torch.Tensor,
        shifts: torch.Tensor,
        dim: int = -1
    ) -> torch.Tensor:
        """
        Apply cyclic shifts to multiple atoms in batch.
        
        Args:
            atom_indices: Tensor of shape (batch_size,) containing atom indices
            shifts: Tensor of shape (batch_size,) containing shift amounts
            dim: Dimension along which to apply shifts (default: -1)
        
        Returns:
            Tensor of shape (batch_size, atom_dim) containing shifted atoms
        """
        batch_size = atom_indices.shape[0]
        shifted_atoms = []
        
        for i in range(batch_size):
            atom_idx = atom_indices[i].item()
            shift = shifts[i].item()
            shifted = self.apply_cyclic_shift(atom_idx, shift, dim)
            shifted_atoms.append(shifted)
        
        return torch.stack(shifted_atoms, dim=0)
    
    def get_all_atoms(self) -> torch.Tensor:
        """
        Get all dictionary atoms.
        
        Returns:
            Tensor of shape (num_atoms, atom_dim)
        """
        return self.atoms
    
    def __repr__(self) -> str:
        return (
            f"MetaDictionary(num_atoms={self.num_atoms}, "
            f"atom_dim={self.atom_dim}, dtype={self.dtype})"
        )


class ModuleTypeDictionary(nn.Module):
    """
    A collection of meta-dictionaries for different module types.
    
    Each transformer module type (Attention Q, K, V, OutputProj, FFN Up, FFN Down)
    has its own set of shared meta-dictionaries. This enables type-specific
    compression while maintaining global sharing across layers.
    
    Args:
        module_type: Name/identifier for the module type
        num_atoms: Number of dictionary atoms per module type
        atom_dim: Dimension of each dictionary atom
        init_scale: Initialization scale for dictionary atoms
    """
    
    def __init__(
        self,
        module_type: str,
        num_atoms: int,
        atom_dim: int,
        init_scale: float = 0.02,
    ):
        super().__init__()
        self.module_type = module_type
        self.dictionary = MetaDictionary(num_atoms, atom_dim, init_scale=init_scale)
    
    def forward(self) -> torch.Tensor:
        """Returns the dictionary atoms for this module type."""
        return self.dictionary()
    
    def __repr__(self) -> str:
        return (
            f"ModuleTypeDictionary(module_type='{self.module_type}', "
            f"num_atoms={self.dictionary.num_atoms}, "
            f"atom_dim={self.dictionary.atom_dim})"
        )

