import torch
import torch.nn as nn

class CNNModel(nn.Module):
    def __init__(self, no_of_leads, no_of_classes):
        super(CNNModel, self).__init__()
        
        # First block
        self.conv1 = nn.Conv2d(1, 32, kernel_size=(1, 7), stride=(1, 1))
        self.bn1 = nn.BatchNorm2d(32)
        self.relu = nn.ReLU()
        
        # Second block
        self.conv2_1 = nn.Conv2d(32, 64, kernel_size=(1, 5), stride=(1, 1))
        self.bn2_1 = nn.BatchNorm2d(64)
        self.dropout = nn.Dropout(p=0.1)
        self.conv2_2 = nn.Conv2d(64, 64, kernel_size=(1, 5), stride=(1, 2))
        
        # Shortcut for second block
        self.maxpool1 = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
        self.conv_shortcut1 = nn.Conv2d(32, 64, kernel_size=(1, 1))
        
        # Third block
        self.conv3_1 = nn.Sequential(
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(64, 64, kernel_size=(1, 5), stride=(1, 1)),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(64, 64, kernel_size=(1, 5), stride=(1, 2))
        )
        
        # Shortcut for third block
        self.maxpool2 = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
        
        # Fourth block
        self.conv4_1 = nn.Sequential(
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(64, 128, kernel_size=(1, 5), stride=(1, 1)),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(128, 128, kernel_size=(1, 5), stride=(1, 2))
        )
        
        # Shortcut for fourth block
        self.maxpool3 = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
        self.conv_shortcut3 = nn.Conv2d(64, 128, kernel_size=(1, 1))
        
        # Fifth block
        self.conv5_1 = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(128, 128, kernel_size=(1, 5), stride=(1, 1)),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(128, 128, kernel_size=(1, 5), stride=(1, 2))
        )
        
        # Shortcut for fifth block
        self.maxpool4 = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
        
        # Final convolution
        self.final_conv = nn.Conv2d(128, 128, kernel_size=(no_of_leads, 1))
        self.final_bn = nn.BatchNorm2d(128)
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        
        # Fully connected layers
        self.fc1 = nn.Sequential(
            nn.Linear(128, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.1)
        )
        
        self.fc2 = nn.Sequential(
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.15)
        )
        
        self.output = nn.Linear(64, no_of_classes)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        # First block
        x1 = self.relu(self.bn1(self.conv1(x)))
        
        # Second block
        x2 = self.conv2_1(x1)
        x2 = self.bn2_1(x2)
        x2 = self.relu(x2)
        x2 = self.dropout(x2)
        x2 = self.conv2_2(x2)
        
        # First shortcut connection
        shortcut = self.maxpool1(x1)
        shortcut = self.conv_shortcut1(shortcut)
        x2 = x2 + shortcut
        
        # Third block
        x3 = self.conv3_1(x2)
        shortcut = self.maxpool2(x2)
        x3 = x3 + shortcut
        
        # Fourth block
        x4 = self.conv4_1(x3)
        shortcut = self.maxpool3(x3)
        shortcut = self.conv_shortcut3(shortcut)
        x4 = x4 + shortcut
        
        # Fifth block
        x5 = self.conv5_1(x4)
        shortcut = self.maxpool4(x4)
        x5 = x5 + shortcut
        
        # Final convolution and pooling
        x = self.final_conv(x5)
        x = self.final_bn(x)
        x = self.relu(x)
        x = self.global_pool(x)
        x = x.view(x.size(0), -1)
        
        # Fully connected layers
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.output(x)
        x = self.sigmoid(x)
        
        return x