#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May  7 00:33:43 2021

@author: valentinahrtonova
"""
from __future__ import division

import os
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
from numpy import asarray, savez_compressed, load
import pandas as pd 


def rescale_image(n, new_min, new_max):
    """Rescale image values to range"""
    
    if not new_min == 0 and n.min() == 0:
        n = (n/n.min()) * new_min
    if not n.max() == 0:
        n = (n/n.max()) * new_max
    return n

def normalize_image(image, img_min, img_max):
    """Window and normalize image"""
    
    norm_image = image + 1024
    
    # Window image
    norm_image[norm_image < img_min] = img_min
    norm_image[norm_image > img_max] = img_max
    image = image.astype("float32")
    
    # Normalize to range [0,1]
    image = (norm_image - img_min) / (img_max - img_min)
    
    return image

def crop_image(image, coords, display=False):
    """Crop image around non-zero area"""

    top_left = np.min(coords, axis=1)
    bottom_right = np.max(coords, axis=1)
    
    # Remove the background
    croped_image = image[top_left[0]:bottom_right[0],
                top_left[1]:bottom_right[1]]
    
    return croped_image

def pad_image(image, new_height, new_width):
    """Pad image with zeros to get given image size"""
    
    height, width = image.shape

    final_image = np.zeros((new_height, new_width))
    pad_left = int((new_width - width) / 2)
    pad_top = int((new_height - height) / 2)
    
    # Replace the pixels with the image's pixels
    final_image[pad_top:pad_top + height, pad_left:pad_left + width] = image
    
    return final_image

"""CREATE SAGITAL SLICES"""

def sagital_from_nifti(folder, df, dataset_file):

    for i in df.shape[0]:
        
        file = df.File[i]
        #file = "verse113"
        folder_path = os.path.join(folder, file)
        verse = nib.load(os.path.join(folder_path,  file + ".nii.gz"))
        verse_seg = nib.load(os.path.join(folder_path,  file + '_seg.nii.gz'))
    
        data = verse.get_fdata()
        data_seg = verse_seg.get_fdata()
        
        tag_master, mask_master, targ_master = list(), list (), list ()
        
        for j in range(df.Min_idx[i],df.Max_idx[i]):
        
            #j = 90
            sagital_slice = np.rot90(data[j, :, :])
            sagital_mask = np.rot90(data_seg[j,:,:])
            mask = rescale_image(sagital_mask, 0, 1)  
            # Find the vertebrae area
            coords = np.array(np.nonzero(mask))
            if coords.size ==0:
                continue
            
            slice = normalize_image(sagital_slice, 0, 3000)
            
            vert = slice * mask    
            
            mask_crop = crop_image(mask, coords)
            vert_crop = crop_image(vert, coords)
            mask_final = pad_image(mask_crop, 256, 256)
            vert_final = pad_image(vert_crop, 256, 256)
            
            tag = file + '_sagital_' + str(j)
                    
            tag_master.append(tag)
            mask_master.append(mask_final)
            targ_master.append(vert_final)   
        
    mask_master = asarray(mask_master)
    targ_master = asarray(targ_master)
    
    print('Loaded: ', mask_master.shape, targ_master.shape)
    # save as compressed numpy array
    savez_compressed(dataset_filename, tag_master, mask_master, targ_master)
    print('Saved dataset: ', dataset_filename) 
 
# Test dataset

test_folder = '/Users/valentinahrtonova/Desktop/Spine_dataset/datasets_test/'
test_filenames = os.listdir(test_folder)

if '.DS_Store' in test_filenames:
    test_filenames.remove('.DS_Store')
    
test_data = {'File': ['verse265', 'verse267', 'verse151'], 
             'Min_idx': [95, 75, 167], 
             'Max_idx': [145, 120, 220], 
             'Min_range': [120, 120, 190], 
             'Max_rand': [300, 290, 390]}  
    
test_df = pd.DataFrame(test_data)

dataset_filename = 'pix2pix_sagital_input_demo.npz'
sagital_from_nifti(test_folder, test_df, dataset_filename)

# Train dataset

train_folder = '/Users/valentinahrtonova/Desktop/Spine_dataset/datasets_train/'
train_filenames = os.listdir(train_folder)

if '.DS_Store' in train_filenames:
    train_filenames.remove('.DS_Store')
    
train_data = {'File': ["verse253", "verse112", "verse113", "verse122", "verse137", "verse145", "verse127", "verse104", "verse135"],
              'Min_idx': [167, 165, 76, 47, 80, 175, 155, 137, 160], 
              'Max_idx': [213, 200, 108, 92, 135, 218, 205, 177, 203], 
              'Min_range': [190, 0, 36, 65, 35, 190, 200, 65, 190], 
              'Max_range': [390, 175, 200, 235, 185, 360, 390, 250, 380]}

train_df = pd.DataFrame(train_data)

dataset_filename = 'pix2pix_sagital_input_demo.npz'
sagital_from_nifti(train_folder, train_df, dataset_filename)
    
"""CREATE AXIAL SLICES"""

def axial_from_nifti(folder, dataset_filename):  
    """ Create dataset from Nifti file"""
        
    filenames = os.listdir(folder)
    
    if '.DS_Store' in filenames:
        filenames.remove('.DS_Store')
    
    tag_master, mask_master, targ_master = list(), list(), list()
    
    for file in filenames:
        
        #file = 'verse265'
        folder_path = os.path.join(folder, file)
        verse = nib.load(os.path.join(folder_path,  file + ".nii.gz"))
        verse_seg = nib.load(os.path.join(folder_path,  file + '_seg.nii.gz'))
        
        header = verse.header
        print(header)
        
        data = verse.get_fdata()
        data_seg = verse_seg.get_fdata()

        
        for i in range(data.shape[2]):
            
            #i = 209
            axial_slice = np.rot90(data[:, :, i])
            axial_mask = np.rot90(data_seg[:,:,i])
            
            # plt.imshow(axial_slice, cmap = 'gray')
            # plt.axis('off')

            # plt.imshow(axial_mask, cmap = 'gray')
            # plt.axis('off')
            
            mask = rescale_image(axial_mask, 0, 1)  
            # Find the vertebrae area
            coords = np.array(np.nonzero(mask))
            if coords.size ==0:
                continue
            
            slice = normalize_image(axial_slice, 0, 3000)
            
            vert = slice * mask
            
            mask_crop = crop_image(mask, coords)
            vert_crop = crop_image(vert, coords)
            mask_final = pad_image(mask_crop, 128, 128)
            vert_final = pad_image(vert_crop, 128, 128)
            
            tag = file + '_axial_' + str(i)
            
            tag_master.append(tag)
            mask_master.append(mask_final)
            targ_master.append(vert_final)       
    
    tag_master = asarray(tag_master)
    mask_master = asarray(mask_master)
    targ_master = asarray(targ_master)
    
    # save as compressed numpy array
    savez_compressed(dataset_filename, tag_master, mask_master, targ_master)
    print('Saved dataset: ', dataset_filename) 

folder = '/Users/valentinahrtonova/Desktop/Spine_dataset/datasets_test/'
dataset_filename = 'pix2pix_axial_input_demo.npz'
axial_from_nifti(folder, dataset_filename)

folder = '/Users/valentinahrtonova/Desktop/Spine_dataset/datasets_train/'
dataset_filename = 'verse_dataset_128_final_train.npz'
axial_from_nifti(folder, dataset_filename)
