"""
Compression Algorithm for CyMeta.

This module implements the preprocessing/compression algorithm that learns
dictionary atoms, index maps, and gating networks from full-precision weights.
"""

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

from cymeta.dictionary import MetaDictionary
from cymeta.index_map import IndexMap
from cymeta.gating import GatingNetwork
from cymeta.reconstruct import reconstruct_weights


class CyMetaCompressor:
    """
    Compresses full-precision weight matrices into CyMeta format.
    
    The compression algorithm:
    1. Learns K dictionary atoms D = {d1, ..., dK}
    2. For each row of W, finds best dictionary index p and cyclic shift s
    3. Saves index map I[row] = (p, s)
    4. Trains tiny gating net G to modulate atom usage
    
    Args:
        num_atoms: Number of dictionary atoms (default: 32)
        max_shift: Maximum absolute shift value (default: None, auto-calculated)
        gating_hidden: Hidden dimension for gating network (default: 16)
        device: Device for computation (default: 'cpu')
    """
    
    def __init__(
        self,
        num_atoms: int = 32,
        max_shift: Optional[int] = None,
        gating_hidden: int = 16,
        device: str = "cpu",
    ):
        self.num_atoms = num_atoms
        self.max_shift = max_shift
        self.gating_hidden = gating_hidden
        self.device = device
    
    def compress(
        self,
        weight_matrix: torch.Tensor,
        input_dim: Optional[int] = None,
        num_iterations: int = 100,
        learning_rate: float = 0.01,
        verbose: bool = False,
    ) -> Tuple[MetaDictionary, IndexMap, GatingNetwork]:
        """
        Compress a full-precision weight matrix into CyMeta format.
        
        Args:
            weight_matrix: Full-precision weight matrix of shape (out_dim, in_dim)
            input_dim: Input dimension for gating network (default: in_dim from weight)
            num_iterations: Number of optimization iterations (default: 100)
            learning_rate: Learning rate for optimization (default: 0.01)
            verbose: Whether to print progress (default: False)
        
        Returns:
            Tuple of (dictionary, index_map, gating_network)
        """
        weight_matrix = weight_matrix.to(self.device)
        out_dim, in_dim = weight_matrix.shape
        
        if input_dim is None:
            input_dim = in_dim
        
        # Set max_shift if not provided
        if self.max_shift is None:
            max_shift = in_dim // 2
        else:
            max_shift = self.max_shift
        
        # Initialize dictionary
        dictionary = MetaDictionary(
            num_atoms=self.num_atoms,
            atom_dim=in_dim,
        ).to(self.device)
        
        # Initialize index map with greedy matching
        index_map = self._initialize_index_map(
            weight_matrix, dictionary, max_shift
        )
        
        # Initialize gating network
        gating_network = GatingNetwork(
            input_dim=input_dim,
            num_atoms=self.num_atoms,
            hidden_dim=self.gating_hidden,
        ).to(self.device)
        
        # Joint optimization
        optimizer = torch.optim.Adam(
            list(dictionary.parameters()) + list(gating_network.parameters()),
            lr=learning_rate,
        )
        
        best_loss = float("inf")
        best_dict = None
        best_gating = None
        
        for iteration in range(num_iterations):
            optimizer.zero_grad()
            
            # Reconstruct weights (using mean token embedding for now)
            # In practice, this would use actual token embeddings
            token_embedding = torch.randn(1, input_dim, device=self.device)
            
            reconstructed = self._reconstruct_with_index_map(
                dictionary, index_map, gating_network, token_embedding
            )
            
            # Reconstruction loss
            loss = torch.nn.functional.mse_loss(reconstructed, weight_matrix)
            
            # Backward pass
            loss.backward()
            optimizer.step()
            
            if loss.item() < best_loss:
                best_loss = loss.item()
                best_dict = dictionary.state_dict().copy()
                best_gating = gating_network.state_dict().copy()
            
            if verbose and (iteration + 1) % 10 == 0:
                print(f"Iteration {iteration + 1}/{num_iterations}, Loss: {loss.item():.6f}")
        
        # Load best states
        if best_dict is not None:
            dictionary.load_state_dict(best_dict)
        if best_gating is not None:
            gating_network.load_state_dict(best_gating)
        
        return dictionary, index_map, gating_network
    
    def _initialize_index_map(
        self,
        weight_matrix: torch.Tensor,
        dictionary: MetaDictionary,
        max_shift: int,
    ) -> IndexMap:
        """
        Initialize index map using greedy matching.
        
        For each row, find the best dictionary atom and shift that minimizes
        reconstruction error.
        
        Args:
            weight_matrix: Target weight matrix
            dictionary: Dictionary with atoms
            max_shift: Maximum shift value
        
        Returns:
            Initialized IndexMap
        """
        out_dim, in_dim = weight_matrix.shape
        atom_indices = torch.zeros(out_dim, in_dim, dtype=torch.int32, device=self.device)
        shift_indices = torch.zeros(out_dim, in_dim, dtype=torch.int32, device=self.device)
        
        atoms = dictionary.get_all_atoms()
        
        # For each position in the weight matrix
        for row in range(out_dim):
            for col in range(in_dim):
                target_value = weight_matrix[row, col]
                best_error = float("inf")
                best_atom_idx = 0
                best_shift = 0
                
                # Try each atom and shift
                for atom_idx in range(self.num_atoms):
                    atom = atoms[atom_idx]
                    
                    # Try different shifts
                    for shift in range(-max_shift, max_shift + 1):
                        shifted_atom = torch.roll(atom, shifts=shift, dims=0)
                        # Use the value at position col (wrapped)
                        wrapped_col = (col - shift) % in_dim
                        value = shifted_atom[wrapped_col]
                        
                        error = abs(target_value.item() - value.item())
                        
                        if error < best_error:
                            best_error = error
                            best_atom_idx = atom_idx
                            best_shift = shift
                
                atom_indices[row, col] = best_atom_idx
                shift_indices[row, col] = best_shift
        
        return IndexMap(atom_indices, shift_indices, max_shift)
    
    def _reconstruct_with_index_map(
        self,
        dictionary: MetaDictionary,
        index_map: IndexMap,
        gating_network: GatingNetwork,
        token_embedding: torch.Tensor,
    ) -> torch.Tensor:
        """
        Reconstruct weights using index map and gating network.
        
        Args:
            dictionary: MetaDictionary
            index_map: IndexMap
            gating_network: GatingNetwork
            token_embedding: Token embedding for gating
        
        Returns:
            Reconstructed weight matrix
        """
        return reconstruct_weights(
            dictionary, index_map, gating_network, token_embedding
        )
    
    def compress_layer(
        self,
        layer: nn.Linear,
        num_iterations: int = 100,
        learning_rate: float = 0.01,
        verbose: bool = False,
    ) -> Dict:
        """
        Compress a PyTorch Linear layer.
        
        Args:
            layer: nn.Linear layer to compress
            num_iterations: Number of optimization iterations
            learning_rate: Learning rate
            verbose: Whether to print progress
        
        Returns:
            Dictionary containing compressed components
        """
        weight = layer.weight.data
        input_dim = layer.in_features
        
        dictionary, index_map, gating_network = self.compress(
            weight,
            input_dim=input_dim,
            num_iterations=num_iterations,
            learning_rate=learning_rate,
            verbose=verbose,
        )
        
        return {
            "dictionary": dictionary,
            "index_map": index_map,
            "gating_network": gating_network,
            "bias": layer.bias,
        }

