import numpy as np
import pandas as pd
from PIL import Image, ImageFile
import os
import torch
from torch.utils.data import Dataset, Sampler
import sys
sys.path.append("/home/anjum/PycharmProjects/kaggle")
# sys.path.append("/home/anjum/rsna_code")  # GCP
from rsna_intracranial_hemorrhage_detection.data_prep import linear_windowing, sigmoid_windowing
ImageFile.LOAD_TRUNCATED_IMAGES = True


data_path = "/mnt/storage_dimm2/kaggle_data/rsna-intracranial-hemorrhage-detection/"
# data_path = "/home/anjum/rsna_data/"  # GCP


class ICHDataset(Dataset):
    def __init__(self, dataset, phase=1, image_filter=None, transforms=None, image_folder=None, png=True):
        df_paths = {
            "train": os.path.join(data_path, "stage_1_train.csv"),
            "test1": os.path.join(data_path, "stage_1_sample_submission.csv"),
            "test2": os.path.join(data_path, "stage_2_sample_submission.csv")
        }

        self.png = png

        if self.png:
            image_dirs = {
                "train": os.path.join(data_path, "png", "train", image_folder),
                "test1": os.path.join(data_path, "png", "test_stage_1", image_folder),
                "test2": os.path.join(data_path, "png", "test_stage_2", image_folder)
            }
        else:
            image_dirs = {
                "train": os.path.join(data_path, "npy", "train", image_folder),
                "test1": os.path.join(data_path, "npy", "test_stage_1", image_folder),
                "test2": os.path.join(data_path, "npy", "test_stage_2", image_folder)
            }

        self.dataset = dataset
        self.phase = phase
        self.transforms = transforms
        self.image_dir = image_dirs[dataset]

        self.df = pd.read_csv(df_paths[dataset]).drop_duplicates()
        self.df['ImageID'] = self.df['ID'].str.slice(stop=12)
        self.df['Diagnosis'] = self.df['ID'].str.slice(start=13)
        self.df_pivot = self.df.pivot(index="ImageID", columns="Diagnosis", values="Label")

        if image_filter is not None:
            self.df_pivot = self.df_pivot.loc[image_filter]

        if self.phase == 0:
            self.labels = self.df_pivot.values
        elif self.phase == 1:
            self.labels = self.df_pivot["any"].values.reshape(-1, 1)
        else:
            self.labels = self.df_pivot[["epidural", "intraparenchymal",
                                         "intraventricular", "subarachnoid", "subdural"]].values

        self.image_ids = self.df_pivot.index.values
        self.class_weights = np.mean(self.labels, axis=0)

    def load_image(self, image_name):
        # window_width, window_length = 80, 40  # Brain window
        window_width, window_length = 200, 80  # Subdural window
        # window_width, window_length = 130, 50  # Subdural window

        if self.png:
            img = np.array(Image.open(os.path.join(self.image_dir, image_name+".png")).convert("RGB"))
            return linear_windowing(img, window_width, window_length)
            # return sigmoid_windowing(img, window_width, window_length)
        else:
            img = np.load(os.path.join(self.image_dir, image_name+".npy"))

            # PIL doesn't work with 16-bit RGB images :(
            # Should be ok though since the useful HU interval is between 0-255
            if img.shape[0] == 0 or img.shape[1] == 0:
                return np.zeros(shape=(512, 512, 3), dtype=np.uint8)
            else:
                # return np.clip(img, 0, 255).astype(np.uint8)  # Use this with the Windowing module
                return linear_windowing(img, window_width, window_length)
                # return sigmoid_windowing(img, window_width, window_length)

    def __getitem__(self, idx):
        img_id = self.image_ids[idx]
        img = self.load_image(img_id)

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

        if self.dataset == "train":
            return img, torch.tensor(self.labels[idx], dtype=torch.float32)
        else:
            return img, torch.tensor([0], dtype=torch.float32)

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


class BalancedRandomSampler(Sampler):
    def __init__(self, data_source):
        """
        Balances the negative and positive samples. All of the positive samples are used, but a random subset of
        the negative samples are used to create a 50:50 dataset
        :param data_source: An ICHDataset
        """
        super().__init__(data_source)
        self.labels = data_source.labels
        self.ids_pos = np.where(self.labels[:, 0] == 1)[0]
        self.ids_neg = np.where(self.labels[:, 0] == 0)[0]

    def __iter__(self):
        ids_neg_sampled = np.random.choice(self.ids_neg, self.ids_pos.shape[0], replace=False)
        ids = np.concatenate([self.ids_pos, ids_neg_sampled])
        np.random.shuffle(ids)
        return iter(ids)

    def __len__(self):
        return self.ids_pos.shape[0] * 2