import os
from fsl.eddy_correct_diff import eddy_corr


def basic_files(cortex_only=True,atlas_type='mega',folder_name=r'C:\Users\Admin\Desktop\v7_calibration\thebase4ever'):
    if atlas_type == 'mega':
        if cortex_only:
            atlas_label = r'C:\Users\Admin\my_scripts\aal\megaatlas\MegaAtlas_cortex_Labels.nii'
        else:
            atlas_label = r'C:\Users\Admin\my_scripts\aal\megaatlas\MegaAtlas_Labels_highres.nii'

        atlas_template = r'C:\Users\Admin\my_scripts\aal\megaatlas\Schaefer_template.nii'

    elif atlas_type == 'aal3':
        atlas_label = r'F:\Hila\aal\aal3\AAL3_highres_atlas.nii'
        atlas_template = r'F:\Hila\aal\aal3\AAL3_highres_template.nii'
        #atlas_label = r'F:\Hila\aal\aal3\registered\AAL3_highres_atlas_corrected.nii'
        #atlas_template = r'F:\Hila\aal\aal3\registered\MNI152_T1_1mm.nii'

    elif atlas_type == 'yeo7_200':
        #atlas_label = r'F:\Hila\aal\yeo7_200\yeo7_200_atlas.nii'
        #atlas_template = r'F:\Hila\aal\yeo7_200\Schaefer_template.nii'
        atlas_label = r'G:\data\atlases\yeo\yeo7_200\yeo7_200_atlas.nii'
        atlas_template = r'G:\data\atlases\yeo\yeo7_200\Schaefer_template.nii'

    elif atlas_type == 'yeo7_1000':
        atlas_label = r'C:\Users\Admin\my_scripts\aal\yeo\yeo7_1000\yeo7_1000_atlas.nii'
        atlas_template = r'C:\Users\Admin\my_scripts\aal\yeo\yeo7_1000\Schaefer_template.nii'

    elif atlas_type == 'yeo17_1000':
        atlas_label = r'C:\Users\Admin\my_scripts\aal\yeo\yeo17_1000\yeo17_1000_atlas.nii'
        atlas_template = r'C:\Users\Admin\my_scripts\aal\yeo\yeo17_1000\Schaefer_template.nii'

    elif atlas_type == 'bna':
        atlas_label = r'G:\data\atlases\BNA\BN_Atlas_274_combined_1mm.nii'
        atlas_template = r'G:\data\atlases\BNA\MNI152_T1_1mm.nii'

    elif atlas_type == 'bna_cor':
        atlas_label = r'G:\data\atlases\BNA\newBNA_Labels.nii'
        atlas_template = r'G:\data\atlases\BNA\MNI152_T1_1mm.nii'

    elif atlas_type == 'yeo7_100':
        atlas_label = r'G:\data\atlases\yeo\yeo7_100\yeo7_100_atlas.nii'
        atlas_template = r'G:\data\atlases\yeo\yeo7_100\Schaefer_template.nii'


    atlas_label = os_path_2_fsl(atlas_label)
    atlas_template = os_path_2_fsl(atlas_template)

    # #folder_name = r'F:\Hila\Ax3D_Pack\V6\after_file_prep'
    # folder_name = r'C:\Users\Admin\Desktop\v7_calibration\thebase4ever'
    # #folder_name = r'C:\Users\Admin\Desktop\Language'
    all_subj_folders = os.listdir(folder_name)
    subj = all_subj_folders

    return subj, folder_name, atlas_template, atlas_label


def subj_files(subj_folder):
    mprage_file_name=[]
    diff_file_name=[]
    pa_file_name=[]
    for file in os.listdir(subj_folder):
        if 'MPRAGE' in file and not (file.startswith('r') or 'brain' in file):
            mprage_file_name = file
        if (file.endswith('001.nii') and 'AP' in file) or  ('diff_corrected_D' in file and file.endswith('.nii') and not 'mask' in file) or ('data.nii' in file):
            diff_file_name = file
        if (file.endswith('001.nii') and 'PA' in file) or ('data_PA.nii' in file):
            pa_file_name = file
    return mprage_file_name, diff_file_name, pa_file_name


def bet_4_regis_mprage(subj_folder,mprage_file_name,diff_file_name):

    subj_mprage = subj_folder + mprage_file_name
    # BET for registered MPRAGE:
    if subj_mprage.endswith('.nii'):
        in_brain = subj_mprage[:-4]
        out_brain = subj_mprage[:-4]+'_brain'
        diff_file_1st = diff_file_name[:-4]+'_1st'
    elif subj_mprage.endswith('.nii.gz'):
        out_brain = subj_mprage[:-7]+'_brain'
        diff_file_1st = diff_file_name[:-7]+'_1st'
        in_brain = subj_mprage[:-7]

    cmd = 'bash -lc "bet {0} {1} {2} {3}"'.format(in_brain, out_brain,'-f 0.30','-g 0.20')
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)
    # save first corrected diff:
    cmd = fr'bash -lc "fslroi {subj_folder}/{diff_file_name} {subj_folder}/{diff_file_1st} 0 1"'
    os.system(cmd)

    return subj_mprage, out_brain


def reg_from_chm_2_mprage(subj_folder,subj_mprage):
    subj_first_charmed = subj_folder + '/diff_corrected_1st.nii'
    out_registered = subj_folder + '/rdiff_corrected_1st.nii'
    out_registered_mat = out_registered[:-4] +'.mat'
    options = '-bins 256 -cost normmi -searchrx -90 90 -searchry -90 90 -searchrz -90 90 -dof 12'

    cmd = 'bash -lc "flirt -ref {0} -in {1} -out {2} -omat {3} {4}"'.format(subj_mprage, subj_first_charmed, out_registered, out_registered_mat, options)
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return subj_first_charmed, out_registered, out_registered_mat


def create_inv_mat(out_registered_mat):
    inv_mat = out_registered_mat[:-4] + '_inv.mat'
    cmd = 'bash -lc "convert_xfm -omat {0} -inverse {1}"'.format(inv_mat, out_registered_mat)
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return inv_mat


def reg_from_mprage_2_chm_inv(subj_folder, mprage_file_name, out_brain, subj_first_charmed, inv_mat):
    if mprage_file_name.endswith('.nii'):
        out_registered = subj_folder + 'r' + mprage_file_name[:-4] + '_brain.nii'
    elif mprage_file_name.endswith('.nii.gz'):
        out_registered = subj_folder + 'r' + mprage_file_name[:-7] + '_brain.nii'

    cmd = 'bash -lc "flirt -in {0} -ref {1} -out {2} -applyxfm -init {3}"'.format(out_brain, subj_first_charmed, out_registered, inv_mat)
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return out_registered


def bet_4_atlas(atlas_template):
    atlas_brain = atlas_template[:-4] + '_brain'
    cmd = 'bash -lc "bet {0} {1} {2} {3}"'.format(atlas_template[:-4], atlas_brain,'-f 0.45','-g -0.1')
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return atlas_brain


def flirt_primary_guess(subj_folder,atlas_template, out_registered):
    options = r'-bins 256 -cost corratio -searchrx -90 90 -searchry -90 90 -searchrz -90 90 -dof 12 -interp trilinear'
    atlas_brain = atlas_template[:-4] + '_brain.nii'

    atlas_registered_flirt = os.path.join(subj_folder+ 'r' + atlas_brain.split(sep="\\")[-1])
    atlas_registered_flirt_mat = atlas_registered_flirt[:-4] + '.mat'

    cmd = f'bash -lc "flirt -ref {out_registered} -in {atlas_brain} -out {atlas_registered_flirt} -omat {atlas_registered_flirt_mat} {options}"'
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return atlas_brain, atlas_registered_flirt, atlas_registered_flirt_mat


def fnirt_from_atlas_2_subj(subj_folder,out_registered, atlas_brain, atlas_registered_flirt_mat, cortex_only = True):
    if cortex_only:
        warp_name = subj_folder + 'atlas2subjmegaatlas.nii'
    else:
        warp_name = subj_folder + 'atlas2subj.nii'
    cmd = f'bash -lc "fnirt --ref={out_registered} --in={atlas_brain} --aff={atlas_registered_flirt_mat} --cout={warp_name}"'
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return warp_name


def apply_fnirt_warp_on_template(subj_folder, atlas_brain, out_registered, warp_name):
    atlas_registered = os.path.join(subj_folder+ 'rr' + atlas_brain.split(sep="\\")[-1])
    cmd = f'bash -lc "applywarp --ref={out_registered} --in={atlas_brain} --out={atlas_registered} --warp={warp_name} --interp=nn"'
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return atlas_registered


def apply_fnirt_warp_on_label(subj_folder, atlas_label, out_registered, warp_name):
    atlas_labels_registered = os.path.join(subj_folder+ 'r' + atlas_label.split(sep="\\")[-1])
    cmd = 'bash -lc "applywarp --ref={0} --in={1} --out={2} --warp={3} --interp={4}"'.format(out_registered, atlas_label, atlas_labels_registered, warp_name, 'nn')
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)

    return atlas_labels_registered


def apply_fnirt_warp_on_masks(subj_folder, atlas_mask, out_registered, warp_name):
    atlas_mask_registered = os.path.join(subj_folder+ 'r' + atlas_mask.split(sep="\\")[-1])
    cmd = 'bash -lc "applywarp --ref={0} --in={1} --out={2} --warp={3} --interp={4}"'.format(out_registered, atlas_mask, atlas_mask_registered, warp_name, 'nn')
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)


def fast_seg(out_registered):
    options = r'-t 1 -n 3 -H 0.1 -I 4 -l 10.0 -o'
    cmd = 'bash -lc "fast {0} {1} {2}"'.format(options, out_registered, out_registered)
    cmd = cmd.replace(os.sep,'/')
    os.system(cmd)


def os_path_2_fsl(path):
    if 'F:' in path:
        path = path.replace('F:','/mnt/f')
    elif 'C:' in path:
        path = path.replace('C:', '/mnt/c')
    elif 'D:' in path:
        path = path.replace('D:', '/mnt/d')
    elif 'G:' in path:
        path = path.replace('G:', '/mnt/g')
    elif 'Y:' in path:
        path = path.replace('Y:', '/mnt/y')
    return path





def all_func_to_run(s, folder_name, atlas_template, atlas_label, fast = True, atlas_registration_only = False):
    subj_name = r'/' + s + r'/'
    subj_folder = folder_name + subj_name
    subj_folder = subj_folder.replace(os.sep, '/')

    mprage_file_name, diff_file_name, pa_file_name = subj_files(subj_folder)
    out_registered = None
    subj_folder = os_path_2_fsl(subj_folder)
    if not atlas_registration_only:
        eddy_corr(subj_folder,diff_file_name,pa_file_name,acqr_file='/mnt/h/comp/Desktop/language_eddy/datain.txt',index_file='/mnt/h/comp/Desktop/language_eddy/index186.txt')

        subj_mprage, out_brain = bet_4_regis_mprage(subj_folder,mprage_file_name,'diff_corrected.nii')

        ''' Registration from MPRAGE to 1st CHARMED scan using inverse matrix of CHARMED to MPRAGE registration:
        From CHARMED to MPRAGE:'''
        subj_first_charmed, out_registered, out_registered_mat = reg_from_chm_2_mprage(subj_folder, subj_mprage)

        '''Creation of inverse matrix:  '''
        inv_mat = create_inv_mat(out_registered_mat)

        '''From MPRAGE to CHARMED using the inverse matrix: '''
        out_registered = reg_from_mprage_2_chm_inv(subj_folder, mprage_file_name, out_brain, subj_first_charmed, inv_mat)

    ''' BET for mni template:
        BET for mni template:
        if not performed before, run:   '''
    # atlas_brain = bet_4_atlas(atlas_template)

    '''Registration from atlas to regisered MPRAGE:
        flirt for atlas to registered MPRAGE for primary guess:  '''
    if not out_registered:
        out_registered = subj_folder + 'rMPRAGE_brain.nii'
    atlas_brain, atlas_registered_flirt, atlas_registered_flirt_mat = flirt_primary_guess(subj_folder, atlas_template,
                                                                                          out_registered)

    '''fnirt for atlas based on flirt results:    '''
    warp_name = fnirt_from_atlas_2_subj(subj_folder, out_registered, atlas_brain, atlas_registered_flirt_mat,
                                        cortex_only=False)

    '''apply fnirt warp on atlas template:  '''
    #atlas_registered = apply_fnirt_warp_on_template(subj_folder, atlas_brain, out_registered, warp_name)

    '''apply fnirt warp on atlas labels:   '''
    atlas_labels_registered = apply_fnirt_warp_on_label(subj_folder, atlas_label, out_registered, warp_name)

    if fast:
        '''FAST segmentation:   '''
        fast_seg(out_registered)

    print('Finished file prep for ' + subj_name[:-1])


if __name__ == '__main__':
    from multiprocessing import Process
    #subj, folder_name, atlas_template, atlas_label = basic_files(False, atlas_type='yeo7_200',folder_name = 'F:\Hila\TDI\moreTheBase4Ever')


    subj, folder_name, atlas_template, atlas_label = basic_files(False, atlas_type='yeo7_100', folder_name = 'F:\Hila\TDI\TheBase4Ever')
    for s in subj:
        all_func_to_run(s, folder_name, atlas_template, atlas_label, fast=False, atlas_registration_only=True)
