import os
from os.path import join,abspath,basename
from glob import glob
from bids import BIDSLayout
from scripts.commonsmk import bids_name
import shutil
import yaml

#with open('/home/dcao6/scratch/autofid_final/Auto-afids-feature-store-model-train/modelling/config/config.yml') as file:
    #config = yaml.load(file, Loader=yaml.FullLoader)
    
configfile: "config/config.yml"

# Copy dataset_description.json files to derivatives directories (for BIDSlayout indexing)
deriv_dirs=2*[0]
deriv_dirs[0]= features = join(config['output_dir'],'derivatives','features')
deriv_dirs[1]= models = join(config['output_dir'],'derivatives','models')

for idir in deriv_dirs:
    if not os.path.exists(idir):
        os.makedirs(idir)
    data_des = abspath(join(join(srcdir("script.sh"),"../.."),'resources','dataset_description_'+basename(idir)+'.json'))
    if not os.path.exists(join(idir,'dataset_description.json')):
        shutil.copy2(data_des, join(idir,'dataset_description.json'))

# Load the BIDS directory
layout = BIDSLayout(join(config['input_dir'],'bids'), derivatives=join(config['input_dir'],'derivatives'),validate=False)

subjects = layout.get_subjects(**config['meta_nifti'])
sessions = layout.get_sessions(**config['meta_nifti'])
runs = layout.get_runs(**config['meta_nifti'])

suffix = config['meta_nifti']['suffix'] + config['meta_nifti']['extension']
suffix_afids = config['derivatives_meta']['afids_groundtruth']['suffix'] + config['derivatives_meta']['afids_groundtruth']['extension']

afids_nums=[str(x).zfill(2) for x in list(range(1,33))]
train_levels = ['coarse', 'medium', 'fine']

if config['registration']:
	out_features=bids_name(root=join(config['output_dir'],'derivatives','features'),prefix='afid-{afid_num}', suffix='features.hkl',**{'desc':'{train_level}','space':'MNI152NLin2009cAsym'})
	out_models=bids_name(root=join(config['output_dir'],'derivatives','models'),prefix='afid-{afid_num}', suffix='model',**{'desc':'{train_level}','space':'MNI152NLin2009cAsym'})
else:
	out_features=bids_name(root=join(config['output_dir'],'derivatives','features'),prefix='afid-{afid_num}', suffix='features.hkl',**{'desc':'{train_level}','space':'T1w'})
	out_models=bids_name(root=join(config['output_dir'],'derivatives','models'),prefix='afid-{afid_num}', suffix='model',**{'desc':'{train_level}','space':'T1w'})

rule all:
    input:
        expand(out_features, afid_num=afids_nums, train_level=train_levels),
        expand(out_models, afid_num=afids_nums, train_level=train_levels),
        #expand(bids_name(root=join(config['output_dir'],'derivatives','models'),prefix='afid-{afid_num}', suffix='model',**{'desc':'coarse'}), afid_num=afids_nums),

include: 'rules/modelTraining.smk'
