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

# Import preprocessing constants
from .preprocessing import INTERNATIONAL_10_20_CHANNELS, NUM_CHANNELS, preprocess_batch as _preprocess


MODEL_CHANNELS = INTERNATIONAL_10_20_CHANNELS
MODEL_INPUT_SAMPLES = 1000
MODEL_SAMPLING_RATE_HZ = 250.0
MODEL_WINDOW_SECONDS = 4.0
MODEL_ENTRY_CLASS = "EEGResNetAttention"
MODEL_NUM_CLASSES = 2
MODEL_CLASS_LABELS = ["Non-IED", "IED"]
MODEL_NON_IED_CLASS_INDEX = 0
MODEL_IED_CLASS_INDICES = [1]
MODEL_DESCRIPTION = "SenuaLab patient-independent ResNet-Attention binary IED detector"
MODEL_BATCH_PREPROCESSOR = "preprocess_batch"
MODEL_REQUIRED_REFERENCE = "common average (applied by model adapter)"
MODEL_REQUIRED_FILTERS = ["1-45 Hz zero-phase Butterworth (applied by model adapter)"]
MODEL_REQUIRED_NORMALIZATION = "global four-second window z-score, clipped to [-8,8]"
MODEL_INPUT_UNIT = "scale-invariant after window z-score"
MODEL_SOURCE_SIGNAL_POLICY = "raw"
MODEL_REQUIRES_FULL_WINDOW = True
MODEL_REQUIRES_ALL_CHANNELS = True
MODEL_DEFAULT_THRESHOLD = 0.6944227814674377
MODEL_DEFAULT_STEP_MS = 500.0
MODEL_DEFAULT_PAD_POLICY = "skip"
MODEL_DEFAULT_BATCH_SIZE = 32
MODEL_DEFAULT_BATCH_MEMORY_MB = 96.0
MODEL_VALIDATION_NOTE = "Threshold selected on the vEpiSet validation subjects; not externally validated."
MODEL_PREPROCESSING_NOTE = "The adapter applies the exact released 1-45 Hz, common-average, resampling, and global-window z-score pipeline."


def preprocess_batch(batch, source_sfreq=None, channel_names=None):
    return _preprocess(
        batch,
        source_sfreq=source_sfreq,
        target_sfreq=250,
        target_samples=1000,
        channel_names=channel_names,
    )


class Simple1DCNN(nn.Module):
    """
    Simple 1D CNN for EEG IED classification.
    
    Architecture:
    - 4 convolutional blocks with BatchNorm, ReLU, and MaxPool
    - Fully connected classifier with dropout
    
    Default configuration uses 19 channels (standard 10-20) to match
    the vEpiSet preprocessing pipeline.
    """
    
    def __init__(self, num_channels=NUM_CHANNELS, num_classes=2, input_length=2000):
        """
        Args:
            num_channels: Number of EEG channels (default 19 for standard 10-20)
            num_classes: Number of output classes (default 2: IED vs Non-IED)
            input_length: Number of time samples (default 2000 = 4 seconds at 500 Hz)
        """
        super(Simple1DCNN, self).__init__()
        
        self.num_channels = num_channels
        self.input_length = input_length
        
        self.features = nn.Sequential(
            # Block 1
            nn.Conv1d(in_channels=num_channels, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # Block 2
            nn.Conv1d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # Block 3
            nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
            
            # Block 4
            nn.Conv1d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
        )
        
        # Calculate size after convolutions
        # Input: 2000
        # Pool 1: 1000
        # Pool 2: 500
        # Pool 3: 250
        # Pool 4: 125
        self.flatten_size = 256 * (input_length // 16)  # 16 = 2^4 (4 pooling layers)
        
        self.classifier = nn.Sequential(
            nn.Linear(self.flatten_size, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x


# ============================================================================
# Advanced Model Architectures for Better IED Detection
# ============================================================================

class SqueezeExcitation1D(nn.Module):
    """Squeeze-and-Excitation block for channel attention."""
    
    def __init__(self, channels, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool1d(1)
        self.fc = 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.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1)
        return x * y.expand_as(x)


class ResidualBlock1D(nn.Module):
    """Residual block with optional squeeze-excitation."""
    
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, 
                 downsample=None, use_se=True):
        super().__init__()
        padding = kernel_size // 2
        
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, 
                               stride=stride, padding=padding, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, 
                               stride=1, padding=padding, bias=False)
        self.bn2 = nn.BatchNorm1d(out_channels)
        
        self.downsample = downsample
        self.se = SqueezeExcitation1D(out_channels) if use_se else nn.Identity()
        
    def forward(self, x):
        identity = x
        
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.se(out)
        
        if self.downsample is not None:
            identity = self.downsample(x)
        
        out += identity
        out = self.relu(out)
        return out


class TemporalAttention(nn.Module):
    """Self-attention mechanism for temporal patterns."""
    
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = channels // num_heads
        self.scale = self.head_dim ** -0.5
        
        self.qkv = nn.Linear(channels, channels * 3, bias=False)
        self.proj = nn.Linear(channels, channels)
        self.norm = nn.LayerNorm(channels)
        
    def forward(self, x):
        # x: [B, C, T] -> [B, T, C]
        x = x.permute(0, 2, 1)
        B, T, C = x.shape
        
        # Residual connection
        residual = x
        x = self.norm(x)
        
        # Multi-head attention
        qkv = self.qkv(x).reshape(B, T, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, B, heads, T, head_dim]
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        
        x = (attn @ v).transpose(1, 2).reshape(B, T, C)
        x = self.proj(x)
        x = x + residual
        
        # Back to [B, C, T]
        return x.permute(0, 2, 1)


class MultiScaleConv(nn.Module):
    """Multi-scale convolution for capturing different frequency patterns."""
    
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.branch1 = nn.Sequential(
            nn.Conv1d(in_channels, out_channels // 4, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm1d(out_channels // 4),
            nn.ReLU(inplace=True)
        )
        self.branch2 = nn.Sequential(
            nn.Conv1d(in_channels, out_channels // 4, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm1d(out_channels // 4),
            nn.ReLU(inplace=True)
        )
        self.branch3 = nn.Sequential(
            nn.Conv1d(in_channels, out_channels // 4, kernel_size=7, padding=3, bias=False),
            nn.BatchNorm1d(out_channels // 4),
            nn.ReLU(inplace=True)
        )
        self.branch4 = nn.Sequential(
            nn.Conv1d(in_channels, out_channels // 4, kernel_size=11, padding=5, bias=False),
            nn.BatchNorm1d(out_channels // 4),
            nn.ReLU(inplace=True)
        )
        
    def forward(self, x):
        return torch.cat([
            self.branch1(x),
            self.branch2(x),
            self.branch3(x),
            self.branch4(x)
        ], dim=1)


class EEGResNetAttention(nn.Module):
    """
    Advanced EEG IED Detector with ResNet backbone and Attention.
    
    Features:
    - Multi-scale initial convolution for capturing different EEG frequencies
    - Residual blocks with Squeeze-and-Excitation attention
    - Temporal self-attention for capturing long-range dependencies
    - Global average pooling for robust feature aggregation
    """
    
    def __init__(self, num_channels=NUM_CHANNELS, num_classes=2, input_length=2000):
        super().__init__()
        
        self.num_channels = num_channels
        self.input_length = input_length
        
        # Multi-scale initial convolution
        self.stem = MultiScaleConv(num_channels, 64)
        self.stem_pool = nn.MaxPool1d(kernel_size=2, stride=2)
        
        # Residual blocks
        self.layer1 = self._make_layer(64, 64, num_blocks=2, stride=1)
        self.layer2 = self._make_layer(64, 128, num_blocks=2, stride=2)
        self.layer3 = self._make_layer(128, 256, num_blocks=2, stride=2)
        self.layer4 = self._make_layer(256, 512, num_blocks=2, stride=2)
        
        # Temporal attention
        self.temporal_attn = TemporalAttention(512, num_heads=8)
        
        # Global pooling and classifier
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(0.5)
        self.fc = nn.Linear(512, num_classes)
        
        # Initialize weights
        self._init_weights()
        
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        downsample = None
        if stride != 1 or in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm1d(out_channels)
            )
        
        layers = [ResidualBlock1D(in_channels, out_channels, stride=stride, downsample=downsample)]
        for _ in range(1, num_blocks):
            layers.append(ResidualBlock1D(out_channels, out_channels))
        
        return nn.Sequential(*layers)
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                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)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        # Multi-scale stem
        x = self.stem(x)
        x = self.stem_pool(x)
        
        # Residual blocks
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        
        # Temporal attention
        x = self.temporal_attn(x)
        
        # Global pooling and classification
        x = self.global_pool(x).squeeze(-1)
        x = self.dropout(x)
        x = self.fc(x)
        
        return x


class EEGTransformer(nn.Module):
    """
    Transformer-based EEG IED Detector.
    
    Uses a CNN backbone for initial feature extraction followed by
    Transformer encoder layers for capturing complex temporal patterns.
    """
    
    def __init__(self, num_channels=NUM_CHANNELS, num_classes=2, input_length=2000,
                 d_model=256, nhead=8, num_layers=4, dim_feedforward=512):
        super().__init__()
        
        self.num_channels = num_channels
        self.input_length = input_length
        
        # CNN backbone for initial feature extraction
        self.conv_backbone = nn.Sequential(
            nn.Conv1d(num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=3, stride=2, padding=1),
            
            nn.Conv1d(64, 128, kernel_size=5, stride=2, padding=2, bias=False),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),
            
            nn.Conv1d(128, d_model, kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm1d(d_model),
            nn.ReLU(inplace=True),
        )
        
        # Calculate sequence length after CNN
        # 2000 -> 1000 -> 500 -> 250 -> 125
        self.seq_len = input_length // 16
        
        # Positional encoding
        self.pos_encoding = self._create_positional_encoding(self.seq_len, d_model)
        
        # Transformer encoder
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=nhead, 
            dim_feedforward=dim_feedforward,
            dropout=0.1,
            activation='gelu',
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # Classification head
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.classifier = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(d_model // 2, num_classes)
        )
        
    def _create_positional_encoding(self, seq_len, d_model):
        pe = torch.zeros(seq_len, d_model)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return nn.Parameter(pe.unsqueeze(0), requires_grad=False)
        
    def forward(self, x):
        # CNN backbone: [B, C, T] -> [B, d_model, T']
        x = self.conv_backbone(x)
        
        # Reshape for transformer: [B, d_model, T'] -> [B, T', d_model]
        x = x.permute(0, 2, 1)
        
        # Add positional encoding
        x = x + self.pos_encoding[:, :x.size(1), :]
        
        # Transformer encoder
        x = self.transformer(x)
        
        # Back to [B, d_model, T'] for pooling
        x = x.permute(0, 2, 1)
        
        # Global pooling and classification
        x = self.global_pool(x).squeeze(-1)
        x = self.classifier(x)
        
        return x


class EEGInceptionNet(nn.Module):
    """
    Inception-style network for EEG IED detection.
    
    Uses parallel convolutions with different kernel sizes to capture
    multi-scale temporal patterns characteristic of IEDs.
    """
    
    def __init__(self, num_channels=NUM_CHANNELS, num_classes=2, input_length=2000):
        super().__init__()
        
        self.num_channels = num_channels
        self.input_length = input_length
        
        # Initial convolution
        self.conv1 = nn.Sequential(
            nn.Conv1d(num_channels, 32, kernel_size=7, stride=2, padding=3, bias=False),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
        )
        
        # Inception blocks
        self.inception1 = InceptionBlock(32, 64)
        self.inception2 = InceptionBlock(64, 128)
        self.inception3 = InceptionBlock(128, 256)
        
        # Pooling between inception blocks
        self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
        
        # Channel attention
        self.channel_attn = SqueezeExcitation1D(256)
        
        # Global pooling and classifier
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(0.5)
        self.fc = nn.Linear(256, num_classes)
        
    def forward(self, x):
        x = self.conv1(x)
        
        x = self.inception1(x)
        x = self.pool(x)
        
        x = self.inception2(x)
        x = self.pool(x)
        
        x = self.inception3(x)
        x = self.channel_attn(x)
        
        x = self.global_pool(x).squeeze(-1)
        x = self.dropout(x)
        x = self.fc(x)
        
        return x


class EEGResNetTransformer(nn.Module):
    """Hybrid model: multi-scale ResNet backbone + Transformer encoder.

    Rationale:
    - CNN/ResNet extracts local, frequency-like patterns efficiently.
    - Transformer models longer-range temporal dependencies after downsampling.
    - Outputs raw logits suitable for nn.CrossEntropyLoss.
    """

    def __init__(
        self,
        num_channels=NUM_CHANNELS,
        num_classes=2,
        input_length=2000,
        d_model=512,
        nhead=8,
        num_layers=4,
        dim_feedforward=1024,
        dropout=0.2,
    ):
        super().__init__()

        self.num_channels = num_channels
        self.input_length = input_length

        # ResNet-ish backbone (same stem idea as EEGResNetAttention)
        self.stem = MultiScaleConv(num_channels, 64)
        self.stem_pool = nn.MaxPool1d(kernel_size=2, stride=2)

        self.layer1 = self._make_layer(64, 64, num_blocks=2, stride=1)
        self.layer2 = self._make_layer(64, 128, num_blocks=2, stride=2)
        self.layer3 = self._make_layer(128, 256, num_blocks=2, stride=2)
        self.layer4 = self._make_layer(256, d_model, num_blocks=2, stride=2)

        # Transformer encoder over time (after downsampling)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation='gelu',
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Learnable positional embeddings (sequence length depends on input length)
        # Downsampling: stem_pool(/2) + layer2(/2) + layer3(/2) + layer4(/2) => /16 total
        self.seq_len = max(1, input_length // 16)
        self.pos_embed = nn.Parameter(torch.zeros(1, self.seq_len, d_model))

        self.pre_head_norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.head = nn.Linear(d_model, num_classes)

        self._init_weights()

    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        downsample = None
        if stride != 1 or in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm1d(out_channels),
            )

        layers = [ResidualBlock1D(in_channels, out_channels, stride=stride, downsample=downsample)]
        for _ in range(1, num_blocks):
            layers.append(ResidualBlock1D(out_channels, out_channels))
        return nn.Sequential(*layers)

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

        nn.init.trunc_normal_(self.pos_embed, std=0.02)

    def forward(self, x):
        # x: [B, C, T]
        x = self.stem(x)
        x = self.stem_pool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)  # [B, d_model, T']

        # [B, d_model, T'] -> [B, T', d_model]
        x = x.permute(0, 2, 1)

        # Handle dynamic T' (e.g. if input_length differs)
        if x.size(1) != self.pos_embed.size(1):
            # Simple interpolate positional embeddings along time
            pos = self.pos_embed.transpose(1, 2)  # [1, d_model, seq]
            pos = F.interpolate(pos, size=x.size(1), mode='linear', align_corners=False)
            pos = pos.transpose(1, 2)
        else:
            pos = self.pos_embed

        x = x + pos
        x = self.transformer(x)

        # Mean pool over time
        x = x.mean(dim=1)
        x = self.pre_head_norm(x)
        x = self.dropout(x)
        return self.head(x)


class InceptionBlock(nn.Module):
    """Inception block with multiple parallel convolutions."""
    
    def __init__(self, in_channels, out_channels):
        super().__init__()
        branch_channels = out_channels // 4
        
        # 1x1 conv
        self.branch1 = nn.Sequential(
            nn.Conv1d(in_channels, branch_channels, kernel_size=1, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True)
        )
        
        # 1x1 -> 3x3 conv
        self.branch2 = nn.Sequential(
            nn.Conv1d(in_channels, branch_channels, kernel_size=1, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True),
            nn.Conv1d(branch_channels, branch_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True)
        )
        
        # 1x1 -> 5x5 conv
        self.branch3 = nn.Sequential(
            nn.Conv1d(in_channels, branch_channels, kernel_size=1, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True),
            nn.Conv1d(branch_channels, branch_channels, kernel_size=5, padding=2, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True)
        )
        
        # MaxPool -> 1x1 conv
        self.branch4 = nn.Sequential(
            nn.MaxPool1d(kernel_size=3, stride=1, padding=1),
            nn.Conv1d(in_channels, branch_channels, kernel_size=1, bias=False),
            nn.BatchNorm1d(branch_channels),
            nn.ReLU(inplace=True)
        )
        
    def forward(self, x):
        return torch.cat([
            self.branch1(x),
            self.branch2(x),
            self.branch3(x),
            self.branch4(x)
        ], dim=1)


# Model factory function
def get_model(model_name='resnet_attention', num_channels=NUM_CHANNELS, num_classes=2, input_length=2000):
    """
    Factory function to get model by name.
    
    Args:
        model_name: One of 'simple_cnn', 'resnet_attention', 'transformer', 'inception', 'resnet_transformer'
        num_channels: Number of input channels
        num_classes: Number of output classes
        input_length: Input sequence length
        
    Returns:
        Model instance
    """
    models = {
        'simple_cnn': Simple1DCNN,
        'resnet_attention': EEGResNetAttention,
        'transformer': EEGTransformer,
        'inception': EEGInceptionNet,
        'resnet_transformer': EEGResNetTransformer,
    }
    
    if model_name not in models:
        raise ValueError(f"Unknown model: {model_name}. Choose from {list(models.keys())}")
    
    return models[model_name](num_channels=num_channels, num_classes=num_classes, input_length=input_length)