import math
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
# Add more imports if required

# Sample Transformation function
# YOUR CODE HERE for changing the Transformation values.
trnscm = transforms.Compose([

    transforms.Grayscale(
        num_output_channels=1
    ),

    transforms.Resize((224, 224)),

    transforms.ToTensor(),

    transforms.Normalize(
        mean=[0.5],
        std=[0.5]
    )
])

##Example Network
class Siamese(torch.nn.Module):
    def __init__(self):

        super(Siamese, self).__init__()

        ################################################################
        # CNN Feature Extractor
        ################################################################

        self.cnn1 = nn.Sequential(

            # Block 1
            nn.Conv2d(
                in_channels=1,
                out_channels=32,
                kernel_size=3,
                padding=1
            ),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            # 224 -> 112

            # Block 2
            nn.Conv2d(
                in_channels=32,
                out_channels=64,
                kernel_size=3,
                padding=1
            ),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            # 112 -> 56

            # Block 3
            nn.Conv2d(
                in_channels=64,
                out_channels=128,
                kernel_size=3,
                padding=1
            ),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),

            # 56 -> 28

            # Block 4
            nn.Conv2d(
                in_channels=128,
                out_channels=256,
                kernel_size=3,
                padding=1
            ),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            ################################################################
            # Makes output independent of image size
            ################################################################

            nn.AdaptiveAvgPool2d((1, 1))
        )

        ################################################################
        # Embedding Head
        ################################################################

        self.fc1 = nn.Sequential(

            nn.Flatten(),

            nn.Linear(256, 128),
            nn.ReLU(inplace=True),

            nn.Dropout(0.3),

            nn.Linear(128, 128)
        )

    ####################################################################
    # Generate embedding for one image
    ####################################################################

    def forward_once(self, x):

        output = self.cnn1(x)

        output = self.fc1(output)

        ################################################################
        # L2 normalize embeddings
        ################################################################

        output = F.normalize(output, p=2, dim=1)

        return output

    ####################################################################
    # Siamese forward
    ####################################################################

    def forward(self, input1, input2):

        output1 = self.forward_once(input1)

        output2 = self.forward_once(input2)

        return output1, output2
      
##########################################################################################################
## Sample classification network (Specify if you are using a pytorch classifier during the training)    ##
## classifier = nn.Sequential(nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Linear...)           ##
##########################################################################################################

# YOUR CODE HERE for pytorch classifier

# Definition of classes as dictionary
classes = ['person1','person2','person3','person4','person5','person6','person7']