"""
Multi-Modal 1D ResNet + Dense Multi-Task Network V2.
Improvements: SE attention, deeper backbone, better regularization.
"""
import torch
import torch.nn as nn
from config import NUM_CLASSES, IN_CHANNELS, SEQ_LEN, RESNET_DIM, DENSE_DIM, FUSED_DIM, NUM_FIT_PARAMS


class SEBlock1d(nn.Module):
    """Squeeze-and-Excitation block for channel attention."""
    
    def __init__(self, channels, reduction=4):
        super().__init__()
        self.squeeze = nn.AdaptiveAvgPool1d(1)
        self.excitation = nn.Sequential(
            nn.Linear(channels, channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels),
            nn.Sigmoid(),
        )
    
    def forward(self, x):
        b, c, _ = x.shape
        w = self.squeeze(x).squeeze(-1)  # (b, c)
        w = self.excitation(w).unsqueeze(-1)  # (b, c, 1)
        return x * w


class ResBlock1d(nn.Module):
    """1D Residual Block with SE attention and optional downsampling."""
    
    def __init__(self, in_channels, out_channels, stride=1, use_se=True):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, 
                               stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, 
                               stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm1d(out_channels)
        
        self.se = SEBlock1d(out_channels) if use_se else nn.Identity()
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_channels, out_channels, kernel_size=1, 
                         stride=stride, bias=False),
                nn.BatchNorm1d(out_channels)
            )
    
    def forward(self, x):
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.se(out)
        out += self.shortcut(x)
        out = self.relu(out)
        return out


class ResNetBackbone(nn.Module):
    """1D ResNet backbone with SE attention."""
    
    def __init__(self, in_channels=IN_CHANNELS, base_dim=64, out_dim=RESNET_DIM):
        super().__init__()
        
        self.conv1 = nn.Conv1d(in_channels, base_dim, kernel_size=7, 
                               stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm1d(base_dim)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
        
        self.layer1 = nn.Sequential(
            ResBlock1d(base_dim, base_dim),
            ResBlock1d(base_dim, base_dim)
        )
        self.layer2 = nn.Sequential(
            ResBlock1d(base_dim, base_dim * 2, stride=2),
            ResBlock1d(base_dim * 2, base_dim * 2)
        )
        self.layer3 = nn.Sequential(
            ResBlock1d(base_dim * 2, base_dim * 4, stride=2),
            ResBlock1d(base_dim * 4, base_dim * 4)
        )
        
        self.avgpool = nn.AdaptiveAvgPool1d(1)
    
    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.avgpool(x)
        x = x.squeeze(-1)
        return x


class DenseParamBranch(nn.Module):
    """Dense MLP branch for processing 4 fit parameters."""
    
    def __init__(self, in_features=NUM_FIT_PARAMS, out_features=DENSE_DIM):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 32),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(32),
            nn.Dropout(0.1),
            nn.Linear(32, out_features),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(out_features),
        )
    
    def forward(self, x):
        return self.net(x)


class ClassificationHead(nn.Module):
    """Classification head with dropout."""
    
    def __init__(self, in_features=FUSED_DIM, num_classes=NUM_CLASSES, dropout=0.4):
        super().__init__()
        self.net = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(in_features, 128),
            nn.ReLU(inplace=True),
            nn.BatchNorm1d(128),
            nn.Dropout(dropout * 0.5),
            nn.Linear(128, num_classes),
        )
    
    def forward(self, x):
        return self.net(x)


class RabiMultiTaskNet(nn.Module):
    """
    Multi-Modal Multi-Task Network V2 with SE attention.
    """
    
    def __init__(self):
        super().__init__()
        self.signal_backbone = ResNetBackbone()
        self.param_branch = DenseParamBranch()
        self.data_head = ClassificationHead()
        self.fit_head = ClassificationHead()
    
    def forward(self, signal, params):
        signal_features = self.signal_backbone(signal)
        param_features = self.param_branch(params)
        fused = torch.cat([signal_features, param_features], dim=1)
        data_quality = self.data_head(fused)
        fit_quality = self.fit_head(fused)
        return data_quality, fit_quality


def count_parameters(model):
    """Count total trainable parameters."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == '__main__':
    from config import DEVICE
    model = RabiMultiTaskNet()
    print(f"Total parameters: {count_parameters(model):,}")
    model = model.to(DEVICE)
    
    x = torch.randn(4, 2, SEQ_LEN).to(DEVICE)
    params = torch.randn(4, 4).to(DEVICE)
    dq, fq = model(x, params)
    print(f"Data quality output: {dq.shape}")
    print(f"Fit quality output: {fq.shape}")
    print(f"Device: {DEVICE}")