import torch.nn as nn
import torch
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.modeling_utils import PreTrainedModel
from alexnet_model.configuration_alexnet import AlexNetConfig

class AlexNetPneumoniaClassification(PreTrainedModel):
    config_class = AlexNetConfig

    def __init__(self, config):
        super(AlexNetPneumoniaClassification, self).__init__(config)
        self.num_labels = config.num_labels
        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0)
        self.conv2 = nn.Conv2d(96, 256, kernel_size=5, stride=1,padding=2)                              
        self.conv3 = nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1)
        self.conv5 = nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(256*6*6, 4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3  = nn.Linear(4096, config.num_labels)
        
    def forward(self, pixel_values, labels=None):
        x = torch.relu(self.conv1(pixel_values))
        x = torch.max_pool2d(x, kernel_size=3, stride=2, padding=0)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, kernel_size=3, stride=2, padding=0)
        x = torch.relu(self.conv3(x))
        x = torch.relu(self.conv4(x))
        x = torch.relu(self.conv5(x))
        x = torch.max_pool2d(x, kernel_size=3, stride=2, padding=0)
        x = x.view(-1, 256*6*6)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        logits = self.fc3(x)        
        
        loss = None
        if labels is not None:
          loss_fct = nn.CrossEntropyLoss()
          loss = loss_fct(logits.view(-1, self.num_labels), labels)
          return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
          )
        
        return SequenceClassifierOutput(
          logits=torch.softmax(logits, dim=1),
        )