from torch.utils.data import Dataset
from PIL import Image
import PIL.ImageOps
import random
import numpy as np
import torch
import torchvision.transforms as transforms


class Config():
    training_dir = "./data/faces/training/"
    testing_dir = "./data/faces/testing/"
    train_batch_size = 64
    train_number_epochs = 100

class SiameseNetworkDataset(Dataset):

    def __init__(self,imageFolderDataset,transform=None,should_invert=True):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert

    def __getitem__(self,index):
        img0_tuple = random.choice(self.imageFolderDataset.imgs)
        #we need to make sure approx 50% of images are in the same class
        should_get_same_class = random.randint(0,1)
        if should_get_same_class:
            while True:
                #keep looping till the same class image is found
                img1_tuple = random.choice(self.imageFolderDataset.imgs)
                if img0_tuple[1]==img1_tuple[1]:
                    break
        else:
            img1_tuple = random.choice(self.imageFolderDataset.imgs)

        img0 = Image.open(img0_tuple[0])
        img1 = Image.open(img1_tuple[0])
        img0 = img0.convert("L")
        img1 = img1.convert("L")

        if self.should_invert:
            img0 = PIL.ImageOps.invert(img0)
            img1 = PIL.ImageOps.invert(img1)

        if self.transform is not None:
            img0 = self.transform(img0)
            img1 = self.transform(img1)

        return img0, img1 , torch.from_numpy(np.array([int(img1_tuple[1]!=img0_tuple[1])],dtype=np.float32))

    def __len__(self):
        return len(self.imageFolderDataset.imgs)

class Triplet_dataset(Dataset):

    def __init__(self, imageFolderDataset, transform=None, should_invert=True):
        self.imageFolderDataset = imageFolderDataset
        self.transform = transform
        self.should_invert = should_invert

    def __getitem__(self, index):
        anchor_tuple = random.choice(self.imageFolderDataset.imgs)
        while True:
            positive_tuple = random.choice(self.imageFolderDataset.imgs)
            if positive_tuple[1] == anchor_tuple[1]:
                break

        while True:
            negative_tuple = random.choice(self.imageFolderDataset.imgs)
            if negative_tuple[1] != anchor_tuple[1]:
                break

        anchor = Image.open(anchor_tuple[0])
        positive = Image.open(positive_tuple[0])
        negative = Image.open(negative_tuple[0])

        anchor = anchor.convert("L")
        positive = positive.convert("L")
        negative = negative.convert("L")

        if self.should_invert:
            anchor = PIL.ImageOps.invert(anchor)
            positive = PIL.ImageOps.invert(positive)
            negative = PIL.ImageOps.invert(negative)

        if self.transform is not None:
            anchor = self.transform(anchor)
            positive = self.transform(positive)
            neagtive = self.transform(neagtive)

        return anchor, positive, negative
