"""
Weight Reconstruction Module for CyMeta Compression.

This module implements the core reconstruction algorithm that combines
dictionary atoms, cyclic shifts, and gating coefficients to reconstruct
effective weight vectors at runtime.
"""

import torch
from typing import Optional, Tuple
from cymeta.dictionary import MetaDictionary
from cymeta.index_map import IndexMap
from cymeta.gating import GatingNetwork


def reconstruct_weights(
    dictionary: MetaDictionary,
    index_map: IndexMap,
    gating_network: GatingNetwork,
    token_embedding: torch.Tensor,
    head: Optional[int] = None,
    position: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Reconstruct effective weight vectors using CyMeta components.
    
    Algorithm:
    1. Get gating coefficients from gating network
    2. For each weight position, gather dictionary atoms and apply cyclic shifts
    3. Linearly combine shifted atoms using gating coefficients
    4. Produce final weight vector
    
    Args:
        dictionary: MetaDictionary containing dictionary atoms
        index_map: IndexMap containing atom and shift indices
        gating_network: GatingNetwork for computing mixture coefficients
        token_embedding: Token embedding tensor for context-dependent gating
        head: Optional attention head index
        position: Optional position encoding
    
    Returns:
        Reconstructed weight tensor of shape (out_dim, in_dim)
    """
    # Get gating coefficients
    gating_coeffs = gating_network(token_embedding, head=head, position=position)

    if gating_coeffs.dim() == 0:
        gating_vector = gating_coeffs.unsqueeze(0)
    elif gating_coeffs.dim() == 1:
        gating_vector = gating_coeffs
    else:
        reduce_dims = tuple(range(gating_coeffs.dim() - 1))
        gating_vector = gating_coeffs.mean(dim=reduce_dims)

    # Reduce gating outputs to a single vector of shape (num_atoms,)
    if gating_coeffs.dim() == 0:
        gating_vector = gating_coeffs.unsqueeze(0)
    elif gating_coeffs.dim() == 1:
        gating_vector = gating_coeffs
    else:
        reduce_dims = tuple(range(gating_coeffs.dim() - 1))
        gating_vector = gating_coeffs.mean(dim=reduce_dims)
    
    # Get dimensions
    out_dim, in_dim = index_map.shape
    num_atoms = dictionary.num_atoms
    atom_dim = dictionary.atom_dim
    
    # Validate dimensions
    if atom_dim != in_dim:
        raise ValueError(
            f"Dictionary atom dimension {atom_dim} must match "
            f"input dimension {in_dim}"
        )
    
    # Initialize reconstructed weight matrix
    device = dictionary.atoms.device
    reconstructed = torch.zeros(out_dim, in_dim, device=device, dtype=dictionary.dtype)
    
    # Get all dictionary atoms
    atoms = dictionary.get_all_atoms()  # Shape: (num_atoms, atom_dim)
    
    # Reconstruct each row of the weight matrix
    for row in range(out_dim):
        # Get indices for this row
        atom_indices, shift_indices = index_map.get_indices(row)
        row_gating = gating_vector
        
        # Reconstruct this row by combining shifted atoms
        row_weights = torch.zeros(in_dim, device=device, dtype=dictionary.dtype)
        
        for col in range(in_dim):
            atom_idx = atom_indices[col].item()
            shift = shift_indices[col].item()
            
            # Get and shift the atom
            shifted_atom = dictionary.apply_cyclic_shift(atom_idx, shift)
            
            # Weight by gating coefficient
            row_weights += row_gating[atom_idx] * shifted_atom
        
        reconstructed[row] = row_weights
    
    return reconstructed


def reconstruct_weights_vectorized(
    dictionary: MetaDictionary,
    index_map: IndexMap,
    gating_network: GatingNetwork,
    token_embedding: torch.Tensor,
    head: Optional[int] = None,
    position: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Vectorized version of weight reconstruction for better efficiency.
    
    This version uses batched operations to reconstruct weights more efficiently,
    though it may use more memory.
    
    Args:
        dictionary: MetaDictionary containing dictionary atoms
        index_map: IndexMap containing atom and shift indices
        gating_network: GatingNetwork for computing mixture coefficients
        token_embedding: Token embedding tensor
        head: Optional attention head index
        position: Optional position encoding
    
    Returns:
        Reconstructed weight tensor of shape (out_dim, in_dim)
    """
    # Get gating coefficients
    gating_coeffs = gating_network(token_embedding, head=head, position=position)
    
    # Get dimensions
    out_dim, in_dim = index_map.shape
    num_atoms = dictionary.num_atoms
    atom_dim = dictionary.atom_dim
    
    # Validate dimensions
    if atom_dim != in_dim:
        raise ValueError(
            f"Dictionary atom dimension {atom_dim} must match "
            f"input dimension {in_dim}"
        )
    
    device = dictionary.atoms.device
    dtype = dictionary.dtype
    
    # Get all dictionary atoms
    atoms = dictionary.get_all_atoms()  # Shape: (num_atoms, atom_dim)
    
    # Reduce gating outputs to a single vector of shape (num_atoms,)
    if gating_coeffs.dim() == 0:
        gating_vector = gating_coeffs.unsqueeze(0)
    elif gating_coeffs.dim() == 1:
        gating_vector = gating_coeffs
    else:
        reduce_dims = tuple(range(gating_coeffs.dim() - 1))
        gating_vector = gating_coeffs.mean(dim=reduce_dims)
    
    # Expand gating coefficients to match rows
    gating_matrix = gating_vector.unsqueeze(0).expand(out_dim, -1)
    
    # Vectorized reconstruction
    # For each row, we need to:
    # 1. Get atom and shift indices
    # 2. Apply shifts to atoms
    # 3. Weight by gating coefficients
    # 4. Sum over atoms
    
    reconstructed = torch.zeros(out_dim, in_dim, device=device, dtype=dtype)
    
    # Process each row
    for row in range(out_dim):
        atom_indices = index_map.atom_indices[row]  # Shape: (in_dim,)
        shift_indices = index_map.shift_indices[row]  # Shape: (in_dim,)
        row_gating = gating_matrix[row]  # Shape: (num_atoms,)
        
        # For each column, get the shifted atom
        # This is still a loop, but we can optimize further with advanced indexing
        row_weights = torch.zeros(in_dim, device=device, dtype=dtype)
        
        for col in range(in_dim):
            atom_idx = atom_indices[col].item()
            shift = shift_indices[col].item()
            
            # Apply cyclic shift
            shifted_atom = torch.roll(atoms[atom_idx], shifts=shift, dims=0)
            
            # Weight by gating coefficient for this atom
            row_weights += row_gating[atom_idx] * shifted_atom
        
        reconstructed[row] = row_weights
    
    return reconstructed


def reconstruct_row(
    dictionary: MetaDictionary,
    atom_indices: torch.Tensor,
    shift_indices: torch.Tensor,
    gating_coeffs: torch.Tensor,
) -> torch.Tensor:
    """
    Reconstruct a single row of weights.
    
    This is a lower-level function for reconstructing individual rows,
    useful for incremental or batched reconstruction.
    
    Args:
        dictionary: MetaDictionary containing dictionary atoms
        atom_indices: Atom indices for this row, shape (in_dim,)
        shift_indices: Shift indices for this row, shape (in_dim,)
        gating_coeffs: Gating coefficients, shape (num_atoms,)
    
    Returns:
        Reconstructed weight row, shape (in_dim,)
    """
    in_dim = atom_indices.shape[0]
    atom_dim = dictionary.atom_dim
    
    if atom_dim != in_dim:
        raise ValueError(
            f"Dictionary atom dimension {atom_dim} must match "
            f"input dimension {in_dim}"
        )
    
    device = dictionary.atoms.device
    dtype = dictionary.dtype
    atoms = dictionary.get_all_atoms()
    
    row_weights = torch.zeros(in_dim, device=device, dtype=dtype)
    
    for col in range(in_dim):
        atom_idx = atom_indices[col].item()
        shift = shift_indices[col].item()
        
        # Apply cyclic shift
        shifted_atom = torch.roll(atoms[atom_idx], shifts=shift, dims=0)
        
        # Weight by gating coefficient
        row_weights += gating_coeffs[atom_idx] * shifted_atom
    
    return row_weights

