"""
RepCNN and related architectures for wakeword detection.

This module implements the RepCNN architecture and variants optimized for
wakeword detection tasks with different computational requirements.
"""

import math
from typing import Dict, List, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F


class DepthwiseSeparableConv2d(nn.Module):
    """Depthwise separable convolution for efficient processing."""
    
    def __init__(self, in_channels: int, out_channels: int, 
                 kernel_size: int, stride: int = 1, padding: int = 0):
        super().__init__()
        self.depthwise = nn.Conv2d(
            in_channels, in_channels, kernel_size, stride, padding, 
            groups=in_channels, bias=False
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, 1, bias=False)
        
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x


class SEBlock(nn.Module):
    """Squeeze-and-Excitation block for attention mechanism."""
    
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        self.squeeze = nn.AdaptiveAvgPool2d(1)
        self.excitation = nn.Sequential(
            nn.Linear(channels, channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels, bias=False),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.squeeze(x).view(b, c)
        y = self.excitation(y).view(b, c, 1, 1)
        return x * y.expand_as(x)


class RepConvBlock(nn.Module):
    """
    Reparameterizable convolution block.
    
    During training, uses multiple paths that can be reparameterized
    into a single convolution during inference for efficiency.
    """
    
    def __init__(self, in_channels: int, out_channels: int, 
                 kernel_size: int = 3, stride: int = 1, padding: int = 1,
                 groups: int = 1, use_se: bool = False):
        super().__init__()
        
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.groups = groups
        self.use_se = use_se
        
        # Main 3x3 convolution
        self.conv_3x3 = nn.Conv2d(
            in_channels, out_channels, kernel_size, stride, padding, 
            groups=groups, bias=False
        )
        
        # 1x1 convolution branch
        self.conv_1x1 = nn.Conv2d(
            in_channels, out_channels, 1, stride, 0, groups=groups, bias=False
        ) if kernel_size > 1 else None
        
        # Identity branch (only if input and output channels match)
        self.identity = (
            nn.BatchNorm2d(out_channels) if in_channels == out_channels and stride == 1
            else None
        )
        
        # Batch normalization
        self.bn = nn.BatchNorm2d(out_channels)
        
        # SE block
        self.se = SEBlock(out_channels) if use_se else None
        
        # Activation
        self.activation = nn.ReLU(inplace=True)
        
        # Flag to track if reparameterized
        self.is_reparameterized = False
        
    def forward(self, x):
        if self.is_reparameterized:
            # Use single reparameterized convolution
            out = self.reparam_conv(x)
            out = self.bn(out)
        else:
            # Use multiple branches during training
            out = self.conv_3x3(x)
            
            if self.conv_1x1 is not None:
                out += self.conv_1x1(x)
            
            if self.identity is not None:
                out += x
            
            out = self.bn(out)
        
        if self.se is not None:
            out = self.se(out)
        
        out = self.activation(out)
        return out
    
    def reparameterize(self):
        """Reparameterize multiple branches into single convolution."""
        if self.is_reparameterized:
            return
        
        # Get weights from different branches
        w_3x3 = self.conv_3x3.weight
        
        # Initialize combined weight
        w_combined = w_3x3.clone()
        
        # Add 1x1 convolution weight (pad to 3x3)
        if self.conv_1x1 is not None:
            w_1x1 = self.conv_1x1.weight
            pad_size = (self.kernel_size - 1) // 2
            w_1x1_padded = F.pad(w_1x1, [pad_size] * 4)
            w_combined += w_1x1_padded
        
        # Add identity weight
        if self.identity is not None:
            identity_weight = torch.zeros_like(w_3x3)
            center = self.kernel_size // 2
            for i in range(self.in_channels):
                identity_weight[i, i, center, center] = 1.0
            w_combined += identity_weight
        
        # Create reparameterized convolution
        self.reparam_conv = nn.Conv2d(
            self.in_channels, self.out_channels, self.kernel_size,
            self.stride, self.padding, groups=self.groups, bias=False
        )
        self.reparam_conv.weight.data = w_combined
        
        # Remove original branches
        self.__delattr__('conv_3x3')
        if hasattr(self, 'conv_1x1'):
            self.__delattr__('conv_1x1')
        if hasattr(self, 'identity'):
            self.__delattr__('identity')
        
        self.is_reparameterized = True


class RepCNN(nn.Module):
    """
    RepCNN architecture for wakeword detection.
    
    Based on RepVGG but optimized for audio spectrograms and wakeword detection.
    Supports reparameterization for efficient inference.
    """
    
    def __init__(self, num_classes: int = 3, input_channels: int = 1,
                 num_mels: int = 40, time_frames: int = 151,
                 width_multiplier: float = 1.0, use_se: bool = False,
                 dropout_rate: float = 0.2):
        """
        Initialize RepCNN model.
        
        Args:
            num_classes: Number of output classes (e.g., 3 for custom/system/negative)
            input_channels: Number of input channels (1 for mono spectrograms)
            num_mels: Number of mel filterbank features
            time_frames: Number of time frames in spectrogram
            width_multiplier: Width multiplier for channels
            use_se: Whether to use Squeeze-and-Excitation blocks
            dropout_rate: Dropout rate for regularization
        """
        super().__init__()
        
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.num_mels = num_mels
        self.time_frames = time_frames
        self.use_se = use_se
        
        # Calculate channel dimensions
        def make_divisible(v, divisor=8):
            return max(divisor, int(v + divisor / 2) // divisor * divisor)
        
        channels = [
            make_divisible(64 * width_multiplier),
            make_divisible(128 * width_multiplier),
            make_divisible(256 * width_multiplier),
            make_divisible(512 * width_multiplier)
        ]
        
        # Initial convolution layer
        self.conv1 = RepConvBlock(
            input_channels, channels[0], kernel_size=3, stride=1, padding=1, use_se=use_se
        )
        
        # Stage 1: Process frequency information
        self.stage1 = nn.Sequential(
            RepConvBlock(channels[0], channels[0], use_se=use_se),
            RepConvBlock(channels[0], channels[1], stride=(2, 1), use_se=use_se),  # Downsample freq only
        )
        
        # Stage 2: Process both frequency and time
        self.stage2 = nn.Sequential(
            RepConvBlock(channels[1], channels[1], use_se=use_se),
            RepConvBlock(channels[1], channels[2], stride=2, use_se=use_se),  # Downsample both
        )
        
        # Stage 3: Final feature extraction
        self.stage3 = nn.Sequential(
            RepConvBlock(channels[2], channels[2], use_se=use_se),
            RepConvBlock(channels[2], channels[3], stride=2, use_se=use_se),
        )
        
        # Global pooling and classification
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(dropout_rate)
        
        # Calculate the expected feature size after convolutions
        self._calculate_feature_size()
        
        self.classifier = nn.Linear(channels[3], num_classes)
        
        # Initialize weights
        self._initialize_weights()
    
    def _calculate_feature_size(self):
        """Calculate feature map size after convolutions."""
        # This is mainly for verification - we use adaptive pooling
        with torch.no_grad():
            dummy_input = torch.zeros(1, self.input_channels, self.num_mels, self.time_frames)
            x = self.conv1(dummy_input)
            x = self.stage1(x)
            x = self.stage2(x)
            x = self.stage3(x)
            self.final_feature_size = x.shape[1:]
    
    def forward(self, x):
        """Forward pass through RepCNN."""
        # Input shape: (batch, channels, mel_bins, time_frames)
        
        x = self.conv1(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        
        # Global average pooling
        x = self.global_pool(x)
        x = x.flatten(1)
        
        # Classification
        x = self.dropout(x)
        x = self.classifier(x)
        
        return x
    
    def get_embeddings(self, x):
        """Get feature embeddings before classification."""
        x = self.conv1(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        
        x = self.global_pool(x)
        x = x.flatten(1)
        
        return x
    
    def reparameterize(self):
        """Reparameterize all RepConv blocks for inference."""
        self.conv1.reparameterize()
        for module in self.stage1:
            if isinstance(module, RepConvBlock):
                module.reparameterize()
        for module in self.stage2:
            if isinstance(module, RepConvBlock):
                module.reparameterize()
        for module in self.stage3:
            if isinstance(module, RepConvBlock):
                module.reparameterize()
    
    def _initialize_weights(self):
        """Initialize model weights."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


class ImprovedRepCNN(RepCNN):
    """
    Improved RepCNN with additional optimizations for wakeword detection.
    
    Includes temporal attention, improved skip connections, and better
    handling of the time dimension in spectrograms.
    """
    
    def __init__(self, num_classes: int = 3, input_channels: int = 1,
                 num_mels: int = 40, time_frames: int = 151,
                 width_multiplier: float = 1.0, use_se: bool = True,
                 dropout_rate: float = 0.2, use_temporal_attention: bool = True):
        """Initialize Improved RepCNN."""
        super().__init__(
            num_classes, input_channels, num_mels, time_frames,
            width_multiplier, use_se, dropout_rate
        )
        
        self.use_temporal_attention = use_temporal_attention
        
        if use_temporal_attention:
            # Add temporal attention mechanism
            channels = [
                int(64 * width_multiplier),
                int(128 * width_multiplier),
                int(256 * width_multiplier),
                int(512 * width_multiplier)
            ]
            
            self.temporal_attention = nn.Sequential(
                nn.Conv1d(channels[3], channels[3] // 4, 1),
                nn.ReLU(inplace=True),
                nn.Conv1d(channels[3] // 4, channels[3], 1),
                nn.Sigmoid()
            )
    
    def forward(self, x):
        """Forward pass with temporal attention."""
        # Input shape: (batch, channels, mel_bins, time_frames)
        
        x = self.conv1(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        
        if self.use_temporal_attention:
            # Apply temporal attention
            b, c, h, w = x.shape
            
            # Global average pool over frequency dimension
            temporal_features = F.adaptive_avg_pool2d(x, (1, w)).squeeze(2)  # (b, c, w)
            
            # Apply temporal attention
            attention_weights = self.temporal_attention(temporal_features)  # (b, c, w)
            
            # Apply attention to original features
            attention_weights = attention_weights.unsqueeze(2)  # (b, c, 1, w)
            x = x * attention_weights
        
        # Global average pooling
        x = self.global_pool(x)
        x = x.flatten(1)
        
        # Classification
        x = self.dropout(x)
        x = self.classifier(x)
        
        return x


class LightweightRepCNN(nn.Module):
    """
    Lightweight RepCNN optimized for edge devices and real-time inference.
    
    Uses depthwise separable convolutions, reduced channels, and optimized
    architecture for minimal computational requirements.
    """
    
    def __init__(self, num_classes: int = 3, input_channels: int = 1,
                 num_mels: int = 40, time_frames: int = 151,
                 width_multiplier: float = 0.5, dropout_rate: float = 0.1):
        """Initialize Lightweight RepCNN."""
        super().__init__()
        
        self.num_classes = num_classes
        self.input_channels = input_channels
        self.num_mels = num_mels
        self.time_frames = time_frames
        
        # Reduced channel dimensions for efficiency
        def make_divisible(v, divisor=8):
            return max(divisor, int(v + divisor / 2) // divisor * divisor)
        
        channels = [
            make_divisible(32 * width_multiplier),
            make_divisible(64 * width_multiplier),
            make_divisible(128 * width_multiplier),
            make_divisible(256 * width_multiplier)
        ]
        
        # Initial regular convolution
        self.conv1 = nn.Sequential(
            nn.Conv2d(input_channels, channels[0], 3, 1, 1, bias=False),
            nn.BatchNorm2d(channels[0]),
            nn.ReLU(inplace=True)
        )
        
        # Depthwise separable convolution blocks
        self.stage1 = nn.Sequential(
            self._make_depthwise_block(channels[0], channels[0]),
            self._make_depthwise_block(channels[0], channels[1], stride=(2, 1))
        )
        
        self.stage2 = nn.Sequential(
            self._make_depthwise_block(channels[1], channels[1]),
            self._make_depthwise_block(channels[1], channels[2], stride=2)
        )
        
        self.stage3 = nn.Sequential(
            self._make_depthwise_block(channels[2], channels[2]),
            self._make_depthwise_block(channels[2], channels[3], stride=2)
        )
        
        # Efficient global pooling and classification
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self.classifier = nn.Linear(channels[3], num_classes)
        
        self._initialize_weights()
    
    def _make_depthwise_block(self, in_channels: int, out_channels: int, 
                             stride: Tuple[int, int] = (1, 1)):
        """Create depthwise separable convolution block."""
        return nn.Sequential(
            # Depthwise convolution
            nn.Conv2d(in_channels, in_channels, 3, stride, 1, 
                     groups=in_channels, bias=False),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            
            # Pointwise convolution
            nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        """Forward pass through Lightweight RepCNN."""
        x = self.conv1(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        
        x = self.global_pool(x)
        x = x.flatten(1)
        
        x = self.dropout(x)
        x = self.classifier(x)
        
        return x
    
    def get_embeddings(self, x):
        """Get feature embeddings before classification."""
        x = self.conv1(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        
        x = self.global_pool(x)
        x = x.flatten(1)
        
        return x
    
    def _initialize_weights(self):
        """Initialize model weights."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


def create_repcnn_model(model_type: str = "standard", **kwargs) -> nn.Module:
    """
    Factory function to create RepCNN models.
    
    Args:
        model_type: Type of RepCNN model ('standard', 'improved', 'lightweight')
        **kwargs: Additional arguments for model initialization
        
    Returns:
        RepCNN model instance
    """
    if model_type == "standard":
        return RepCNN(**kwargs)
    elif model_type == "improved":
        return ImprovedRepCNN(**kwargs)
    elif model_type == "lightweight":
        return LightweightRepCNN(**kwargs)
    else:
        raise ValueError(f"Unknown model type: {model_type}")


def count_parameters(model: nn.Module) -> Dict[str, int]:
    """Count model parameters."""
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    return {
        "total_parameters": total_params,
        "trainable_parameters": trainable_params,
        "non_trainable_parameters": total_params - trainable_params
    }


def estimate_model_flops(model: nn.Module, input_shape: Tuple[int, ...]) -> int:
    """
    Estimate FLOPs for model inference (simplified calculation).
    
    Args:
        model: PyTorch model
        input_shape: Input tensor shape (C, H, W)
        
    Returns:
        Estimated FLOPs
    """
    def conv_flops(in_channels, out_channels, kernel_size, output_size):
        return in_channels * out_channels * kernel_size * kernel_size * output_size[0] * output_size[1]
    
    def linear_flops(in_features, out_features):
        return in_features * out_features
    
    # This is a simplified estimation - for accurate FLOPs counting,
    # use specialized libraries like thop or fvcore
    flops = 0
    
    # Approximate based on model type
    if isinstance(model, (RepCNN, ImprovedRepCNN)):
        # Rough estimation for RepCNN
        flops = input_shape[1] * input_shape[2] * 1000000  # ~1M FLOPs per pixel
    elif isinstance(model, LightweightRepCNN):
        # Lower FLOPs for lightweight model
        flops = input_shape[1] * input_shape[2] * 500000  # ~0.5M FLOPs per pixel
    
    return flops