#!/usr/bin/python
import os, glob
from multiprocessing import Pool
from functools import partial


# Perform motion tracking for cine MR
def track_cine(data_dir, par_dir, tamplate_dir):
    
    motion_dir = os.path.join(data_dir, 'motion')
    if os.path.exists(motion_dir):
       os.system('rm -r {0}'.format(motion_dir))
       os.mkdir(motion_dir)
    else:
       os.mkdir(motion_dir)

    # Split 4D nifti into time phases
    os.system('splitvolume '
              '{0}/lvsa_.nii.gz '
              '{1}/lvsa_ '
              '-sequence'
              .format(data_dir, motion_dir))

    n_frame = len(glob.glob('{0}/lvsa_??.nii.gz'.format(motion_dir)))

    # Inter-frame motion estimation
    for fr in range(1, n_frame):
        target = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr-1)
        source = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr)
        par = '{0}/ffd_motion.cfg'.format(par_dir)
        dof = '{0}/ffd_{1:02d}_to_{2:02d}.dof.gz'.format(motion_dir, fr-1, fr)

        if not os.path.exists(dof):
            os.system('mirtk register '
                      '{0} '
                      '{1} '
                      '-parin {2} '
                      '-dofout {3}'
                      .format(target, source, par, dof))

    # Compose inter-frame transformation fields
    for fr in range(2, n_frame):
        dofs = ''
        for k in range(1, fr+1):
            dof = '{0}/ffd_{1:02d}_to_{2:02d}.dof.gz'.format(motion_dir, k-1, k)
            dofs += dof + ' '
            dof_out = '{0}/ffd_comp_00_to_{1:02d}.dof.gz'.format(motion_dir, fr)
            if not os.path.exists(dof_out):
                os.system('mirtk compose-dofs '
                          '{0} '
                          '{1}'
                          .format(dofs, dof_out))

    # Refine motion fields
    # Composition of inter-frame motion fields can lead to accumulative errors. At this step, we refine the motion fields
    # by re-registering the n-th frame with the ED frame.
    for fr in range(2, n_frame):
        target = '{0}/lvsa_00.nii.gz'.format(motion_dir)
        source = '{0}/lvsa_{1:02d}.nii.gz'.format(motion_dir, fr)
        par = '{0}/ffd_refine.cfg'.format(par_dir)
        dofin = '{0}/ffd_comp_00_to_{1:02d}.dof.gz'.format(motion_dir, fr)
        dof = '{0}/ffd_00_to_{1:02d}.dof.gz'.format(motion_dir, fr)
        if not os.path.exists(dof):
            os.system('mirtk register '
                      '{0} '
                      '{1} '
                      '-parin {2} '
                      '-dofin {3} '
                      '-dofout {4}'
                      .format(target, source, par, dofin, dof))

    # Obtain the RV mesh with the same number of points as the template mesh
    os.system('srreg '
              '{2}/vtks/RV_ED.vtk '
              '{0}/RVendo_ED.vtk '
              '-dofout {1}/RV_srreg.dof.gz '
              '-symmetric'
              .format(tamplate_dir, motion_dir, data_dir))
    
    os.system('mirtk transform-points '
              '{0}/RVendo_ED.vtk '
              '{1}/RV_ED_srreg.vtk '
              '-dofin {1}/RV_srreg.dof.gz '
              '-invert'
              .format(tamplate_dir, motion_dir))
    
    os.system('sareg '
              '{0}/RV_ED_srreg.vtk '
              '{1}/vtks/RV_ED.vtk '
              '-dofout {0}/RV_sareg.dof.gz '
              '-symmetric'
              .format(motion_dir, data_dir))
    
    os.system('snreg '
              '{0}/RV_ED_srreg.vtk '
              '{1}/vtks/RV_ED.vtk '
              '-dofin {0}/RV_sareg.dof.gz '
              '-dofout {0}/RV_snreg.dof.gz '
              '-ds 20 -symmetric'
              .format(motion_dir, data_dir))
    
    os.system('mirtk transform-points '
              '{0}/RV_ED_srreg.vtk '
              '{0}/RV_fr00.vtk '
              '-dofin {0}/RV_snreg.dof.gz'
              .format(motion_dir))

    # Transform the mesh
    for fr in range(1, n_frame):
        os.system('mirtk transform-points '
                  '{0}/RV_fr00.vtk '
                  '{0}/RV_fr{1:02d}.vtk '
                  '-dofin {0}/ffd_00_to_{1:02d}.dof.gz'
                  .format(motion_dir, fr))
     
    # Convert vtks to text files
    for fr in range(0, n_frame):
        os.system('vtk2txt '
                  '{0}/RV_fr{1:02d}.vtk '
                  '{0}/RV_fr{1:02d}.txt'
                  .format(motion_dir, fr))


def apply_PC(subject, data_dir, param_dir, template_dir):
    
    print('  co-registering {0}'.format(subject))    
    
    subject_dir = os.path.join(data_dir, subject)
        
    if os.path.isdir(subject_dir):
        
        track_cine(subject_dir, param_dir, template_dir)
        
        print('  finish motion tracking in subject {0}'.format(subject))
        
    else:  
        print('  {0} is not a valid directory, do nothing'.format(subject_dir))


def motionTracking(dir_0, dir_2, dir_3, coreNo, parallel):
               
    if parallel:
    
        print('Motion tracking running on {0} cores'.format(coreNo))
        
        pool = Pool(processes = coreNo) 
        
        # partial only in Python 2.7+
        pool.map(partial(apply_PC, 
                         data_dir=dir_0,  
                         param_dir=dir_2, 
                         template_dir=dir_3), 
                         sorted(os.listdir(dir_0)))       
                
    else:
        
        print('Motion tracking from segmentations running subsequently')
                
        data_dir, param_dir, template_dir = dir_0, dir_2, dir_3
                                     
        for subject in sorted(os.listdir(data_dir)):
            
            subject_dir = os.path.join(data_dir, subject)

            if not os.path.isdir(subject_dir):
                
                print('  {0} is not a valid folder, Skip'.format(subject_dir))
                
                continue 
                   
            print('  motion tracking {0}'.format(subject))  
            
            track_cine(subject_dir, param_dir, template_dir)
            
            print('  finish motion tracking in {0}'.format(subject))