#!/usr/bin/env python
import os
from multiprocessing import Process, Manager, Value, Pool
import multiprocessing as mp
import subprocess
import argparse
import pathlib
import sys


def ANTS_rigid_body_trans(b0_nii, result, mask_file, reference):

    print("Performing ants rigid body transformation...")
    input_file = b0_nii
    case_name = os.path.basename(input_file)
    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-'
    output_file = os.path.join(os.path.dirname(input_file), output_name)

    trans_matrix = "antsRegistrationSyNQuick.sh -d 3 -f " + reference + " -m " + input_file + " -t r -o " + output_file
    output1 = subprocess.check_output(trans_matrix, shell=True)

    omat_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-0GenericAffine.mat'
    omat_file = os.path.join(os.path.dirname(input_file), omat_name)

    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-Warped.nii.gz'
    transformed_file = os.path.join(os.path.dirname(input_file), output_name)

    result.append((transformed_file, omat_file, mask_file))

SUFFIX_TXT = "txt"
SUFFIX_NIFTI_GZ = "nii.gz"
parser = argparse.ArgumentParser()
parser.add_argument('-b0', action='store', dest='b0', type=str,
                        help="txt file containing list of /path/to/b0, one path in each line")
parser.add_argument('-mask', action='store', dest='mask', type=str,
                        help="txt file containing list of /path/to/mask, one path in each line")
parser.add_argument('-ref', action='store', dest='ref', type=str,
                        help="reference b0 file for registration")

args = parser.parse_args()
reference = str(args.ref)

try:
    args = parser.parse_args()
    if len(sys.argv) == 1:
        parser.print_help()
        parser.error('too few arguments')
        sys.exit(0)

except SystemExit:
    sys.exit(0)

if args.b0:
    f = pathlib.Path(args.b0)
    if f.exists():
        print ("File exist")
        filename = args.b0
    else:
        print ("File not found")
        sys.exit(1)

    # Input caselist.txt
    if filename.endswith(SUFFIX_TXT):
        with open(filename) as f:
            target_list = f.read().splitlines()

if args.mask:
    f = pathlib.Path(args.mask)
    if f.exists():
        print ("File exist")
        filename = args.mask
    else:
        print ("File not found")
        sys.exit(1)

    # Input caselist.txt
    if filename.endswith(SUFFIX_TXT):
        with open(filename) as f:
            mask_list = f.read().splitlines()

with Manager() as manager:
    result = manager.list()              
    ants_jobs = []
    for i in range(0, len(target_list)):
        p_ants = mp.Process(target=ANTS_rigid_body_trans, args=(target_list[i],
                                                             result, mask_list[i], reference))
        ants_jobs.append(p_ants)
        p_ants.start()
        
    for process in ants_jobs:
        process.join()

    result = list(result)

transformed_cases = []
omat_list = []
masks_new_list = []

for subject_ANTS in result:
    transformed_cases.append(subject_ANTS[0])
    omat_list.append(subject_ANTS[1])
    masks_new_list.append(subject_ANTS[2])

# Apply the same tranformation to the mask file
for i in range(0, len(transformed_cases)):

    input_file = transformed_cases[i]
    case_name = os.path.basename(input_file)
    output_name = case_name[:len(case_name) - (len(SUFFIX_NIFTI_GZ) + 1)] + '-mask.nii.gz'
    output_file = os.path.join(os.path.dirname(masks_new_list[i]), output_name)
    apply_mask_trans = "antsApplyTransforms -d 3 -i " + masks_new_list[i] + " -r " + input_file + " -o " \
                            + output_file + " --transform [" + omat_list[i] + "]"

    output2 = subprocess.check_output(apply_mask_trans, shell=True)