import os, glob

import nibabel as nib

import numpy as np

import cv2

from bayesnoiseremoval import BayesPreprocessor as bp

# STEP 1 - Load and visualize data

imgPath = '../../mr_train/'

maskPath = '../../train_labels/'

imagePathInput = os.path.join(imgPath)

maskPathInput = os.path.join(maskPath)

imgOutput = './zdim/imagesbayes/'

maskOutput = './zdim/masks/'

imageSliceOutput = os.path.join(imgOutput)

maskSliceOutput = os.path.join(maskOutput)

# STEP 2 - Image normalization

HOUNSFIELD_MIN = -1000

HOUNSFIELD_MAX = 2000

HOUNSFIELD_RANGE = HOUNSFIELD_MAX - HOUNSFIELD_MIN

# STEP 3 - Slicing and saving

SLICE_X = False

SLICE_Y = False

SLICE_Z = True

SLICE_DECIMATE_IDENTIFIER = 3

# Normalize image

def normalizeImageIntensityRange(img):

    img[img < HOUNSFIELD_MIN] = HOUNSFIELD_MIN

    img[img > HOUNSFIELD_MAX] = HOUNSFIELD_MAX

    return (img - HOUNSFIELD_MIN) / HOUNSFIELD_RANGE

def readImageVolume(imgPath, normalize=False):

    img = nib.load(imgPath).get_fdata()

    if normalize:

        return normalizeImageIntensityRange(img)

    else:

        return img

# Save volume slice to file

def saveSlice(img, fname, path):

    # img = np.uint8(img * 255)

    fout = os.path.join(path, f'{fname}.png')

    cv2.imwrite(fout, img)

    print(f'[+] Slice saved: {fout}', end='\r')

def sliceAndSaveVolumeImage(vol, fname, path):

    (dimx, dimy, dimz) = vol.shape

    print(dimx, dimy, dimz)

    cnt = 0

    if SLICE_X:

        cnt += dimx

        print('Slicing X: ')

        for i in range(dimx):

            saveSlice(bp.bayes_noise_removal(image=vol[i, :, :], o=140),

                      fname + f'-slice{str(i).zfill(SLICE_DECIMATE_IDENTIFIER)}_x', path)

    if SLICE_Y:

        cnt += dimy

        print('Slicing Y: ')

        for i in range(dimy):

            saveSlice(bp.bayes_noise_removal(image=vol[:, i, :], o=140),

                      fname + f'-slice{str(i).zfill(SLICE_DECIMATE_IDENTIFIER)}_y', path)

    if SLICE_Z:

        cnt += dimz

        print('Slicing Z: ')

        for i in range(dimz):

            saveSlice(bp.bayes_noise_removal(image=vol[:, :, i], o=140),

                      fname + f'-slice{str(i).zfill(SLICE_DECIMATE_IDENTIFIER)}_z', path)

    return cnt

for index, filename in enumerate(sorted(glob.iglob(imagePathInput + '*.nii.gz'))):

    img = readImageVolume(filename)

    print(filename, img.shape, np.sum(img.shape), np.min(img), np.max(img))

    numOfSlices = sliceAndSaveVolumeImage(img, 'heart' + str(index), imageSliceOutput)

    print(f'\n{filename}, {numOfSlices} slices created \n')

# # Read and process image mask volumes

# for index, filename in enumerate(sorted(glob.iglob(maskPathInput + '*.nii'))):

#     img = readImageVolume(filename, False)

#     print(filename, img.shape, np.sum(img.shape), np.min(img), np.max(img))

#     numOfSlices = sliceAndSaveVolumeImage(img, 'tooth' + str(index), maskSliceOutput)

#     print(f'\n{filename}, {numOfSlices} slices created \n')