import numpy as np
import os
import cv2
import pandas as pd
from torch.utils.data import Dataset
# import jpeg4py as jpeg
from utils import get_windowing, window_image
import pydicom

IGNORE_IDS = [
    'ID_6431af929',
]

windows_range = {
    'brain': [40, 80],
    'bone': [600, 2800],
    'subdual': [75, 215]
}

LABEL_COLS = ["epidural", "intraparenchymal", "intraventricular", "subarachnoid", "subdural", "any"]
LABEL_COLS_WITHOUT_ANY = ["epidural", "intraparenchymal", "intraventricular", "subarachnoid", "subdural"]


def load_dicom_image(path):
    data = pydicom.read_file(path)
    image = data.pixel_array
    window_center, window_width, intercept, slope = get_windowing(data)
    images = []
    image_windowed = window_image(image, window_center, window_width, intercept, slope)
    images.append(image_windowed)

    for k, v in windows_range.items():
        image_windowed = window_image(image, v[0], v[1], intercept, slope)
        images.append(image_windowed)

    images = np.asarray(images).transpose((1, 2, 0))
    images = images / 255
    return images


def load_image(path):
    image = cv2.imread(path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image


def load_random_windows(path, id):
    random_window = np.random.choice(['brain', 'bone', 'subdual'], 1)[0]
    return load_image(os.path.join(path, random_window, id + ".jpg"))


def load_multi_images(root, image_name):
    images = []
    for i, (k, v) in enumerate(windows_range.items()):
        image = cv2.imread(os.path.join(root, k, image_name), 0)
        images.append(image)

    images = np.asarray(images).transpose((1, 2, 0))

    return images


# def load_jpeg_image(path):
#     image = jpeg.JPEG(path).decode()
#     return image


import random
def get_balance_set(df):
    patients = set(df["patient_id"].unique())
    patients_pos = set(df[df["any"] == 1]["patient_id"].unique())
    patients_neg = patients - patients_pos
    patients_neg_balance = random.sample(patients_neg, len(patients_pos))
    patients_balance = patients_pos.union(patients_neg_balance)

    print(len(patients), len(patients_pos), len(patients), len(patients_balance))

    return df[df["patient_id"].isin(patients_balance)]


from sklearn.preprocessing import MinMaxScaler
meta_data_cols = [
    'image_position_patient_0', 'image_position_patient_1', 'image_position_patient_2',
    'image_orientation_patient_0', 'image_orientation_patient_2', 'image_orientation_patient_3',
    'image_orientation_patient_4', 'image_orientation_patient_5'
]


class RSNADataset(Dataset):
    """
    Read JPG images
    """
    def __init__(self, csv_file, root, with_any, transform, mode='train', image_type='jpg'):
        if isinstance(csv_file, pd.DataFrame):
            df = csv_file
        else:
            print(csv_file)
            df = pd.read_csv(csv_file)
        if mode == 'train':
            # df = df
            df = get_balance_set(df)
        if mode in ['train', 'valid']:
            meta_data = pd.read_csv(f"/data/df_dicom_metadata_train.csv", usecols=meta_data_cols + ['sop_instance_uid'])
        else:
            meta_data = pd.read_csv(f"/data/df_dicom_metadata_test.csv", usecols=meta_data_cols + ['sop_instance_uid'])
            df["sop_instance_uid"] = "ID_" + df["sop_instance_uid"]
        meta_data = meta_data[meta_data['sop_instance_uid'].isin(df['sop_instance_uid'])]
        df = df.merge(meta_data, on='sop_instance_uid', how='left')
        ID_col = "Image" if "Image" in df.columns else "ID" if "ID" in df.columns else "sop_instance_uid"
        df = df[~df[ID_col].isin(IGNORE_IDS)]
        self.ids = df[ID_col].values
        self.metadata = df[meta_data_cols].values
        self.with_any = with_any
        if with_any:
            self.labels = df[LABEL_COLS].values
        else:
            self.labels = df[LABEL_COLS_WITHOUT_ANY].values
        self.root = root
        self.transform = transform
        self.image_type = image_type

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

    def __getitem__(self, idx):
        id = self.ids[idx]
        label = self.labels[idx].astype(np.float32)

        meta = self.metadata[idx].astype(np.float32)

        if not "ID" in id:
            id = "ID_" + id

        image = os.path.join(self.root, id + "." + self.image_type)
        image = load_image(image)

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        image = np.transpose(image, (2, 0, 1)).astype(np.float32)

        return {
            'images': image,
            'targets': label,
            'meta': meta
        }


class RSNARandomWindowDataset(RSNADataset):
    """
    Random select bone, brain and subdual during the training
    """

    def __getitem__(self, idx):
        id = self.ids[idx]
        label = self.labels[idx].astype(np.float32)

        image = load_random_windows(self.root, id)

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        image = np.transpose(image, (2, 0, 1)).astype(np.float32)

        return {
            'images': image,
            'targets': label
        }


class RSNADicomDataset(RSNADataset):
    """
    load dicom image directly. windows are applied on the fly.
    """
    def __init__(self, csv_file, root, with_any, transform, mode='train'):
        super(RSNADicomDataset, self).__init__(csv_file, root, with_any, transform, mode)

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

    def __getitem__(self, idx):
        id = self.ids[idx]
        label = self.labels[idx].astype(np.float32)

        image = os.path.join(self.root, id + ".dcm")
        image = load_dicom_image(image)

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        image = np.transpose(image, (2, 0, 1)).astype(np.float32)

        return {
            'images': image,
            'targets': label
        }


class RSNAMultiWindowsDataset(Dataset):
    """
    Read all window images then concatinate.
    """
    def __init__(self, csv_file, root, with_any, transform):
        if isinstance(csv_file, pd.DataFrame):
            df = csv_file
        else:
            df = pd.read_csv(csv_file)
        ID_col = "Image" if "Image" in df.columns else "ID" if "ID" in df.columns else "sop_instance_uid"
        df = df[~df[ID_col].isin(IGNORE_IDS)]
        self.ids = df[ID_col].values
        self.with_any = with_any
        if with_any:
            self.labels = df[LABEL_COLS].values
        else:
            self.labels = df[LABEL_COLS_WITHOUT_ANY].values
        self.root = root
        self.transform = transform

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

    def __getitem__(self, idx):
        id = self.ids[idx]
        label = self.labels[idx].astype(np.float32)

        # image = os.path.join(self.root, id + ".jpg")
        image = load_multi_images(self.root, id + ".jpg")

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        image = np.transpose(image, (2, 0, 1)).astype(np.float32)

        return {
            'images': image,
            'targets': label
        }