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


# ---------------------------
# Transformation Function
# ---------------------------
# Same transforms as used during training in Colab
trnscm = transforms.Compose([
    transforms.Resize((100, 100)),
    transforms.ToTensor()
])


# ---------------------------
# Siamese Network Definition
# ---------------------------
class Siamese(torch.nn.Module):
    def __init__(self):
        super(Siamese, self).__init__()
        
        # CNN layers (same as your Colab model)
        self.cnn1 = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(1, 4, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(4),

            nn.ReflectionPad2d(1),
            nn.Conv2d(4, 8, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(8),

            nn.ReflectionPad2d(1),
            nn.Conv2d(8, 8, kernel_size=3),
            nn.ReLU(inplace=True),
            nn.BatchNorm2d(8)
        )

        # Fully connected layers
        self.fc1 = nn.Sequential(
            nn.Linear(8 * 100 * 100, 500),
            nn.ReLU(inplace=True),
            nn.Linear(500, 500),
            nn.ReLU(inplace=True),
            nn.Linear(500, 5)
        )
        
    def forward_once(self, x):
        # Forward pass for one image
        output = self.cnn1(x)
        output = output.view(output.size()[0], -1)
        output = self.fc1(output)
        return output
        
    def forward(self, x1, x2):
        # Forward pass for both images
        output1 = self.forward_once(x1)
        output2 = self.forward_once(x2)
        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...)           ##
##########################################################################################################

# Not used for face similarity — so keep it as None
classifier = None


# ---------------------------
# Class labels (optional)
# ---------------------------
classes = ['person1', 'person2', 'person3', 'person4', 'person5', 'person6', 'person7']