# -*- coding: utf-8 -*-

import os
import math
import random
from collections import Counter

from PIL import Image, ImageEnhance

import torch
from torchvision.datasets import ImageFolder
from datasets import DatasetDict

from tabulate import tabulate

class VisionDataset:
    
    @staticmethod
    def datasetFromDirectory(path: str, test_ratio = 0.15, balanced = True, augmentation = False):
        """
        Load data from one directory
        ↪ Directory
            ↪ sars_cov_2
            ↪ normal
            ↪ pneumonia
        """

        return VisionDataset.__splitDatasets(
            ImageFolder(path),
            test_ratio,
            balanced,
            augmentation,
        )
    
    @staticmethod
    def datasetFromDirectories(train: str, test: str, balanced = True):
        """
        Args:
            train (:obj:`str`): Path to the training dataset folder
            test  (:obj:`str`): Path to the test dataset folder
        Returns:
            :class:`DatasetDict`: The dataset containing Train and Test
        Example:
            ↪ Train/
                ↪ sars_cov_2/
                ↪ normal/
                ↪ pneumonia/
            ↪ Test/
                ↪ sars_cov_2/
                ↪ normal/
                ↪ pneumonia/
        Information:
            We build the dataset accordingly to the directory tree.
            So, each sub-directories will be accounted like a class.
        """

        return DatasetDict({
            "train": VisionDataset.__balance(
                ImageFolder(train),
                balanced
            ),
            "test": ImageFolder(test),
        })

    @staticmethod
    def getConfig(dataset: torch.utils.data.Dataset):

        labels2ids = {}
        ids2labels = {}

        for i, class_name in enumerate(dataset.features['label'].names):
            labels2ids[class_name] = str(i)
            ids2labels[str(i)] = class_name

        return labels2ids, ids2labels

    """
    🧬 Apply data augmentation on the input image
    Source: https://medium.com/lunit/photometric-data-augmentation-in-projection-radiography-bed3ae9f55c3
    """
    @staticmethod
    def __augmentation(image, beta=0.33):

        # Random augmentation in X % of cases
        if random.randint(0,100) < int(beta*100):
        
            # Random Contrast
            im3 = ImageEnhance.Contrast(image)
            im3.enhance(random.uniform(0.5, 1.0)).show()
        
            # Random Noise
            # Not implemented yet!

        return image

    """
    ⚖️ Balance the dataset according to the less represented class
    """
    @staticmethod
    def __balance(train_ds, balanced):
        
        # If balanced isn't enabled
        if balanced == False:
            return train_ds

        ct_train = Counter([label for _, label in train_ds])
        train_classes = [ct_train[a] for a in sorted(ct_train)]

        # Get the less represented label in train
        less_represented_train = min(train_classes)

        labels = {}

        # For each image
        for img, label in train_ds:

            # Create the label array if isn't
            if label not in labels:
                labels[label] = []

            # Add the image to the array
            labels[label].append(img)

        # New dataset
        balanced_ds = []

        # For each label
        for label in labels:

            # Get images
            imgs = labels[label]

            # For each image
            for img in imgs[0:less_represented_train]:

                # Create a tuple: image and label 
                t = (img, label)

                # Add it
                balanced_ds.append(t)

        print("The less represented label in train as " + str(less_represented_train) + " occurrences")
        print("Size of train after balancing is " + str(len(balanced_ds)))

        return torch.utils.data.Subset(
            balanced_ds,
            range(0,len(balanced_ds))
        ).dataset

    """
    ✂️ Split the dataset into sub-datasets
    """
    @staticmethod
    def __splitDatasets(dataset: torch.utils.data.Dataset, test_ratio = 0.15, balanced = True, augmentation = False):

        print("Split Datasets...")

        label2id, id2label = VisionDataset.getConfig(dataset)

        # Train Ratio
        train_ratio = 1 - test_ratio
        
        # Generate index from 0 to N where N is the len of the dataset
        indices = torch.randperm(len(dataset)).tolist()

        # Index of the validation corpora
        elements_test_dev = math.floor(len(indices) * .15)

        # Get end indexes
        train_index = int(len(dataset) * train_ratio)

        # TRAIN
        train_ds = torch.utils.data.Subset(dataset, indices[0:train_index])
        ct_train = Counter([label for _, label in train_ds])
        train_classes = [ct_train[a] for a in sorted(ct_train)]
        
        # If balanced is enabled 
        train_ds = VisionDataset.__balance(train_ds, balanced)
        
        # Compute again the stats
        ct_train = Counter([label for _, label in train_ds])
        train_classes = [ct_train[a] for a in sorted(ct_train)]

        # If data augmentation is enabled
        if augmentation == True:

            new_ds = []

            # For each annotated image
            for img, label in train_ds:

                # Augment it
                new_ds.append((VisionDataset.__augmentation(img), label))
            
            # Replace by the augmented data
            train_ds = torch.utils.data.Subset(
                new_ds,
                range(0,len(new_ds))
            ).dataset

        # TEST
        test_ds = torch.utils.data.Subset(dataset, indices[train_index:])
        ct_test = Counter([label for _, label in test_ds])
        test_classes = [ct_test[a] for a in sorted(ct_test)]

        # Make a table with all the information about the repartition of the dataset
        table_repartition = [
            ["Train"] + train_classes + [str(len(train_ds))],
            ["Test"] + test_classes + [str(len(test_ds))],
        ]
        repartitions_table = tabulate(table_repartition, ["Dataset"] + list(id2label.values()) + ["Total"], tablefmt="pretty")
        print(repartitions_table)

        return DatasetDict({
            "train": train_ds,
            "test": test_ds.dataset,
        })
