"""
Weight Retrieval Layer - Core Innovation of WRAG 2.0

This layer dynamically retrieves and composes weights during the forward pass,
rather than using fixed weights like traditional neural networks.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class WeightRetrievalLayer(nn.Module):
    """
    Dynamically retrieves and composes layer weights based on input.
    
    Traditional layer: x → W → output (W is fixed)
    This layer: x → retrieve W based on x → W → output (W is dynamic)
    
    Args:
        d_in: Input dimension
        d_out: Output dimension
        num_shards: Number of weight shards in the bank
        k: Number of shards to retrieve and compose
    """
    
    def __init__(self, d_in, d_out, num_shards=10, k=3):
        super().__init__()
        
        self.d_in = d_in
        self.d_out = d_out
        self.num_shards = num_shards
        self.k = k
        
        # Weight bank: stores multiple weight matrices (shards)
        # Shape: [num_shards, d_in, d_out]
        self.weight_bank = nn.Parameter(
            torch.randn(num_shards, d_in, d_out) * 0.02
        )
        
        # Query network: converts input to retrieval query
        # This learns what kind of weights to retrieve for each input
        self.query_net = nn.Sequential(
            nn.Linear(d_in, d_in // 2),
            nn.ReLU(),
            nn.Linear(d_in // 2, num_shards)
        )
        
        # Bias (optional, like standard linear layer)
        self.bias = nn.Parameter(torch.zeros(d_out))
    
    def forward(self, x, return_weights=False):
        """
        Forward pass with dynamic weight retrieval.
        
        Args:
            x: Input tensor [batch, d_in]
            return_weights: If True, return dynamic weight matrix instead of applying it
        
        Returns:
            If return_weights=False: output [batch, d_out], scores [batch, num_shards]
            If return_weights=True: dynamic_W [batch, d_in, d_out], scores [batch, num_shards]
        """
        batch = x.shape[0]
        
        # Step 1: Compute retrieval scores
        retrieval_scores = self.query_net(x)  # [batch, num_shards]
        
        # Step 2: Select top-k shards
        top_k_scores, top_k_indices = torch.topk(
            retrieval_scores, k=self.k, dim=1
        )  # [batch, k]
        
        # Step 3: Softmax to get composition weights
        composition_weights = F.softmax(top_k_scores, dim=1)  # [batch, k]
        
        # Step 4: Retrieve selected shards efficiently
        # Instead of indexing which causes memory explosion, use einsum
        # Create one-hot encoding for top-k selection
        selection_mask = torch.zeros(batch, self.num_shards, device=x.device)
        selection_mask.scatter_(1, top_k_indices, composition_weights)
        
        # Step 5: Compose dynamic weight matrix using einsum
        # weight_bank: [num_shards, d_in, d_out]
        # selection_mask: [batch, num_shards]
        dynamic_W = torch.einsum(
            'bs,sio->bio', 
            selection_mask, 
            self.weight_bank
        )  # [batch, d_in, d_out]
        
        if return_weights:
            return dynamic_W, retrieval_scores
        
        # Step 6: Apply dynamic weights to input
        output = torch.bmm(x.unsqueeze(1), dynamic_W).squeeze(1)  # [batch, d_out]
        
        # Add bias
        output = output + self.bias
        
        return output, retrieval_scores
    
    def get_retrieval_stats(self, x):
        """
        Get statistics about which shards are being retrieved.
        Useful for debugging and analysis.
        """
        if len(x.shape) == 3:
            query_input = x.mean(dim=1)
        else:
            query_input = x
        
        retrieval_scores = self.query_net(query_input)
        top_k_scores, top_k_indices = torch.topk(
            retrieval_scores, k=self.k, dim=1
        )
        
        return {
            'top_k_indices': top_k_indices.cpu().numpy(),
            'top_k_scores': top_k_scores.cpu().numpy(),
            'all_scores': retrieval_scores.cpu().numpy()
        }


class SimpleWeightRetrievalNet(nn.Module):
    """
    Simple network using WeightRetrievalLayers for testing.
    
    Architecture:
        Input → WRLayer1 → ReLU → WRLayer2 → Output
    """
    
    def __init__(self, input_dim, hidden_dim, output_dim, num_shards=10, k=3):
        super().__init__()
        
        self.layer1 = WeightRetrievalLayer(
            input_dim, hidden_dim, num_shards, k
        )
        self.layer2 = WeightRetrievalLayer(
            hidden_dim, output_dim, num_shards, k
        )
    
    def forward(self, x):
        x = self.layer1(x)
        x = F.relu(x)
        x = self.layer2(x)
        return x


# Quick test
if __name__ == "__main__":
    print("Testing WeightRetrievalLayer...")
    
    # Create layer
    layer = WeightRetrievalLayer(
        d_in=784,  # MNIST flattened
        d_out=128,
        num_shards=10,
        k=3
    )
    
    # Test forward pass
    batch_size = 4
    x = torch.randn(batch_size, 784)
    
    output = layer(x)
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    
    # Check retrieval stats
    stats = layer.get_retrieval_stats(x)
    print(f"Retrieved shards: {stats['top_k_indices']}")
    
    print("\n✓ WeightRetrievalLayer works!")