import numpy as np

import pandas as pd

from PIL import Image, ImageFile

from scipy import ndimage

import pydicom

import os

from tqdm import tqdm

from time import time

ImageFile.LOAD_TRUNCATED_IMAGES = True

data_path = "/mnt/storage_dimm2/kaggle_data/rsna-intracranial-hemorrhage-detection/"

def get_metadata(image_dir):

    labels = [

        'BitsAllocated', 'BitsStored', 'Columns', 'HighBit',

        'ImageOrientationPatient_0', 'ImageOrientationPatient_1', 'ImageOrientationPatient_2',

        'ImageOrientationPatient_3', 'ImageOrientationPatient_4', 'ImageOrientationPatient_5',

        'ImagePositionPatient_0', 'ImagePositionPatient_1', 'ImagePositionPatient_2',

        'Modality', 'PatientID', 'PhotometricInterpretation', 'PixelRepresentation',

        'PixelSpacing_0', 'PixelSpacing_1', 'RescaleIntercept', 'RescaleSlope', 'Rows', 'SOPInstanceUID',

        'SamplesPerPixel', 'SeriesInstanceUID', 'StudyID', 'StudyInstanceUID',

        'WindowCenter', 'WindowWidth', 'Image',

    ]

    data = {l: [] for l in labels}

    for image in tqdm(os.listdir(image_dir)):

        data["Image"].append(image[:-4])

        ds = pydicom.dcmread(os.path.join(image_dir, image))

        for metadata in ds.dir():

            if metadata != "PixelData":

                metadata_values = getattr(ds, metadata)

                if type(metadata_values) == pydicom.multival.MultiValue and metadata not in ["WindowCenter", "WindowWidth"]:

                    for i, v in enumerate(metadata_values):

                        data[f"{metadata}_{i}"].append(v)

                else:

                    if type(metadata_values) == pydicom.multival.MultiValue and metadata in ["WindowCenter", "WindowWidth"]:

                        data[metadata].append(metadata_values[0])

                    else:

                        data[metadata].append(metadata_values)

    return pd.DataFrame(data).set_index("Image")

def build_triplets(metadata):

    metadata.sort_values(by="ImagePositionPatient_2", inplace=True, ascending=False)

    studies = metadata.groupby("StudyInstanceUID")

    triplets = []

    for study_name, study_df in tqdm(studies):

        padded_names = np.pad(study_df.index, (1, 1), 'edge')

        for i, img in enumerate(padded_names[1:-1]):

            t = [padded_names[i], img, padded_names[i + 2]]

            triplets.append(t)

    return pd.DataFrame(triplets, columns=["red", "green", "blue"])

class CropHead(object):

    def __init__(self, offset=10):

        """

        Crops the head by labelling the objects in an image and keeping the second largest object (the largest object

        is the background). This method removes most of the headrest

        Originally made as a image transform for use with PyTorch, but too slow to run on the fly :(

        :param offset: Pixel offset to apply to the crop so that it isn't too tight

        """

        self.offset = offset

    def crop_extents(self, img):

        try:

            if type(img) != np.array:

                img_array = np.array(img)

            else:

                img_array = img

            labeled_blobs, number_of_blobs = ndimage.label(img_array)

            blob_sizes = np.bincount(labeled_blobs.flatten())

            head_blob = labeled_blobs == np.argmax(blob_sizes[1:]) + 1  # The number of the head blob

            head_blob = np.max(head_blob, axis=-1)

            mask = head_blob == 0

            rows = np.flatnonzero((~mask).sum(axis=1))

            cols = np.flatnonzero((~mask).sum(axis=0))

            x_min = max([rows.min() - self.offset, 0])

            x_max = min([rows.max() + self.offset + 1, img_array.shape[0]])

            y_min = max([cols.min() - self.offset, 0])

            y_max = min([cols.max() + self.offset + 1, img_array.shape[1]])

            return x_min, x_max, y_min, y_max

        except ValueError:

            return 0, 0, -1, -1

    def __call__(self, img):

        """

        Crops a CT image to so that as much black area is removed as possible

        :param img: PIL image

        :return: Cropped image

        """

        x_min, x_max, y_min, y_max = self.crop_extents(img)

        try:

            if type(img) != np.array:

                img_array = np.array(img)

            else:

                img_array = img

            return Image.fromarray(np.uint8(img_array[x_min:x_max, y_min:y_max]))

        except ValueError:

            return img

    def __repr__(self):

        return self.__class__.__name__ + '(offset={})'.format(self.offset)

def prepare_dicom(dcm, default_window=False):

    """

    Converts a DICOM object to a 16-bit Numpy array (in Housnfield units) or a uint8 image if the default window is used

    :param dcm: DICOM Object

    :param default_window: Flag to use the window settings specified in the metadata

    :return: Numpy array in either int16 or uint8

    """

    try:

        # https://www.kaggle.com/jhoward/cleaning-the-data-for-rapid-prototyping-fastai

        if dcm.BitsStored == 12 and dcm.PixelRepresentation == 0 and dcm.RescaleIntercept > -100:

            x = dcm.pixel_array + 1000

            px_mode = 4096

            x[x >= px_mode] = x[x >= px_mode] - px_mode

            dcm.PixelData = x.tobytes()

            dcm.RescaleIntercept = -1000

        pixels = dcm.pixel_array.astype(np.float32) * dcm.RescaleSlope + dcm.RescaleIntercept

    except ValueError as e:

        print("ValueError with", dcm.SOPInstanceUID, e)

        return np.zeros((512, 512))

    # Pad the image if it isn't square

    if pixels.shape[0] != pixels.shape[1]:

        (a, b) = pixels.shape

        if a > b:

            padding = ((0, 0), ((a - b) // 2, (a - b) // 2))

        else:

            padding = (((b - a) // 2, (b - a) // 2), (0, 0))

        pixels = np.pad(pixels, padding, mode='constant', constant_values=0)

    # Return image windows as per the metadata parameters

    if default_window:

        width = dcm.WindowWidth

        if type(width) != pydicom.valuerep.DSfloat:

            width = width[0]

        level = dcm.WindowCenter

        if type(level) != pydicom.valuerep.DSfloat:

            level = level[0]

        img_windowed = linear_windowing(pixels, level, width)

        return img_windowed

    # Return array Hounsfield units only

    else:

        return pixels.astype(np.int16)

def prepare_png(dataset, folder_name, channels=(0, 1, 2), crop=False):

    """

    Create PNG images using 3 specified window settings

    :param dataset: One of "train", "test_stage_1" or "test_stage_2"

    :param folder_name: Name of the output folder

    :param channels: Tuple to specifiy what windows to use for RGB channels

    :param crop: Flag to crop image to only the head

    :return:

    """

    start = time()

    image_dirs = {

        "train": os.path.join(data_path, "stage_1_train_images"),

        "test_stage_1": os.path.join(data_path, "stage_1_test_images"),

        "test_stage_2": os.path.join(data_path, "stage_2_test_images")

    }

    windows = [

        (None, None),  # No windowing

        (80, 40),  # Brain

        (200, 80),  # Subdural

        (40, 40),  # Stroke

        (2800, 600),  # Temporal bone

        (380, 40),  # Soft tissue

        (2000, 600),  # Bone

    ]

    output_path = os.path.join(data_path, "png", dataset, f"{folder_name}")

    crop_head = CropHead()

    if not os.path.exists(output_path):

        os.makedirs(output_path)

    for image_name in tqdm(os.listdir(image_dirs[dataset])):

        ds = pydicom.dcmread(os.path.join(image_dirs[dataset], image_name))

        rgb = []

        for c in channels:

            if c == 0:

                ch = prepare_dicom(ds, default_window=True)

            else:

                ch = prepare_dicom(ds)

                ch = linear_windowing(ch, windows[c][0], windows[c][1])

            rgb.append(ch)

        img = np.stack(rgb, -1)

        if crop:

            x_min, x_max, y_min, y_max = crop_head.crop_extents(img > 0)

            img = img[x_min:x_max, y_min:y_max]

            if img.shape[0] == 0 or img.shape[1] == 0:

                img = np.zeros(shape=(512, 512, 3), dtype=np.uint8)

        im = Image.fromarray(img.astype(np.uint8))

        im.save(os.path.join(output_path, image_name[:-4] + ".png"))

    print("Done in", (time() - start) // 60, "minutes")

def prepare_png_adjacent(dataset, folder_name, crop=True):

    """

    Prepare 3 channel adjacent images in Hounsfield units clipped between 0-255 HU

    The target image is the green channel. The reg and blue channels are spatially adjacent slices

    :param dataset: One of "train", "test_stage_1" or "test_stage_2"

    :param folder_name: Name of the output folder

    :param crop: Flag to crop image to only the head

    """

    start = time()

    triplet_dfs = {

        "train": os.path.join(data_path, "train_triplets.csv"),

        "test_stage_1": os.path.join(data_path, "stage_1_test_triplets.csv"),

        "test_stage_2": os.path.join(data_path, "stage_2_test_triplets.csv")

    }

    image_dirs = {

        "train": os.path.join(data_path, "stage_1_train_images"),

        "test_stage_1": os.path.join(data_path, "stage_1_test_images"),

        "test_stage_2": os.path.join(data_path, "stage_2_test_images")

    }

    output_path = os.path.join(data_path, "png", dataset, f"{folder_name}")

    if not os.path.exists(output_path):

        os.makedirs(output_path)

    triplets = pd.read_csv(triplet_dfs[dataset])

    crop_head = CropHead()

    for _, row in tqdm(triplets.iterrows(), total=len(triplets), desc=dataset):

        rgb = []

        for ch in ["red", "green", "blue"]:

            dcm = pydicom.dcmread(os.path.join(image_dirs[dataset], row[ch] + ".dcm"))

            rgb.append(prepare_dicom(dcm))

        img = np.stack(rgb, -1)

        img = np.clip(img, 0, 255).astype(np.uint8)

        if crop:

            x_min, x_max, y_min, y_max = crop_head.crop_extents(img > 0)

            img = img[x_min:x_max, y_min:y_max]

            if img.shape[0] == 0 or img.shape[1] == 0:

                img = np.zeros(shape=(512, 512, 3), dtype=np.uint8)

        im = Image.fromarray(img)

        im.save(os.path.join(output_path, row["green"] + ".png"))

    print("Done in", (time() - start) // 60, "minutes")

def dicom_to_npy(dataset, folder_name):

    """

    Saves DICOM images as 16-bit Numpy arrays

    :param dataset: One of "train", "test_stage_1" or "test_stage_2"

    :param folder_name: Name of the output folder

    """

    image_dirs = {

        "train": os.path.join(data_path, "stage_1_train_images"),

        "test_stage_1": os.path.join(data_path, "stage_1_test_images"),

        "test_stage_2": os.path.join(data_path, "stage_2_test_images")

    }

    output_path = os.path.join(data_path, "npy", dataset, f"{folder_name}")

    print("Saving slices to", output_path)

    if not os.path.exists(output_path):

        os.makedirs(output_path)

    for image_name in tqdm(os.listdir(image_dirs[dataset])):

        dcm = pydicom.dcmread(os.path.join(image_dirs[dataset], image_name))

        np.save(os.path.join(output_path, image_name[:-4]), prepare_dicom(dcm))

def prepare_npy_adjacent(dataset, folder_name, crop=True):

    """

    Prepare 3 channel adjacent images in Hounsfield units (unclipped)

    :param dataset: One of "train", "test_stage_1" or "test_stage_2"

    :param folder_name: Name of the output folder

    :param crop: Flag to crop image to only the head

    """

    start = time()

    triplet_dfs = {

        "train": os.path.join(data_path, "train_triplets.csv"),

        "test_stage_1": os.path.join(data_path, "stage_1_test_triplets.csv"),

        "test_stage_2": os.path.join(data_path, "stage_2_test_triplets.csv")

    }

    image_dirs = {

        "train": os.path.join(data_path, "npy", dataset, "single_hu_slices"),

        "test_stage_1": os.path.join(data_path, "npy", dataset, "single_hu_slices"),

        "test_stage_2": os.path.join(data_path, "npy", dataset, "single_hu_slices")

    }

    output_path = os.path.join(data_path, "npy", dataset, f"{folder_name}")

    if not os.path.exists(output_path):

        os.makedirs(output_path)

    triplets = pd.read_csv(triplet_dfs[dataset])

    crop_head = CropHead()

    for _, row in tqdm(triplets.iterrows(), total=len(triplets), desc=dataset):

        r = np.load(os.path.join(image_dirs[dataset], row["red"] + ".npy"))

        g = np.load(os.path.join(image_dirs[dataset], row["green"] + ".npy"))

        b = np.load(os.path.join(image_dirs[dataset], row["blue"] + ".npy"))

        img = np.stack([r, g, b], -1)

        if crop:

            x_min, x_max, y_min, y_max = crop_head.crop_extents(img > 0)

            img = img[x_min:x_max, y_min:y_max]

            if img.shape[0] == 0 or img.shape[1] == 0:

                img = np.zeros(shape=(512, 512, 3), dtype=np.int16)

        np.save(os.path.join(output_path, row["green"]), img)

    print("Done in", (time() - start) // 60, "minutes")

def sigmoid(x):

    return 1 / (1 + np.exp(-x))

def linear_windowing(img, window_width, window_length):

    """

    Applies a linear window on an array

    :param img: Image array (in Hounsfield units)

    :param window_width:

    :param window_length:

    :return:

    """

    if window_width and window_length:

        lower = window_length - (window_width / 2)

        upper = window_length + (window_width / 2)

        img = np.clip(img, lower, upper)

        img = (img - lower) / (upper - lower)

        return (img*255).astype(np.uint8)

    else:

        return img

def sigmoid_windowing(img, window_width, window_length, u=255, epsilon=255):

    """

    Applies a sigmoid window on an array

    From Practical Window Setting Optimization for Medical Image Deep Learning https://arxiv.org/pdf/1812.00572.pdf

    :param img: Image array (in Hounsfield units)

    :param window_width:

    :param window_length:

    :param u:

    :param epsilon:

    :return:

    """

    if window_width and window_length:

        weight = (2 / window_width) * np.log((u / epsilon) - 1)

        bias = (-2 * window_length / window_width) * np.log((u / epsilon) - 1)

        img = u * sigmoid(weight * img + bias)

        return img.astype(np.uint8)

    else:

        return img

if __name__ == '__main__':

    # Generate metadata dataframes

    train_metadata = get_metadata(os.path.join(data_path, "stage_1_train_images"))

    test_metadata = get_metadata(os.path.join(data_path, "stage_1_test_images"))

    train_metadata.to_parquet(f'{data_path}/train_metadata.parquet.gzip', compression='gzip')

    test_metadata.to_parquet(f'{data_path}/stage_1_test_metadata.parquet.gzip', compression='gzip')

    # Build triplets of adjacent images

    train_triplets = build_triplets(train_metadata)

    test_triplets = build_triplets(test_metadata)

    train_triplets.to_csv(os.path.join(data_path, "train_triplets.csv"))

    test_triplets.to_csv(os.path.join(data_path, "stage_1_test_triplets.csv"))

    # Prepare adjacent images

    prepare_png_adjacent("train", "adjacent_hu_cropped")

    prepare_png_adjacent("test_stage_1", "adjacent_hu_cropped")

    # Prepare 3 window images (brain-subdural-bone)

    prepare_png("train", "brain-subdural-bone", channels=(1, 2, 6), crop=True)

    prepare_png("test_stage_1", "brain-subdural-bone", channels=(1, 2, 6), crop=True)

    # Stage 2 preparations

    test_metadata = get_metadata(os.path.join(data_path, "stage_2_test_images"))

    test_metadata.to_parquet(f'{data_path}/stage_2_test_metadata.parquet.gzip', compression='gzip')

    test_triplets = build_triplets(test_metadata)

    test_triplets.to_csv(os.path.join(data_path, "stage_2_test_triplets.csv"))

    prepare_png_adjacent("test_stage_2", "adjacent_hu_cropped")

    prepare_png("test_stage_2", "brain-subdural-bone", channels=(1, 2, 6), crop=True)