import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image

def ensure_gray(image: Image.Image) -> Image.Image:
    if image.mode != "L":
        image = image.convert("L")
    return image

trnscm = transforms.Compose([
    transforms.Lambda(ensure_gray),
    transforms.Resize((100, 100)),
    transforms.ToTensor(),
])

class Siamese(nn.Module):
    def __init__(self):
        super().__init__()
        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),
        )
        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):
        out = self.cnn1(x)
        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        return out

    def forward(self, x1, x2):
        return self.forward_once(x1), self.forward_once(x2)

class FaceClassifier(nn.Module):
    def __init__(self, input_dim=5, num_classes=7):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes),
        )

    def forward(self, x):
        return self.fc(x)

# Update to match your captured_face_images ImageFolder order
classes = ["person1","person2","person3","person4","person5","person6","person7"]