#%%
import pandas as pd 
import os.path as op
import re
from nilearn import masking
from glob import glob
from nilearn.image import load_img, math_img, resample_to_img
import sys

dir_path = sys.argv[1]

# dir_path = "/tmp/data"
mask_path = "/extras/brainmask_12.8.nii"
print("Analyzing data in {}".format(dir_path))

#%%
# extract for resampled at 4 and smoothed 4
resolutions = ['r4s4', 'r4s8', 'r8s4', 'r8s8']
tissues = ['p1', 'p2']

mask = load_img(mask_path)
df_total = pd.DataFrame()

subj_path =  glob(op.join(dir_path, '**/r4s4*p1sub*'), recursive=True)[0]
subj_file = op.basename(subj_path)
if not subj_path:
    raise ValueError('no subject found')
sess = op.dirname(subj_path)
df = pd.DataFrame()
subj_names = []
session_names = []
sub_paths = [] 
sub_paths.append(sess)
s_name = 'sub-' + re.findall('sub-(.*?)_T1', subj_file)[0]
subj_names.append(s_name)
if '/ses' in sess:
    sess_name = "ses-" + re.findall('ses-(.*?)_T1', subj_file)[0]
    session_names.append(sess_name)
else:
    sess_name = 'ses-1'
    session_names.append(sess_name)
print("Extracting features for {}".format(sess))
for resol in resolutions:
    print("resolution & smoothing: {}".format(resol))
    for tissue in tissues:
        print("Tissue (p1=GM, p2=WM, p3=CSF): {}".format(tissue))
        subjects = glob(op.join(sess, "**", '*{}r_{}sub*'.format(
                        resol, tissue)), recursive=True)
        resampled_mask = resample_to_img(mask, subjects[0])
        bin_mask = math_img('img >= 0.5', img=resampled_mask)
        masked = masking.apply_mask(imgs=subjects, mask_img=bin_mask)
        tis_r_s = '_'.join([tissue, resol])
        file_name = '_'.join([s_name, sess_name,
                                '_features']) + '.csv'
        cols = list(range(0, len(masked[0])))
        df_temp = pd.DataFrame(masked, columns=cols)
        df_temp = df_temp.add_prefix(tis_r_s + '_')
        df = pd.concat([df, df_temp], axis=1)
df['subj_path'] = sub_paths
df['subject'] = subj_names
df['session'] = session_names
df.to_csv(op.join(sess, file_name))
# %%