import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


class ODCNN(nn.Module, PyTorchModelHubMixin):
    def __init__(self, dropout_rate=0.5):
        super().__init__()

        # Input 3 channels, 80 bands
        # Conv 1: 7x3 filters -> 10 maps
        self.conv1 = nn.Conv2d(3, 10, kernel_size=(3, 7))
        self.relu1 = nn.ReLU()  #  ReLU improvement
        self.pool1 = nn.MaxPool2d(kernel_size=(3, 1), stride=(3, 1))

        # Conv 2: 3x3 filters -> 20 maps
        self.conv2 = nn.Conv2d(10, 20, kernel_size=(3, 3))
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=(3, 1), stride=(3, 1))

        # Flatten size calculation based on architecture
        # (20 feature maps * 8 freq bands * 7 time frames)
        self.flatten_size = 20 * 8 * 7

        # Dropout on FC inputs
        self.dropout = nn.Dropout(p=dropout_rate)

        # 256 Hidden Units
        self.fc1 = nn.Linear(self.flatten_size, 256)
        self.relu_fc = nn.ReLU()

        # Output Unit
        self.fc2 = nn.Linear(256, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)

        x = x.view(x.size(0), -1)

        x = self.dropout(x)
        x = self.fc1(x)
        x = self.relu_fc(x)

        x = self.dropout(x)
        x = self.fc2(x)
        x = self.sigmoid(x)

        return x


if __name__ == "__main__":
    from torchinfo import summary

    model = ODCNN()
    summary(model, (1, 3, 80, 15))