import numpy as np
import subprocess
import os
from os.path import join as opj
import re
import nibabel as nib

# paths = np.genfromtxt('/home1/varunk/results_again_again/anat_file_paths.txt', dtype='str') #Didn't work

# paths = np.genfromtxt('/home1/varunk/results_again_again/skullstrip_anat.txt', dtype='str') 

paths = np.genfromtxt('/home1/varunk/results_again_again/bet_resample_reg_brain_paths.txt', dtype='str')

# paths = np.genfromtxt('/home1/varunk/results_again_again/bet_resample_brain_paths.txt', dtype='str')
paths = np.sort(paths)
save_destination = '/home1/varunk/results_again_again/registered_anat_qc/'
if not os.path.exists(save_destination):
    os.mkdir(save_destination)
    
os.chdir(save_destination)

# print(np.sort(paths)[0:10])
paths

"""
Explanation: FSL - Brain Extraction Tool(BET) Quality Check
This script was created when the BET performed realy poorly on subjects of ABIDE. It is a very simple script that does the following:
* Read the brain nifti files using the bash's find command and save it to a file as -
find . -name "sub-*.nii.gz" &gt; filename.txt
* Read the file to extract all the filenames.
* Count the number of brain voxels in each nifti file.
* Then find the number of brains whose number brain voxels are some specific standard deviations away from the mean.
I used another metric that didn't work
* Count the voxels that doesnot lie inside the standard mask.
- This didn't work as ater the transformation to the standard space the brains were totally skewed but lied inside the standard mask.
End of explanation
"""


# paths = paths[0:3]
res_voxels = []
for path in paths:
    base_dir = '/home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg'
    anat_path = base_dir + path[1:]
    print(anat_path)
    sub_id_extracted = re.search('.+_subject_id_(\d+)', anat_path).group(1)
#     out_file = 'sub_' + sub_id_extracted + '_residual'
    
    res_file = opj(os.getcwd(),out_file)
    proc = subprocess.Popen(['fslstats', anat_path, '-V'], 
                             stdout=subprocess.PIPE)
    voxels = int((proc.communicate()[0]).decode("utf-8").split(' ')[0])
    res_voxels.append(voxels)

res_voxels

%matplotlib inline
import matplotlib.pyplot as plt

bins = np.arange(min(res_voxels), max(res_voxels), 500) # fixed bin size

res = plt.hist(res_voxels, 
         bins=bins, 
         alpha=0.5, 
         color='#887E43',
         label='Defender')

np.mean(res_voxels),np.std(res_voxels)

#number_of_outliers = 
(res_voxels > 10.0*np.std(res_voxels)).sum()

paths[np.where((res_voxels > 20.0*np.std(res_voxels)) == True)]

np.mean(res_voxels),np.std(res_voxels)

#number_of_outliers = 
s = 3
(res_voxels > np.mean(res_voxels) + s*np.std(res_voxels)).sum(), (res_voxels < np.mean(res_voxels) - s*np.std(res_voxels)).sum()

# index of outliers
np.where((res_voxels > 6.0*np.std(res_voxels)) == True)

# Subjects that are outliers
paths[np.where((res_voxels > 15.0*np.std(res_voxels)) == True)]

# Subjects that are outliers
set(paths[np.where((res_voxels > 12.0*np.std(res_voxels)) == True)])

# Subjects that are outliers
set(paths[np.where((res_voxels > 11.0*np.std(res_voxels)) == True)]) -  set(paths[np.where((res_voxels > 12.0*np.std(res_voxels)) == True)])

set(paths[np.where((res_voxels > 10.0*np.std(res_voxels)) == True)]) - set(paths[np.where((res_voxels > 11.0*np.std(res_voxels)) == True)])

paths[np.where((res_voxels > 10.0*np.std(res_voxels)) == True)].shape 

subject_ids_corrupted = [] 
for path in paths[np.where((res_voxels > 6.0*np.std(res_voxels)) == True)]:
    sub_id_extracted = re.search('.+_subject_id_(\d+)', path).group(1)
    subject_ids_corrupted.append(str(int(sub_id_extracted)))

subject_ids_corrupted

# Find the error:
import nibabel as nib

res_voxels = []

res_file = '/home1/varunk/results_again_again/registered_anat_qc/0050977_residual.nii.gz'
proc = subprocess.Popen(['fslstats', res_file, '-V'], 
                         stdout=subprocess.PIPE)
voxels = int((proc.communicate()[0]).decode("utf-8").split(' ')[0])
res_voxels.append(voxels)

res_voxels

!fslmaths /home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg/_subject_id_0050977/skullStrip/sub-0050977_T1w_resample_brain.nii.gz -mul /home1/varunk/results_again_again/registered_anat_qc/mask_inverted.nii.gz 0050977_residual

paths = paths[0:3]
for path in paths:
    base_dir = '/home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg'
    anat_path = base_dir + path[1:]
    print(anat_path)

# (stddata[0]).decode("utf-8").split(' ')[0]

"""
Explanation: Metric:  Number of non zero voxels (Brain Voxels)
End of explanation
"""


mask = '/home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg/resample_mni/MNI152_T1_2mm_brain_resample_mask.nii.gz'

proc = subprocess.Popen(['fslmaths', mask, '-mul', '-1', '-add' ,'1', 'mask_inverted'], 
                         stdout=subprocess.PIPE)
stdoutdata= proc.communicate()
print("The commandline is: {}".format(subprocess.list2cmdline(proc.args)))

cwd = os.getcwd()

mask_inverted_path = opj(cwd, 'mask_inverted.nii.gz')


# /home1/varunk/results_again_again/registered_anat_qc/mask_inverted.nii.gz

# paths

%%time
# paths = paths[0:3]
res_voxels = []
for path in paths:
    base_dir = '/home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg'
    anat_path = base_dir + path[1:]
#     print(anat_path)
    sub_id_extracted = re.search('.+_subject_id_(\d+)', anat_path).group(1)
    out_file = 'sub_' + sub_id_extracted + '_residual'
    proc = subprocess.Popen(['fslmaths', anat_path, '-mul', mask_inverted_path, out_file], 
                         stdout=subprocess.PIPE)
    stdoutdata= proc.communicate()
    print("The command executed is: {}".format(subprocess.list2cmdline(proc.args)))
#     print('Created the residual file for Subject', sub_id_extracted)

    

    res_file = opj(os.getcwd(),out_file)
    proc = subprocess.Popen(['fslstats', res_file, '-V'], 
                             stdout=subprocess.PIPE)
    voxels = int((proc.communicate()[0]).decode("utf-8").split(' ')[0])
    res_voxels.append(voxels)
    
#     print(stdoutdata)
    
# , stderrdata ,    
# , stderr=subprocess.STDOUT

!fslmaths /home1/varunk/results_again_again/ABIDE1_Preprocess/motion_correction_bet/coreg_reg/_subject_id_0050002/skullStrip/sub-0050002_T1w_resample_brain.nii.gz -mul /home1/varunk/results_again_again/registered_anat_qc/mask_inverted.nii.gz sub_0050002_residual

res_voxels

np.mean(res_voxels),np.std(res_voxels)

#number_of_outliers = 
(res_voxels > 6*np.std(res_voxels)).sum()

# index of outliers
np.where((res_voxels > 6*np.std(res_voxels)) == True)

# Subjects that are outliers
paths[np.where((res_voxels > 6*np.std(res_voxels)) == True)]

"""
Explanation: The metric that did not work.
It could not actually figure out the bad brains
Create the inverted mask
End of explanation
"""


