import json
import os
from typing import List, Tuple

import numpy as np
import pandas as pd
import torch
from PIL import Image
from loguru import logger
from sklearn.utils import compute_class_weight
from torch import Tensor
from torch.utils.data import DataLoader
from torchvision.transforms import Compose

from lib.datasets import DatasetSplit
from lib.utils import glob_dir


class LabeledDataset(torch.utils.data.Dataset):
    """
    The LabeledDataset wraps a folder structure of a dataset with train, val, and test folders
    and the corresponding labels.csv with the meta information.
    """

    def __init__(self, data_dir: str, split: DatasetSplit, transforms: Compose):
        """
        Create a new LabeledDataset.
        @param data_dir: The root directory of the dataset.
        @param split: Which dataset split should be considered.
        @param transforms: The transformations that are applied to the images before returning.
        """
        data_dir = os.path.join(data_dir, split.value['folder'])
        logger.debug(f'Setup LabeledDataset on path: {data_dir}')
        # Load meta-data from labels.csv and filter by the split
        df = pd.read_csv(os.path.join(os.path.dirname(data_dir), 'labels.csv'))
        self.meta_data = df[df['split'] == split.value['value']]
        # Search for all files available in the filesystem
        files_found = glob_dir(data_dir)
        file_names_found = [os.path.basename(f).split('-')[-1] for f in files_found]
        logger.debug(f'Files found: {len(file_names_found)}')
        # Extract base name from image_path in meta-data
        self.meta_data['image_path'] = self.meta_data['image_path'].apply(lambda p: p.split('\\')[-1].split('/')[-1])
        # Filter for all available images in files_found
        self.meta_data = self.meta_data[self.meta_data['image_path'].isin(file_names_found)]
        logger.debug(f'Files not available: {len(file_names_found) - len(self.meta_data)}')
        self.meta_data['image_path'] = self.meta_data['image_path'].apply(
            lambda p: files_found[file_names_found.index(p)] if p in file_names_found else p)
        # Extract labels from meta-data
        self.labels = list(self.meta_data.drop(['image_path', 'split'], axis=1).columns.values)[1:]
        self.transforms = transforms
        logger.info(f'Finished "{os.path.basename(data_dir)}" Dataset (total={len(self.meta_data)}) Setup!')

    def __len__(self) -> int:
        return len(self.meta_data)

    def __getitem__(self, idx) -> Tuple[Tensor, np.ndarray]:
        df = self.meta_data.iloc[idx]
        img = Image.open(df["image_path"])
        img = self.transforms(img)
        labels = df[self.labels].astype(float).to_numpy()
        return img, labels

    @property
    def classes(self) -> List[str]:
        """
        @return: The class labels as list.
        """
        return self.labels

    @property
    def num_classes(self) -> int:
        """
        @return: The number of classes.
        """
        return len(self.classes)

    @property
    def is_multi_label(self) -> bool:
        """
        @return: If the dataset represents a multiple labels at once.
        """
        return (self.meta_data[self.labels].sum(axis=1) > 1).any()

    @property
    def class_weight(self) -> Tensor:
        """
        Compute the class weights for the dataset.
        @return: The class weights as tensor.
        """
        logger.debug(f'Compute class weights with {self.num_classes} classes: {", ".join(self.classes)}')
        logger.debug(f'Sum of samples per class:\n{self.meta_data[self.labels].sum(axis=0)}')
        if self.is_multi_label:
            y = self.meta_data[self.labels].to_numpy()
            class_weight = np.empty([self.num_classes, 2])
            for idx in range(self.num_classes):
                class_weight[idx] = compute_class_weight(class_weight='balanced', classes=[0., 1.], y=y[:, idx])
            class_weight = np.mean(class_weight, axis=1)
        else:
            y = np.argmax(self.meta_data[self.labels].to_numpy(), axis=1)
            class_weight = compute_class_weight(class_weight='balanced', classes=range(self.num_classes), y=y)
        class_weights_per_label = {self.labels[idx]: class_weight[idx] for idx in range(len(self.labels))}
        logger.debug(f'Class-Weights for each Label={json.dumps(class_weights_per_label, indent=3)}')
        return torch.from_numpy(class_weight)
