"""
Average variance explained across runs for regression models.
"""
import numpy as np
import nibabel as nib

def avg_runs(model_type, sub, total_run, filepath_mask2, results_save_dir):
    var_data_total = []
    
    for testrun in range(1, total_run+1):
        data_dir = results_save_dir+sub+'_var_expl_map_'+model_type+'_testrun'+str(testrun)+'.nii.gz'
        var_map = nib.load(data_dir) 
        var_data = var_map.get_fdata()
        var_data_total.append(var_data)
    
    var_data_avg = np.mean(var_data_total, 0)
    var_map_affine = var_map.affine 
    var_img_avg = nib.Nifti1Image(var_data_avg, var_map_affine)
    nib.save(var_img_avg, results_save_dir+sub+'_var_expl_map_'+model_type+'_avgruns.nii.gz')
    
