import torch
import torch.nn as nn
from transformers import PreTrainedModel
from .configuration_chronogrid import ChronoGridConfig

def make_cnn_backbone(img_size):
    return nn.Sequential(
        # Block 1 — kernel-3 (captures short transients)
        nn.Conv1d(img_size, 64,  kernel_size=3, padding=1),
        nn.BatchNorm1d(64), nn.ReLU(),
        # Block 2 — kernel-5 (captures broader waveform distortions)
        nn.Conv1d(64,  128, kernel_size=5, padding=2),
        nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2),  # → [B, 128, 113]
        # Block 3 — kernel-3 (deep feature abstraction)
        nn.Conv1d(128, 256, kernel_size=3, padding=1),
        nn.BatchNorm1d(256), nn.ReLU(), nn.MaxPool1d(2),  # → [B, 256, 56]
    )

class _ScaledDotAttn(nn.Module):
    '''Scaled dot-product multi-head self-attention with residual + LayerNorm.'''
    def __init__(self, d_model=256, num_heads=8, dropout=0.1):
        super().__init__()
        self.mha  = nn.MultiheadAttention(d_model, num_heads, dropout=dropout,
                                           batch_first=True)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x):
        out, w = self.mha(x, x, x, need_weights=True, average_attn_weights=True)
        return self.norm(x + out), w          # w: [B, seq, seq]

class ChronoGridModelForSequenceClassification(PreTrainedModel):
    config_class = ChronoGridConfig

    def __init__(self, config):
        super().__init__(config)
        self.num_classes = config.num_classes
        
        # Bugfix for transformers newer versions requiring this attribute
        if not hasattr(self, "all_tied_weights_keys"):
            self.all_tied_weights_keys = {}
            
        # 1. Initialize your architecture
        self.cnn = make_cnn_backbone(config.img_size)
        self.bilstm = nn.LSTM(256, 128, num_layers=2, batch_first=True, bidirectional=True, dropout=0.3)
        self.attention = _ScaledDotAttn(d_model=256, num_heads=8)
        self.head = nn.Sequential(
            nn.Dropout(0.3),
            nn.Linear(256, 128), nn.ReLU(),
            nn.Linear(128, self.num_classes),
        )

    def forward(self, x, labels=None):
        # x shape: [B, 227, 227]
        x = self.cnn(x).permute(0, 2, 1)     
        x, _ = self.bilstm(x)                 
        x, attn_w = self.attention(x)         
        feat = x.mean(dim=1)                  
        logits = self.head(feat)
        
        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_classes), labels.view(-1))

        return (loss, logits) if loss is not None else (logits,)