import numpy as np

import torch
import torch.nn.functional as F
from torchvision import transforms

from typing import Iterable, Union
from pathlib import Path


class FaceId(torch.nn.Module):
    def __init__(
        self, model_path: Path, device: str, input_shape: Iterable[int] = (112, 112)
    ):
        super().__init__()

        self.input_shape = input_shape
        self.net = torch.load(model_path, map_location=torch.device("cpu"))
        self.net.eval()

        self.transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
            ]
        )

        for n, p in self.net.named_parameters():
            assert (
                not p.requires_grad
            ), f"Parameter {n}: requires_grad: {p.requires_grad}"

        self.device = torch.device(device)
        self.to(self.device)

    def forward(
        self, img_id: Union[np.ndarray, Iterable[np.ndarray]], normalize: bool = True
    ) -> torch.Tensor:
        if isinstance(img_id, Iterable):
            img_id = [self.transform(x) for x in img_id]
            img_id = torch.stack(img_id, dim=0)
        else:
            img_id = self.transform(img_id)
            img_id = img_id.unsqueeze(0)

        img_id = img_id.to(self.device)

        img_id_112 = F.interpolate(img_id, size=self.input_shape)
        latent_id = self.net(img_id_112)
        return F.normalize(latent_id, p=2, dim=1) if normalize else latent_id