import torch
import torchvision.transforms as T
import matplotlib.pyplot as plt

from torch.utils.data import Dataset, DataLoader
from torchvision.datasets import ImageFolder
from facenet_pytorch import MTCNN, InceptionResnetV1
from pathlib import Path
from typing import Union, Callable

data_root = Path('.')
# create the MTCNN network
transform = MTCNN(image_size=160, margin=18)

ds = ImageFolder(root=data_root / 'faces', transform=transform)
# our dataset is so small that the batch_size can equal to its lenght
dl = DataLoader(ds, batch_size=len(ds))

ds[1]

from torchvision.utils import make_grid
import matplotlib.pyplot as plt

fig = plt.figure(figsize=(10, 10))
for (imgs, labels) in dl:
    plt.imshow(make_grid(imgs, nrow=4).permute(1,2,0).numpy())
    break

resnet = InceptionResnetV1(pretrained='vggface2').eval()

with torch.no_grad():
    for (imgs, labels) in dl:
        embs = resnet(imgs)
        break

embs.shape

class UnNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        """
        Args:model[0][2]
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
        Returns:
            Tensor: Normalized image.
        """
        x = tensor.clone()
        for t, m, s in zip(x, self.mean, self.std):
            t.mul_(s).add_(m)
            # The normalize code -> t.sub_(m).div_(s)
        return x

# My un-normalization function is not quite right since the images are normalized with their own values and not with the std and mean of the whole dataset.
un_normalize = UnNormalize([ 0.485, 0.456, 0.406 ], (0.247, 0.243, 0.261))

import seaborn as sns
import numpy as np

similarity_matrix = torch.zeros(embs.shape[0], embs.shape[0])

for i in range(embs.shape[0]):
    for j in range(embs.shape[0]):
        similarity_matrix[i,j] = torch.cosine_similarity(embs[i].view(1, -1), embs[j].view(1, -1))


fig = plt.figure(figsize=(15, 15))

sns.heatmap(similarity_matrix.numpy(), annot = True,)

numicons = 8
for i in range(numicons):
    axicon = fig.add_axes([0.12+0.082*i,0.01,0.05,0.05])
    axicon.imshow(un_normalize(ds[i][0]).permute(1,2,0).numpy())
    axicon.set_xticks([])
    axicon.set_yticks([])

    axicon = fig.add_axes([0, 0.15 + 0.092 * i,.05,0.05])
    axicon.imshow(un_normalize(ds[len(ds) - 1 - i][0]).permute(1,2,0).numpy())
    axicon.set_xticks([])
    axicon.set_yticks([])

from matplotlib.offsetbox import OffsetImage, AnnotationBbox

def pca(x: torch.Tensor, k: int = 2) -> torch.Tensor:
    """
    From http://agnesmustar.com/2017/11/01/principal-component-analysis-pca-implemented-pytorch/
    """
    # preprocess the data
    X_mean = torch.mean(x, 0)
    x = x - X_mean.expand_as(x)
    # svd
    U, S, V = torch.svd(torch.t(x))
    return torch.mm(x, U[:, :k])

points = pca(embs, k=2)
plt.rcParams["figure.figsize"] = (12,12)

fig, ax = plt.figure(), plt.subplot(111)
plt.scatter(points[:,0], points[:,1])
for i, p in enumerate(points):
    x, y = p[0], p[1]
    img = un_normalize(ds[i][0])
    img_np = img.permute(1, 2, 0).numpy().squeeze()
    ab = AnnotationBbox(OffsetImage(img_np, zoom=0.6), (x, y), frameon=False)
    ax.add_artist(ab)

plt.plot()

from dataclasses import dataclass, field
from typing import List, Callable
from PIL import Image

@dataclass
class FaceUnlock:
    
    images: List[Image.Image] = field(default_factory = list)
    th: float = 0.8
    transform: Callable =  MTCNN(image_size=160, margin=18)
    embedder: torch.nn.Module = InceptionResnetV1(pretrained='vggface2').eval()
    center: torch.Tensor = None

    def __post_init__(self):
        faces = torch.stack(list(map(self.transform, self.images)))
        embds = self.embedder(faces)

        self.center = embds.sum(0) / embds.shape[0]
    
    def __call__(self, x: Image.Image) -> bool:
        face = self.transform(x)
        emb = self.embedder(face.unsqueeze(0))

        similarity = torch.cosine_similarity(emb.view(1, -1), self.center.view(1, -1))
        is_me =  similarity > self.th

        return is_me, similarity

# load pictures of myself
me = data_root / 'faces' / 'me'
images = list(map(Image.open, me.glob('*')))
# initialize face unlock with my images
face_unlock = FaceUnlock(images)

from ipywidgets import interact, interactive, fixed, interact_manual

def unlock_with_filepath(path):
    img = Image.open(path)
    
    is_me, similarity =  face_unlock(img)
    print(f"{'🔓' if is_me else '🔒'} similarity={similarity.item():.3f}")

    fig = plt.figure()
    plt.imshow(img)
    plt.plot()

test_root = data_root / 'faces_test'

interact(unlock_with_filepath, path=list(test_root.glob('*')))