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

from joblib import Parallel, delayed

ImageFile.LOAD_TRUNCATED_IMAGES = True

data_path = "/home/lab/bac/kaggle_data/rsna/"

class CropHead(object):

    def __init__(self, offset=5):

        """

        Originally made as a image transform, 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 __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

        """

        try:

            if type(img) != np.array:

                img_array = np.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 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 window_img(dcm, width=None, level=None, norm=True):

    try:

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

    except ValueError:

        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)

    if not width:

        width = dcm.WindowWidth

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

            width = width[0]

    if not level:

        level = dcm.WindowCenter

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

            level = level[0]

    lower = level - (width / 2)

    upper = level + (width / 2)

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

    if norm:

        return (img - lower) / (upper - lower)

    else:

        return img

def dcm_to_png(row, image_dirs, dataset, width, level, crop, crop_head, output_path):

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

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

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

    r = window_img(r_dcm, width, level)

    g = window_img(g_dcm, width, level)

    b = window_img(b_dcm, width, level)

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

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

    im = Image.fromarray(img)

    if crop:

        im = crop_head(im)

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

def prepare_png_images(dataset, folder_name, width=None, level=None, crop=True):

    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")

    }

    image_dirs = {

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

        "test_stage_1": os.path.join(data_path, "stage_1_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()

    Parallel(n_jobs=4)(delayed(dcm_to_png)(row, image_dirs, dataset, width, level, crop, crop_head, output_path) for _, row in tqdm(

        tqdm(triplets.iterrows(), total=len(triplets), desc=dataset)

    ))

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

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

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

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

    #     r = window_img(r_dcm, width, level)

    #     g = window_img(g_dcm, width, level)

    #     b = window_img(b_dcm, width, level)

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

    #     img = (img * 255).astype(np.uint8)

    #     im = Image.fromarray(img)

    #

    #     if crop:

    #         im = crop_head(im)

    #

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

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

if __name__ == '__main__':

    prepare_png_images("train", "adjacent-brain-cropped", 80, 40, crop=True)

    # prepare_png_images("test_stage_1", "adjacent-brain-cropped", 80, 40, crop=True)