from __future__ import print_function

import os
import shutil
import subprocess
from multiprocessing import Pool, cpu_count


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


def fast(src_path, dst_path, temp_dir, temp_path):
    command = ["fast", "-t", "1", "-n", "3", "-H", "0.1", "-I", "1", "-l", "20.0",
               "-o", temp_dir, src_path]
    subprocess.call(command, stdout=open(os.devnull), stderr=subprocess.STDOUT)
    shutil.copyfile(temp_path, dst_path)
    shutil.rmtree(os.path.dirname(temp_dir))
    return


def unwarp_segment(arg, **kwarg):
    return segment(*arg, **kwarg)


def segment(src_path, dst_path, temp_dir, temp_path):
    print("Segment on: ", src_path)
    try:
        fast(src_path, dst_path, temp_dir, temp_path)
    except RuntimeError:
        print("\tFalid on: ", src_path)
    return


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

data_src_paths, data_dst_paths = [], []
temp_dirs, temp_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))
        subj_name = subject.split(".")[0]
        temp_dir = os.path.join(dst_label_dir, subj_name, subj_name)
        create_dir(os.path.dirname(temp_dir))
        temp_dirs.append(temp_dir)
        temp_paths.append(temp_dir + "_pve_1.nii.gz")

# Test
# print(data_src_paths[0], data_dst_paths[0],
#       temp_dirs[0], temp_paths[0])
# segment(data_src_paths[0], data_dst_paths[0],
#         temp_dirs[0], temp_paths[0])

# Multi-processing
subj_num = len(data_src_paths)
paras = zip(data_src_paths, data_dst_paths, temp_dirs, temp_paths)
pool = Pool(processes=cpu_count())
pool.map(unwarp_segment, paras)