'''

I/O functions for RSNA 2019

Updated: Yuval 12/10/19

'''

import os

from skimage.io import imread

import numpy as np

import torch

import pydicom

import SimpleITK as sitk

DEFAULT_IMAGES_BASE_PATH = '/media/nvme/data/RSNA/'

train_path='stage_1_train_images/'

test_path='stage_1_test_images/'

file_format='ID_{}.dcm'

def load_one_image(pid,

                   equalize=False,

                   base_path=DEFAULT_IMAGES_BASE_PATH,

                   file_path=train_path,

                   file_format=file_format,

                   window_eq=False,rescale=False):

    '''

    Load one image from a DICOM file 

        Args:

            pid             : string, image id

            equalize         : Equalize - return (image-mean)/std

            base_path       : Base File path for the images 

            file_path       : file path for the images (from base_path)

            file_format     : format string to get from pid to file name. default='ID_{}.dcm'

            window_eq=False : Do window equaliztion: (for backward competability, don't use it anymore use WSO)

                              False - No equalization

                              True  - Do equalization with window = [(40,80),(80,200),(600,2800)]

                              tuple/list shaped as above 

            rescale=False    : Use DICOM parameters for rescale, done automaticaly if windows_eq!=False

       Update:Yuval 12/10/19       

    '''

    dcm_data = pydicom.dcmread(base_path+ file_path+ file_format.format(pid))

    try:

        pixels = dcm_data.pixel_array

    except:

        try:

            read_image = sitk.ReadImage(base_path+ file_path+ file_format.format(pid))

            pixels = sitk.GetArrayFromImage(read_image)[0].astype(np.float32)

        except:

            print('got some error with pid {} {}'.format(pid,base_path+ file_path+ file_format.format(pid)))

            pixels = np.zeros((512,512))

    m = max(pixels.shape)

    p = np.ones((m,m),dtype=np.float)*pixels[0,0]

    p[(m-pixels.shape[0])//2:(m-pixels.shape[0])//2+pixels.shape[0],

      (m-pixels.shape[1])//2:(m-pixels.shape[1])//2+pixels.shape[1]]=pixels

    pixels = p

    if window_eq:

        if window_eq==3:

            pixels = dcm_window3(dcm_data,pixels)

            if equalize:

                for i in range(3):

                    pixels[i] = (pixels[i]-pixels[i].mean())/(pixels[i].std()+1e-6)

        else:

            pixels = dcm_window(dcm_data,pixels)

    else:

        if rescale:

            b = dcm_data.RescaleIntercept

            m = dcm_data.RescaleSlope

            pixels = m * pixels + b       

        if equalize:

            pixels = (pixels-pixels.mean())/(pixels.std()+1e-6)

    return pixels

def dcm_window(dcm,pixel_array,window = None):

    '''

    return a windowed image.

    if now window is provided, use the DICOM first window data

    if window length ==3, use window data

    '''

    ymin=0

    ymax=255

    if window is None:

        c = int(dcm.WindowCenter[0] if type(dcm.WindowCenter) == pydicom.multival.MultiValue else dcm.WindowCenter)

        w = int(dcm.WindowWidth[0] if type(dcm.WindowWidth) == pydicom.multival.MultiValue else dcm.WindowWidth)

    else:

        c , w = window

    b = dcm.RescaleIntercept

    m = dcm.RescaleSlope

    x = m * pixel_array + b

    # windowing C.11.2.1.2.1 Default LINEAR Function

    y = np.zeros_like(x)

    y[x <= (c - 0.5 - (w - 1) / 2)] = ymin

#        ymin + np.tanh(((x[x <= (c - 0.5 - (w - 1) / 2)] - (c - 0.5)) / (w - 1) + 0.5))* (ymax - ymin)  

    y[x > (c - 0.5 + (w - 1) / 2)] = ymax

#        ymax + np.tanh(((x[x > (c - 0.5 + (w - 1) / 2)] - (c - 0.5)) / (w - 1) + 0.5))* (ymax - ymin)   

    y[(x > (c - 0.5 - (w - 1) / 2)) & (x <= (c - 0.5 + (w - 1) / 2))] = \

      ((x[(x > (c - 0.5 - (w - 1) / 2)) & (x <= (c - 0.5 + (w - 1) / 2))] - (c - 0.5)) / (w - 1) + 0.5) * (ymax - ymin) + ymin

#    y = ((x - (c - 0.5)) / (w - 1) + 0.5) * (ymax - ymin) + ymin

    return y/255.0

def dcm_window3(dcm,pixel_array,windows = None):

    y = np.zeros((3,)+pixel_array.shape,dtype=np.float)

    if windows is None:

        windows=[(40,80),(80,200),(600,2800)]

    for i in range(3):

        y[i]=dcm_window(dcm,pixel_array,window = windows[i])

    return y