from __future__ import print_function

import os
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
from multiprocessing import Pool, cpu_count
from scipy.ndimage.interpolation import zoom


def plot_middle(data, slice_no=None):
    if not slice_no:
        slice_no = data.shape[-1] // 2
    plt.figure()
    plt.imshow(data[..., slice_no], cmap="gray")
    plt.show()
    return


def create_dir(path):
    if not os.path.isdir(path):
        os.makedirs(path)
    return


def load_nii(path):
    return nib.load(path).get_data()


def save_nii(data, path):
    nib.save(nib.Nifti1Image(data, np.eye(4)), path)
    return


def brain_mask(data, mask):
    return np.multiply(data, mask)


def resize(data, target_shape=[96, 112, 96]):
    factor = [float(t) / float(s) for t, s in zip(target_shape, data.shape)]
    resized = zoom(data, zoom=factor, order=1, prefilter=False)
    return resized


def norm(data):
    # obj_idx = np.where(data > 0)
    # obj = data[obj_idx]
    # obj_mean, obj_std = np.mean(obj), np.std(obj)
    # obj = (obj - obj_mean) / obj_std
    # data[obj_idx] = obj
    data = data / float(np.max(data))
    return data


def unwarp_postprocess(arg, **kwarg):
    return postprocess(*arg, **kwarg)


def postprocess(src_path, dst_path, temp_path=None, is_mask=False):
    print("Wroking on: ", src_path)
    try:
        data = load_nii(src_path)
        if is_mask:
            mask = load_nii(temp_path)
            data = brain_mask(data, mask)
        data = resize(data)
        # data = norm(data)
        save_nii(data, dst_path)
    except RuntimeError:
        print("\tFalid on: ", src_path)


parent_dir = os.path.dirname(os.getcwd())
data_dir = os.path.join(parent_dir, "data")
data_src_dir = os.path.join(data_dir, "ADNISegment")
data_dst_dir = os.path.join(data_dir, "ADNISegmentPost")
data_labels = ["AD", "NC"]
create_dir(data_dst_dir)

data_src_paths, data_dst_paths = [], []
for label in data_labels:
    src_label_dir = os.path.join(data_src_dir, label)
    dst_label_dir = os.path.join(data_dst_dir, label)
    create_dir(dst_label_dir)
    for subject in os.listdir(src_label_dir):
        data_src_paths.append(os.path.join(src_label_dir, subject))
        data_dst_paths.append(os.path.join(dst_label_dir, subject))

temp_path = os.path.join(data_dir, "Template", "bianca_exclusion_mask.nii.gz")

# Test
# postprocess(data_src_paths[0], data_dst_paths[0], temp_path)

# Multi-processing
is_mask = False
subj_num = len(data_src_paths)
paras = zip(data_src_paths, data_dst_paths,
            [temp_path] * subj_num, [is_mask] * subj_num)
pool = Pool(processes=cpu_count())
pool.map(unwarp_postprocess, paras)