import logging
from pathlib import Path
from typing import Callable, Optional

import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import Subset
from torchvision.datasets import VisionDataset

from base.torchvision_dataset import TorchvisionDataset

from .preprocessing import create_semisupervised_setting


class EsmeraSplit_Dataset(TorchvisionDataset):
    def __init__(
        self,
        root: str,
        ratio_known_normal: float = 0.0,
        ratio_known_outlier: float = 0.0,
        ratio_pollution: float = 0.0,
        inference: bool = False,
    ):
        super().__init__(root)

        # Define normal and outlier classes
        self.n_classes = 2  # 0: normal, 1: outlier
        self.normal_classes = tuple([0])
        self.outlier_classes = tuple([1])
        self.inference_set = None

        # MNIST preprocessing: feature scaling to [0, 1]
        # FIXME understand mnist feature scaling and check if it or other preprocessing is necessary for elpv
        transform = transforms.ToTensor()
        target_transform = transforms.Lambda(lambda x: int(x in self.outlier_classes))

        if inference:
            self.inference_set = EsmeraSplitInference(
                root=self.root,
                transform=transform,
            )
        else:
            # Get train set
            train_set = EsmeraSplitTraining(
                root=self.root,
                transform=transform,
                target_transform=target_transform,
                train=True,
            )

            # Create semi-supervised setting
            idx, _, semi_targets = create_semisupervised_setting(
                train_set.targets.cpu().data.numpy(),
                self.normal_classes,
                self.outlier_classes,
                self.outlier_classes,
                ratio_known_normal,
                ratio_known_outlier,
                ratio_pollution,
            )
            train_set.semi_targets[idx] = torch.tensor(
                np.array(semi_targets, dtype=np.int8)
            )  # set respective semi-supervised labels

            # Subset train_set to semi-supervised setup
            self.train_set = Subset(train_set, idx)

            # Get test set
            self.test_set = EsmeraSplitTraining(
                root=self.root,
                train=False,
                transform=transform,
                target_transform=target_transform,
            )


def split_array_into_subarrays(array, split_height, split_width):
    original_shape = array.shape
    height, width = original_shape[-2], original_shape[-1]
    assert height % split_height == 0, "The height is not divisible by the split_height"
    assert width % split_width == 0, "The width is not divisible by the split_width"
    num_splits_height = height // split_height
    num_splits_width = width // split_width
    reshaped_array = array.reshape(
        -1, num_splits_height, split_height, num_splits_width, split_width
    )
    transposed_array = reshaped_array.transpose(0, 1, 3, 2, 4)
    final_array = transposed_array.reshape(-1, split_height, split_width)
    return final_array


class EsmeraSplitTraining(VisionDataset):
    def __init__(
        self,
        root: str,
        transforms: Optional[Callable] = None,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        train=False,
        split=0.7,
        seed=0,
        height=16,
        width=256,
    ):
        super(EsmeraSplitTraining, self).__init__(
            root, transforms, transform, target_transform
        )

        experiments_data = []
        experiments_targets = []
        validation_files = []
        experiment_files = []

        logger = logging.getLogger()

        for experiment_file in Path(root).iterdir():
            if experiment_file.is_dir() and experiment_file.name == "validation":
                for validation_file in experiment_file.iterdir():
                    if validation_file.suffix != ".npy":
                        continue
                    validation_files.append(experiment_file)
            if experiment_file.suffix != ".npy":
                continue
            experiment_files.append(experiment_file)
            experiment_data = np.load(experiment_file)

            if (
                experiment_data.shape[1] % height != 0
                or experiment_data.shape[2] % width != 0
            ):
                logger.error(
                    f"Experiment {experiment_file.name} has shape {experiment_data.shape} which is not divisible by {height}x{width}"
                )
            experiment_data = split_array_into_subarrays(experiment_data, height, width)

            # experiment_data = np.lib.format.open_memmap(experiment_file, mode='r+')
            experiment_targets = (
                np.ones(experiment_data.shape[0], dtype=np.int8)
                if "smoke" in experiment_file.name
                else np.zeros(experiment_data.shape[0], dtype=np.int8)
            )
            experiments_data.append(experiment_data)
            experiments_targets.append(experiment_targets)

        filtered_validation_files = []
        for validation_file in validation_files:
            validation_file_name = validation_file.name
            file_exists_in_experiments = any(
                experiment_file.name == validation_file_name
                for experiment_file in experiment_files
            )
            if not file_exists_in_experiments:
                filtered_validation_files.append(validation_file)
        validation_files = filtered_validation_files

        logger.info(
            f"Train/Test experiments: {[experiment_file.name for experiment_file in experiment_files]}"
        )
        logger.info(
            f"Validation experiments: {[validation_file.name for validation_file in validation_files]}"
        )

        lidar_projections = np.concatenate(experiments_data)
        smoke_presence = np.concatenate(experiments_targets)

        np.random.seed(seed)

        shuffled_indices = np.random.permutation(lidar_projections.shape[0])
        shuffled_lidar_projections = lidar_projections[shuffled_indices]
        shuffled_smoke_presence = smoke_presence[shuffled_indices]

        split_idx = int(split * shuffled_lidar_projections.shape[0])

        if train:
            self.data = shuffled_lidar_projections[:split_idx]
            self.targets = shuffled_smoke_presence[:split_idx]

        else:
            self.data = shuffled_lidar_projections[split_idx:]
            self.targets = shuffled_smoke_presence[split_idx:]

        self.data = np.nan_to_num(self.data)

        self.data = torch.tensor(self.data)
        self.targets = torch.tensor(self.targets, dtype=torch.int8)

        self.semi_targets = torch.zeros_like(self.targets, dtype=torch.int8)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        """Override the original method of the MNIST class.
        Args:
            index (int): Index

        Returns:
            tuple: (image, target, semi_target, index)
        """
        img, target, semi_target = (
            self.data[index],
            int(self.targets[index]),
            int(self.semi_targets[index]),
        )

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode="F")

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target, semi_target, index


class EsmeraSplitInference(VisionDataset):
    def __init__(
        self,
        root: str,
        transforms: Optional[Callable] = None,
        transform: Optional[Callable] = None,
    ):
        super(EsmeraSplitInference, self).__init__(root, transforms, transform)
        logger = logging.getLogger()

        self.experiment_file_path = Path(root)

        if not self.experiment_file_path.is_file():
            logger.error(
                "For inference the data path has to be a single experiment file!"
            )
            raise Exception("Inference data is not a loadable file!")

        self.data = np.load(self.experiment_file_path)
        self.data = split_array_into_subarrays(self.data, 16, 256)
        self.data = np.nan_to_num(self.data)
        self.data = torch.tensor(self.data)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        """Override the original method of the MNIST class.
        Args:
            index (int): Index

        Returns:
            tuple: (image, index)
        """
        img = self.data[index]

        # doing this so that it is consistent with all other datasets
        # to return a PIL Image
        img = Image.fromarray(img.numpy(), mode="F")

        if self.transform is not None:
            img = self.transform(img)

        return img, index
