import os
import numpy as np
import nibabel
from skimage.io import imsave, imread

data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'data', 'raw')

image_rows = int(512/2)
image_cols = int(512/2) #we will undersample our training 2D images later (for memory and speed)


def create_train_data():
    train_data_path = os.path.join(data_path, 'train')
    images = os.listdir(train_data_path)

    imgs_train=[]  #training images
    imgsliv_train=[] #training masks (corresponding to liver)
    print(('-'*30))
    print('Creating training images...')
    print(('-'*30))
    a=[]
    b=[]
    for k in range(len(images)):
        if k%2==0:
            a.append(np.sort(images)[k]) #file names corresponding to training masks
        else:
            b.append(np.sort(images)[k]) #file names corresponding to training images
        
    for liver,orig in zip(a,b):
        imgl=nibabel.load(os.path.join(train_data_path,liver)) #we load 3D training mask
        imgo=nibabel.load(os.path.join(train_data_path,orig)) #we load 3D training image
        for k in range(imgl.shape[2]):
            dimgl=np.array(imgl.get_data()[::2,::2,k]) #axial cuts are made along the z axis with undersampling
            dimgo=np.array(imgo.get_data()[::2,::2,k])
            if len(np.unique(dimgl))!=1: #we only recover the 2D sections containing the liver
                imgsliv_train.append(dimgl)
                imgs_train.append(dimgo)
            
                
    imgs = np.ndarray((len(imgs_train), image_rows, image_cols), dtype=np.uint8)
    imgs_mask = np.ndarray((len(imgsliv_train), image_rows, image_cols), dtype=np.uint8)
    for index,img in enumerate(imgs_train):
        imgs[index,:,:]=img
    for index,img in enumerate(imgsliv_train):
        imgs_mask[index,:,:]=img

    np.save(os.path.join(data_path, '..', 'processed', 'imgs_train.npy'), imgs)
    np.save(os.path.join(data_path, '..', 'processed', 'imgsliv_train.npy'), imgs_mask)
    print('Saving to .npy files done.')


def load_train_data():
    imgs_train = np.load(os.path.join(data_path, '..', 'processed', 'imgs_train.npy'))
    imgs_mask_train = np.load(os.path.join(data_path, '..', 'processed', 'imgsliv_train.npy'))
    return imgs_train, imgs_mask_train


def create_test_data():
    test_data_path = os.path.join(data_path, 'test')
    images = os.listdir(test_data_path)
    print(('-'*30))
    print('Creating testing images...')
    print(('-'*30))
    imgs_test=[]
    imgsliv_test=[]
    for image_name in images:
        print(image_name)
        img=nibabel.load(os.path.join(test_data_path,image_name))
        print((img.shape))
        for k in range(img.shape[2]):  
            dimg=np.array(img.get_data()[::2,::2,k])
            if 'liver' in image_name:
                imgsliv_test.append(dimg)
            
            elif 'orig' in image_name:
                imgs_test.append(dimg)
    
    imgst= np.ndarray((len(imgs_test), image_rows, image_cols), dtype=np.uint8)
    imgs_maskt= np.ndarray((len(imgsliv_test), image_rows, image_cols), dtype=np.uint8)
    for index,img in enumerate(imgs_test):
        imgst[index,:,:]=img
    for index,img in enumerate(imgsliv_test):
        imgs_maskt[index,:,:]=img

    np.save(os.path.join(data_path, '..', 'processed', 'imgs_test.npy'), imgst)
    np.save(os.path.join(data_path, '..', 'processed', 'imgsliv_test.npy'), imgs_maskt)
    print('Saving to .npy files done.')


def load_test_data():
    imgst = np.load(os.path.join(data_path, '..', 'processed', 'imgs_test.npy'))
    imgs_id = np.load(os.path.join(data_path, '..', 'processed', 'imgsliv_test.npy'))
    return [imgst, imgs_id]

if __name__ == '__main__':
    create_train_data()
    create_test_data()