#!/usr/bin/python3

import sys
import os
import tempfile
import subprocess
import math
import time
import numpy as np
import pandas as pd
import nibabel as nib
from tempfile import mkdtemp
from nilearn.connectome import ConnectivityMeasure
from sklearn.preprocessing import StandardScaler
from nilearn import plotting
from matplotlib import pyplot as plt
from pandas import read_csv

global verbose
verbose = True
global timers
timers = []
def timer(command, name=""):
    msg = ""
    if command == "tic":
        global tic
        tic = time.perf_counter()
    elif command == "toc":
        toc = time.perf_counter()
        msg = "Time elapsed for {}: {} s".format(name, round(toc-tic, 2))
        printv("\n{}\n".format(msg))
        timers.append(msg)
    elif command == "report":
        print("\n------------------------------------------------------------")
        print("Timer report of {}:".format(os.path.basename(sys.argv[0])))
        for timer in timers:
            print("  {}".format(timer))
        print("------------------------------------------------------------\n")

def printv(msg, update=False):
    """Prints a message if verbose is True. If update is true, replace last line."""
    try: # verbose is not always defined
        verbose
    except NameError:
        pass
    else:
        if verbose:
            if update:
                print(msg, end="\r", flush=True)
            else:
                print(msg)

def create_dir_if_not_exist(directory):
    if not os.path.exists(directory):
        print("Creating directory {}...".format(directory))
        os.makedirs(directory)


def load_list_from_file(file_name):
    list = [line.rstrip('\n') for line in open(file_name)]
    return list


def load_time_series(file_list_fname):
    """Load time series from pre-specified file"""
    time_series = []
    file_list = load_list_from_file(file_list_fname)
    n = 1
    n_series = len(file_list)

    timer("tic")
    for file in file_list:
        printv("Loading time series from file {} ({}/{} | {} %)".format(file, n, n_series, round((n/n_series)*100, 2)), update=True)
        # Using Pandas' read_csv() speeds things up enormously in comparison to np.loadtxt
        ts = read_csv(file, sep='\s+')
        tsnp = np.array(ts) # Convert to numpy array
        time_series.append(tsnp)
        n += 1
    timer("toc", name="loading all time series")

    return time_series, file_list


def save_list_to_file(alist, fname):
    with open(fname, 'w') as f:
        f.write("\n".join(alist))


def save_whole_sample_nifti(matrices_array, output_path, output_fname=None, coi_indices=None, coi_fname=None, clean=True):
    """
    Take an array of matrices (e.g. output from ConnecitivityMeasure()),
    rearrange it to build a multi-subject NIfTI, e.g. for 361 ROIs and
    999 subjects: (999, 361, 361) -> (361, 361, 1, 999). Default filename is
    output_path/corr_matrices_full.nii.gz and can be overriden with output_file.

    If coi_indices are supplied, also save a COI matrix at
    output_path/coi_matrices.nii.gz (default) or coi_fname.

    NOTE: This is to be considered an outdated function and kept in only for reference.
    It's preferable to use cifti_dim_to_nifti_dim() and save_matrix(), which is a _lot_
    faster and cleaner.
    """

    n_subjects = matrices_array.shape[0]
    n_rows = matrices_array.shape[1]
    n_columns = matrices_array.shape[2]

    if coi_indices:
        coi_array = reduce_to_coi(matrices_array, coi_indices)
        if not coi_fname:
            coi_fname = os.path.join(output_path, "coi_matrices.nii.gz")

    if clean:
        print("Cleaning matrices before rearranging...")
        clean_matrices_array = np.array(np.empty(matrices_array.shape))
        subject = 0
        while subject < n_subjects:
            matrix_clean = clean_matrix(matrices_array[subject])
            clean_matrices_array[subject] = matrix_clean
            subject += 1
        matrices_array = clean_matrices_array
        if not output_fname:
            output_fname = os.path.join(output_path, "corr_matrices_clean.nii.gz")
    else:
        if not output_fname:
            output_fname = os.path.join(output_path, "corr_matrices_full.nii.gz")
    timer("tic")
    rearray = np.array(np.zeros((n_rows, n_columns, 1, n_subjects)))
    subject = 0
    while subject < n_subjects:
        print("Rearranging matrix of subject {}/{} ({} %).".format(subject+1, n_subjects, round((subject+1)/n_subjects*100, 2)), end="\r")
        row = 0
        while row < n_rows:
            column = 0
            while column < n_columns:
                rearray[row][column][0][subject] = matrices_array[subject][row][column]
                column += 1
            row += 1
        subject += 1
        sys.stdout.flush # Needed to enable same-line updates on console
    print("\nSaving group NIfTI to {}.".format(output_fname))
    affine = np.array([[-1., 0., 0., n_rows-1], [ 0., 1., 0., -0.], [0., 0., 1., -0.], [0., 0., 0., 1.]])
    create_dir_if_not_exist(output_path)
    save_matrix(rearray, output_fname, affine=affine)
    if coi_fname:
        print("\nCreating COI matrices...")
        coi_array = reduce_to_coi(rearray, coi_indices)
        print("Saving COI group NIfTI to {}.".format(coi_fname))
        save_matrix(coi_array, coi_fname, affine=affine)
    timer("toc", name="rearranging matrices and saving whole-sample NIfTI")


def save_matrix(matrix, fname, affine=np.eye(4)):
    """Create a NIfTI image from matrix and save it to a file."""
    image = nib.nifti1.Nifti1Image(matrix, affine=affine)
    create_dir_if_not_exist(os.path.dirname(fname))
    nib.save(image, fname)


def save_individual_matrices(matrices, subjects, output_dir, clean=False, pconn_dummy=False):
    """
    Take a list of (correlation) matrices, optionally clean them (if we are only interested
    in the first row and column) and save them as individual NIfTI (and, optionally CIFTI) files
    """
    n_matrices=len(matrices)
    if n_matrices != len(subjects):
        print("ERROR: Mismatch between number of matrices ({}) and number of subjects ({}).".format(n_matrices, len(subjects)))
        exit(1)
    n=0
    if clean:
        clean_matrices = [] # If we want clean matrices it makes sense to return a list of cleaned matrices
    if pconn_dummy:
        img_dummy = nib.load(pconn_dummy)
    create_dir_if_not_exist(output_dir)
    for matrix in matrices:
        fname=os.path.join(output_dir, "{}.nii".format(subjects[n].rstrip('.txt')))
        print("Saving matrix {}/{} to {} ({} %)".format(n+1, n_matrices, fname, round(((n+1)/n_matrices)*100, 2)))
        if clean:
            matrix_to_save = clean_matrix(matrix)
            clean_matrices.append(matrix_clean)
        else:
            matrix_to_save = matrix
        if pconn_dummy:
            img_new_fname = os.path.join(output_dir, "{}.pconn.nii".format(subjects[n]))
            img_new = nib.Cifti2Image(matrix, header=img_dummy.header, file_map=img_dummy.file_map)
            print("    - Saving pconn to {}".format(img_new_fname))
            img_new.to_filename(img_new_fname)
        save_matrix(matrix_to_save, fname)
        n += 1
    if clean:
        return clean_matrices
#  fname = os.path.join(output_path, "{}.nii".format(len(clean_matrices)))
#  print("Saving combined NIfTI file for all matrices (n = {}) to {}.".format(len(clean_matrices), fname))
#  image = nib.cifti2.Cifti2Image(clean_matrices, affine=np.eye(4))
#  nib.save(image, fname)


def compute_correlations(time_series, kind='partial correlation'):
    correlation_measure = ConnectivityMeasure(kind=kind)
    print("Fit-transforming time series...")
    timer("tic")
    correlation_matrices = correlation_measure.fit_transform(time_series)
    timer("toc", name="fitting all time series")
    return correlation_measure, correlation_matrices


def clean_matrix(matrix):
    tmp_matrix = matrix.copy()
    # Set all values but first row and first column to zero (our ROI)
    tmp_matrix[1:, 1:] = 0
    # Set correlation of ROI with itself to 0, too
    tmp_matrix[0, 0] = 0 # Actually not needed anymore due to np.fill_diagional later on
    return tmp_matrix


def make_temp_dir():
    tmp_base = "/tmp/roi_connectome"
    if not os.path.exists(tmp_base):
        os.makedirs(tmp_base)
    tmp_dir = mkdtemp(dir=tmp_base)
    return tmp_dir

def save_pconn(data, pconn_dummy, output_file):
    img_dummy = nib.load(pconn_dummy)
    img_new = nib.Cifti2Image(data, header=img_dummy.header, file_map=img_dummy.file_map)
    img_new.to_filename(output_file)

def plot_all(matrix, title, time_series, coords_file, labels_file, output_dir, clean=True, vmin=None, vmax=None, nan_matrix=False, pconn_dummy=False, pconn_fname='correlation_measure-mean_.pconn.nii'):
    # Before anything else, save matrix (i.e. corelation_measure.mean) as binary npy
    # file (e.g. to manually create pconn CIFTIs)
    create_dir_if_not_exist(output_dir)
    matrix_fname = os.path.join(output_dir, 'correlation_measure-mean_.npy')
    np.save(matrix_fname, matrix)
    # Build pconn file
    if pconn_dummy:
        pconn_fname = os.path.join(output_dir, pconn_fname)
        print("Saving pconn file to {}...".format(pconn_fname))
        save_pconn(matrix, pconn_dummy, pconn_fname)

    plot_title = title + ", n = {}".format(len(time_series))
    plot_title_orig = plot_title # We might modify plot_title later on
    coordinates = np.loadtxt(coords_file)
    labels = load_list_from_file(labels_file)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if nan_matrix is False:
        nan_matrix = matrix
    timer("tic")
    np.fill_diagonal(matrix, 0)
   # plotting.plot_matrix(nan_matrix, colorbar=True, figure=(40, 40), labels=labels, auto_fit=True, vmin=vmin, vmax=vmax)
   # plt.title(plot_title, fontsize=50)
   # fig_fname = "{}/correlation_matrix.svg".format(output_dir)
   # plt.savefig(fig_fname)
   # plt.clf()

    # We need to introduce this as clean_matrix does something unexpected with np.nan values
    # and plot_connectome specifically cannot deal with this
    if clean:
        print("Cleaning matrix ...")
        mean_roi_matrix = clean_matrix(matrix)
        plot_title = plot_title + " (clean)"
        fname_clean = "clean"
    else:
        mean_roi_matrix = matrix
        fname_clean = "nonclean"

    fig_fname = os.path.join(output_dir, "correlation_matrix_{}".format(fname_clean))
    print("Plotting matrix...")
    plotting.plot_matrix(mean_roi_matrix, colorbar=True, figure=(40, 40), labels=labels, auto_fit=True, vmin=vmin, vmax=vmax)
    plt.title(plot_title, fontsize=50)
    print("Saving plot to {}.svg...".format(fig_fname))
    plt.savefig("{}.svg".format(fig_fname))
    #plt.savefig("{}.png".format(fig_fname), dpi=1080) # PNG export does not seem to work in this case. Who knows..
    timer("toc", name="plotting and saving matrices")
    plt.clf()
    print("Plotting connectome ...")
    timer("tic")
    # Manually create new figure because for some reason plotting.plot_connectome() won't
    # accept figure size like plotting.plot_matrix()
    connectome_figure = plt.figure(figsize=(10, 5))
    # plot_connectome does not process np.nan values, we still have to pass np.nan_to_num(matrix) to plot_all function
    plotting.plot_connectome(np.nan_to_num(mean_roi_matrix), coordinates, figure=connectome_figure, colorbar=True,
                             node_size=30, title=plot_title, edge_vmin=vmin, edge_vmax=vmax)
    fig_fname = os.path.join(output_dir, "roi_connectome_{}".format(fname_clean))
    plt.savefig("{}.svg".format(fig_fname))
    plt.savefig("{}.png".format(fig_fname), dpi=1080)
    timer("toc", name="plotting and saving connectome")
    plt.clf()

    html_fname = os.path.join(output_dir, "roi_connectome_90_{}.html".format(fname_clean))
    print("Constructing interactive HTML connectome...")
    timer("tic")
    web_connectome = plotting.view_connectome(mean_roi_matrix, coordinates, edge_threshold="99%", node_size = 6, symmetric_cmap=False)
    html_fname = os.path.join(output_dir, "roi_connectome_{}.html".format(fname_clean))
    web_connectome.save_as_html(html_fname)
    timer("toc", name="plotting and saving HTML connectome")


def clean_palm_results(fname, labels_file, coords_file, alpha=0.95):
    nifti = nib.load(fname)
    data = nifti.get_data()
    labels = load_list_from_file(labels_file)
    coords = load_list_from_file(coords_file)

#  for column_id, column in enumerate(data[0]): # Only the first row is interesting to us
#    if column[0][0] <= alpha: # There is just one "subject" so non-iterating over the last index is fine
#      data[0][column_id][0] = 0 # NA might be better?
#      if column_id > 0:
#        labels[column_id] = ""
    data[data < alpha] = np.nan # Fancy indexing ftw
    nan_indeces = []
    for index, value in enumerate(data[0]): # Get NaN indices to clean labels/coords
        if np.isnan(value):
            if index > 0: # Retain ROI label
                labels[index] = ""
                # Effectively "remove" all NaN nodes by setting their coordinates to our ROI's
                # It works, there has to be a more elegant way, though
                coords[index] = coords[0]
#        coords[index] = "NA"
#  coords = [ coord for coord in coords if coord != "NA" ]
    print("Length of coordinates list: {}".format(len(coords)))
    tmp_dir = make_temp_dir()
    labels_out_fname = os.path.join(tmp_dir, 'clean_labels.txt')
    coords_out_fname = os.path.join(tmp_dir, 'clean_coords.txt')
    with open(labels_out_fname, 'w+') as f:
        for label in labels:
            f.write("%s\n" % label)
    with open(coords_out_fname, 'w+') as f:
        for label in coords:
            f.write("%s\n" % label)
    # Rearrange NIfTI matrix format to matrix format expected by nilearn
    palm_matrix = np.array(np.zeros((data.shape[3], data.shape[0], data.shape[1])))
    for row_id, row in enumerate(data):
        for column_id, column in enumerate(row):
            for subject_id, subject in enumerate(column):
                palm_matrix[subject_id][row_id][column_id] = subject
                palm_matrix[subject_id][column_id][row_id] = subject # has to be symmetric
    palm_matrix_lite = palm_matrix[~np.isnan(palm_matrix)]
    n_supra_values = len(palm_matrix_lite)
    print("Number of suprathreshold values: {}".format(n_supra_values))
    return palm_matrix, palm_matrix_lite, labels_out_fname, coords_out_fname, n_supra_values


def symmetrize(matrix, mirror_lower=False):
    """
    Return a symmetrized matrix.

    If mirror_lower is true, mirror the lower triangle to the upper triangle.
    This method works with matrices containg nan values, too.
    """
    if mirror_lower:
        matrix_symmetric = np.tril(matrix) + np.triu(matrix.T, 1)
    else:
        matrix_symmetric = matrix + matrix.T - np.diag(matrix.diagonal())
    return matrix_symmetric

def symmetrize_matrices(matrices, mirror_lower=False):
    """Helper function to symmetrize multiple matrices at once, e.g.
    symmetrize all 999 matrices in an array of shape (999, 513, 513)"""

    global verbose
    symm_matrices = np.zeros(matrices.shape)
    for i in range(0, matrices.shape[0]):
        printv("Symmetrizing matrix {} of {}...".format(i+1, matrices.shape[0]), update=True)
        symm_matrices[i] = symmetrize(matrices[i, :, :])

    return symm_matrices

def nifti_dim_to_cifti_dim(matrix):
    """Converts cols:rows:1:subjects to subjects:cols:rows"""
    matrix_cifti = matrix[:, :, 0, :] # Get rid of dimension #3
    matrix_cifti = np.moveaxis(matrix_cifti, -1, 0) # Put last dimension first

    return matrix_cifti


def cifti_dim_to_nifti_dim(matrix):
    """Converts subjects:cols:rows to cols:rows:1:subjects"""
    matrix_nifti = np.moveaxis(matrix, 0, -1) # Subjects:cols:rows -> cols:rows:subjects
    matrix_nifti = np.expand_dims(matrix_nifti, 2) # Cols:rows:subjects -> Colos:rows:1:subjects

    return matrix_nifti


def reduce_to_coi(matrices, indices):
    """
    Reduce the number of correlations in a matrices array to correlations of interest (COI).
    This is, for instance, useful to reduce the number of statistical tests we would
    have to correct for later on.

    Input: matrices to be "cleaned" and an iterable object of row/column indices to be preserved
    Returns: reduced matrices with all non-interesting values replaced with np.nan
    """
    from matplotlib.cbook import flatten

    # Leave the original matrix untouched
    matrices_clean = matrices.copy()

    indices_list = [] # Create dummy list to make iteration fool-proof
    indices_list.append(indices)
    indices_flattened = list(flatten(indices_list)) # Unnest nested iterables
    ndim = len(matrices.shape)
    nrow = matrices.shape[1] # Stricly speaking, this would be ncol for NIfTI style

    if ndim == 3: # Assume subjects:columns:rows ("CIFTI style")
        print("Assuming CIFTI style with {} rows.".format(nrow))
        for idx in range(0, nrow):
            matrices_clean[:, idx, idx+1:] = np.nan # Nan everything above the diagonal
            if idx not in indices_flattened:
                matrices_clean[:, :, idx] = np.nan # Nan everything except our COIs
    elif ndim == 4: # Assume columns:rows:1:subjects ("NIfTI style")
        print("Assuming NIfTI style with {} cols".format(nrow))
        for idx in range(0, nrow):
            matrices_clean[idx, idx+1:, :, :] = np.nan # Nan everything above the diagonal
            if idx not in indices_flattened:
                matrices_clean[:, idx, :, :] = np.nan # Nan everything else except our COIs
    else:
        raise ValueError("Number of dimensions of input matrix must be either 3 or 4.")

    return matrices_clean

def get_nimg_data(fname):
    data = nib.load(fname).get_fdata()
    return data


def wbc_cifti_convert_to_nifti(input_fname, output_fname, wb_command="wb_command"):
    subprocess.run([wb_command, "-cifti-convert", "-to-nifti", input_fname, output_fname], capture_output=True)

def wbc_parcellate_cifti(input_fname, parcellation_fname, output_fname, wb_command="wb_command"):
    result = subprocess.run([wb_command, "-cifti-parcellate", input_fname, parcellation_fname, "COLUMN", output_fname], capture_output=True, text=True)

    return result

def parcellate_ciftis(subj_ids=[], fnames=None, parcellation_fname=None, task=None, data_dir=None, runs=('LR', 'RL'), reg="MSMAll", wb_command="wb_command"):
    '''
    Function to batch-parcellate CIFTI files.
    Takes either a list of ids and assumes a standard directory structure
    (which then requires data_dir and task to be set) or a list of .dtseries.nii filenames.
    The former's output will mimic HCPPipelines output and MSMAll registration method is assumed,
    the latter's output filenames are the input filenames with their suffix replaced by .ptseries.nii
    If that is too rigid, use wbc_parcellate_cifi.
    '''

    if 'dlabel.nii' not in parcellation_fname:
        raise TypeError("parcellation_fname must be set to a .dlabel.nii file.")
    else:
        parcellation_name = os.path.split(parcellation_fname)[1].replace('.dlabel.nii', '')

    if fnames is None:
        fnames = []
    output_fnames = []
    missing_fnames = []
    task = task.upper()

    if len(reg) > 0:
        rege = "_{}".format(reg)
    else:
        rege = reg

    if "REST" in task:
        prefix = 'rfMRI'
        rege = "{}_hp2000_clean".format(rege)
    else:
        prefix = 'tfMRI'


    if len(subj_ids) > 1:
        if data_dir is None or task is None:
            raise TypeError("data_dir and task must be set if list of IDs are specified.")
        for id in subj_ids:
            for run in runs:
                fnames.append(os.path.join(data_dir, id, "MNINonLinear/Results/", "{}_{}_{}".format(prefix,task, run), "{}_{}_{}_Atlas{}.dtseries.nii".format(prefix, task, run, rege)))
                output_fnames.append(os.path.join(data_dir, id, "MNINonLinear/Results/", "{}_{}_{}".format(prefix, task, run), "{}_{}_{}_Atlas_hp200_s2_MSMAll_{}.ptseries.nii".format(prefix, task, run, parcellation_name)))

    else:
        if len(fnames) == 0:
            raise TypeError("Either set subj_ids to a list of IDs and data_dir to a data directory with the standard HCP data structure and task to a task name OR set fnames to a list of file names.")
        output_fnames = [ fname.replace(".dtseries.nii", ".ptseries.nii") for fname in fnames ]

    for idx in range(0, len(fnames)):
        input_fname = fnames[idx]
        output_fname = output_fnames[idx]
        print("Parcellating CIFTI {}/{}...".format(idx+1, len(fnames)))
        result = wbc_parcellate_cifti(input_fname, parcellation_fname, output_fname, wb_command=wb_command)
        if result.returncode > 1:
            print("Parcellation failed for file {}: {} {}".format(input_fname, result.stdout, result.stderr))
            missing_fnames.append(input_fname)

    if len(missing_fnames) > 0:
        print("\nParcellation failed for the following files:")
        for fname in missing_fnames:
            print("\t{}".format(fname))
    else:
        print("\nParcellation completed successfully for all files.")



def run_film_gls(input_fname, output_dir, design_fname, contrast_fname, film_gls="film_gls"):
    subprocess.run([film_gls, "--in={}".format(input_fname), "--rn={}".format(os.path.join(output_dir, "film_gls_output")),
                    "--pd={}".format(design_fname), "--con={}".format(contrast_fname), "--thr=1", "--mode=volumetric", "-v"])
    return os.path.join(output_dir, "film_gls_output", "prewhitened_data.nii.gz")


def prewhiten_cifti(cifti_fname, design_fname, contrast_fname, tmpdir, wb_command="wb_command", film_gls="film_gls"):
    fakenifti_fname = os.path.join(tmpdir.name, "fakenifti.nii.gz")
    # We could use singlesubject_cifti_to_fake_nifti(), but we have to save a temporary file anyway for film_gls and wb_command is certainly more robust
    wbc_cifti_convert_to_nifti(cifti_fname, fakenifti_fname, wb_command=wb_command)
    prewhitened_data_fname = run_film_gls(fakenifti_fname, tmpdir.name, design_fname, contrast_fname, film_gls=film_gls)
    return prewhitened_data_fname


def demean_and_standardize(array):
    scaler = StandardScaler()
    scaler.fit(array)
    demeaned_and_standardized_array = scaler.transform(array)
    return demeaned_and_standardized_array


def get_rest_timeseries(subj_ids, data_dir, sessions=('REST1', 'REST2'), runs=('LR', 'RL'), reg='MSMAll', parcellation_name='RenTianGlasser'):
    files_dict = {}
    ts_dict = {}

    if len(reg) > 0:
        rege = "_{}".format(reg)
    else:
        rege = reg

    for id in subj_ids:
        files_dict[id] = []

        for session in sessions:
            for run in runs:
                fname = os.path.join(data_dir, str(id), "MNINonLinear", "Results", "rfMRI_{}_{}".format(session, run), "rfMRI_{}_{}_Atlas_hp200_s2_MSMAll_{}.ptseries.nii".format(session, run, parcellation_name))
                files_dict[id].append(fname)

    n = 1
    N = len(files_dict)
    for id, fnames in files_dict.items():
        ts_dict[id] = []
        print("Processing {}... ({}/{} | {} %)".format(id, n, N, round(n/N*100, 2)))
        for fname in fnames:
            ts = demean_and_standardize(get_nimg_data(fname))
            for ts_item in ts:
                ts_dict[id].append(ts_item)
        n += 1

    return ts_dict


def get_ev_timeseries(ids_path_or_dict, ev_files, task=None, runs=None, parcellation=None, data_dir='/home/HCP/S1200/', tr=0.72, prewhiten=False, wb_command="wb_command", film_gls="film_gls"):
    """
    Gets event timeseries from parcellated CIFTI files (returned in form of a dictionary: {"ID": list(run1, run2, ...)}
    We first convert input to a unified dictionary structure to accomodate for different input types:
    1. List of IDs, which has as its prerequisite a directory structure based on HCP Pipeline's
       task prepare script. Here, we can make use of the accompanying design files to prewhiten
    2. A path containing files in format ID_TASK_RUN[_PARCELLATION].ptseries.nii (e.g.
       100307_GAMBLING_LR_RenTianGlasser.ptseries.nii or 100307_GAMBLING_LR.ptseries.nii)
    3. For custom file structures, we can just provide the dictionary in its unified format:
       {"ID": "RUN": "FILENAME"}, e.g.:
       {"100307": "LR": "./100307/LR.ptseries.nii"
                  "RL": "./100307/RL.ptseries.nii"}
    """
    import os
    files_dict = {}
    ev_data_dict = {}
    n = 1
    task = task.upper()

    if type(ids_path_or_dict) is str: # Assume path with files in format ID_TASK_RUN_PARCELLATION.ptseries.nii (e.g. 100307_GAMBLING_LR_RenTianGlasser.ptseries.nii)
        import glob
        file_list = glob.glob(os.path.join(ids_path_or_dict, "*.ptseries.nii"))
        for cifti_fname in file_list:
            components = os.path.basename(cifti_fname).split("_")
            id = components[0]
            task = components[1]
            run = components[2].split(".")[0] # in case parcellation is not part of filename
            if len(components) == 4:
                parcellation = components[4].split(".")
            if id not in files_dict:
                files_dict[id] = {run: cifti_fname}
            else:
                files_dict[id][run] = cifti_fname
        runs = 1 # Bogus value (i.e. not None) to pass checks
    elif type(ids_path_or_dict) is list: # Assume list of IDs and a directory structure from HCP Pipeline's task prepare script
        for id in ids_path_or_dict:
            for run in runs:
                run_path = os.path.join(data_dir, str(id), "MNINonLinear", "Results", "tfMRI_{}_{}".format(task, run))
                cifti_fname = os.path.join(run_path, "tfMRI_{}_{}_Atlas_hp200_s2_MSMAll_{}.ptseries.nii".format(task, run, parcellation))
                if id not in files_dict:
                    files_dict[id] = {run: cifti_fname}
                else:
                    files_dict[id][run] = cifti_fname
    elif type(ids_path_or_dict) is dict: # Assume unified dictionary structure, so nothing has to be prepared
        files_dict = ids_path_or_dict
        task = 1 # Set bogus values (i.e. not None) to pass checks
        runs = 1
    else:
        raise TypeError("First argument has to be either a list of IDs, a path or a dictionary.")

    for param in [task, runs]:
        if param is None:
            raise ValueError("{} must be specified (either explicitly or implicity through filename structure or preformatted dictionary).".format(param))

    N = len(files_dict.keys())
    for id, runs in files_dict.items():
        ev_data_dict[id] = [] # Create subject-specific list to store EV data in
        T = 0
        print("Processing {}... ({}/{} | {} %)".format(id, n, N, round(n/N*100, 2)))
        for run in runs:
            run_path = os.path.join(data_dir, str(id), "MNINonLinear", "Results", "tfMRI_{}_{}".format(task, run))
            cifti_fname = files_dict[id][run]
            if prewhiten:
                if parcellation is None:
                    raise ValueError("Parcellation needs to be specified in order for prewhitening process to find design files.")
                tmpdir = tempfile.TemporaryDirectory()
                feat_dir = os.path.join(run_path, "tfMRI_{}_{}_hp200_s2_level1_{}.feat".format(task, run, parcellation))
                design_fname = os.path.join(feat_dir, "design.mat")
                contrast_fname = os.path.join(feat_dir, "design.con")
                # Preprocessing: Prewhiten
                prewhitened_data_fname = prewhiten_cifti(cifti_fname, design_fname, contrast_fname, tmpdir, wb_command=wb_command, film_gls=film_gls)
                nifti_data = get_nimg_data(prewhitened_data_fname)
                # Downstream operations will be easier if we use timepoints:parcels shape instead of parcels:1:1:timepoints
                subject_data = nifti_data[:, 0, 0, :] # Fancy indexing gets rid of the superfluous middle dimensions -> parcels:timepoints
                subject_data = np.swapaxes(subject_data, 0, 1) # Finally swap axes -> timepoints:parcels
                tmpdir.cleanup() # Delete temporary directory
            else:
                subject_data = get_nimg_data(cifti_fname)

            # Demean and standardize (aka "fit-transform") the data to deal with multiple runs
            subject_data_std = demean_and_standardize(subject_data)

            for ev_file in ev_files:
                ev_fname = os.path.join(run_path, "EVs", ev_file)
                evs = np.loadtxt(ev_fname, ndmin=2) # ndmin=2 added for single-line EV txt files
                e = 1
                for ev in evs:
                    start = math.ceil((ev[0]/tr)+1) # We can afford to lose one timepoint at the beginning, it might even improve SNR
                    end = math.floor((((ev[0]+ev[1])/tr)+1)) # We can afford to lose one timepoint at the end, it might even improve SNR
                    ev_data = subject_data_std[start:end]
                    printv("    Timepoints in EV part {}/{} of {} (run {}): {}".format(e, len(evs), ev_file, run, len(ev_data)))
                    T+=len(ev_data)
                    for ev_data_item in ev_data: # Don't create nested lists, just append each item to the primary list
                        ev_data_dict[id].append(ev_data_item)
                    e+=1
        printv("    -> Total timepoints for subject {}: {}".format(id, T))
        n+=1
    return ev_data_dict


def save_data_dict(data_dict, path='./'):
    n = 1
    N = len(data_dict)
    fname_list = []
    os.makedirs(path, exist_ok=True)
    for id, data in data_dict.items():
        fname = "{}.txt".format(id)
        printv("Saving text file {}/{} ({} %)".format(n, N, round(n/N*100, 2)), update=True)
        np.savetxt(os.path.join(path, fname), data)
        fname_list.append(fname)
        n+=1
    with open(os.path.join(path,"ts_files"), 'w+') as f:
        for filename in fname_list:
            f.write("%s\n" % filename)


def singlesubject_cifti_to_fake_nifti(cifti_data):
    """
    Take a CIFTI-style data array and return it as a fake-NIfTI style one.
    """
    #newarray = np.array(np.zeros((cifti_data.shape[1], 1, 1, cifti_data.shape[0])))
    #for idx in range(0, len(cifti_data)):
    #    for idy in range(0, len(cifti_data[idx])):
    #        newarray[idy][0][0][idx] = cifti_data[idx][idy]
    newarray = np.transpose(cifti_data) # x:y -> y:x
    newarray = np.expand_dims(newarray, 1) # y:x -> y:1:x
    newarray = np.expand_dims(newarray, 1) # y:1:x -> y:1:1:x
    affine = [[1, 0, 0, 0],[0, 1, 0, 0],[0, 0, 1, 0],[0, 0, 0, 1]]
    img = nib.nifti1.Nifti1Image(newarray, affine=affine)

    return img

def subtract_arrays(positive, negative):
    """Takes two NumPy arrays 'positive' and 'negative' and returns the following:
    'positive - negative'. This was created to attain a combined PALM results file,
    where one contrast's p values are negative."""
    if isinstance(positive, str):
        positive = get_nimg_data(positive)
    if isinstance(negative, str):
        negative = get_nimg_data(negative)
    result = positive - negative
    return result

def plot_palm_new(palm_results, title, coords, labels, alpha=1.3, scale=False, n_best_values=5, lower_is_better=False, output_dir=None, pconn_dummy=False, pconn_fname='connectome.pconn.nii'):
    """
    Take PALM results (file or NumPy array) and plot its suprathreshold values (default
    threshold value: alpha=1.3) in a correlation matrix and as a connectome visualisation.

    Use lower_is_better=True if you use standard p values and not log p or 1-p
    """
    if isinstance(palm_results, str):
        # We assume the argument is a filename if it's a string
        data = get_nimg_data(palm_results)
        info = ["PALM results: {}".format(palm_results)] # This list holds all the informational messages we may later write into info.txt
    else:
        data = palm_results
        info = ["PALM results:"] # This list holds all the informational messages we may later write into info.txt
    # Get rid of superfluous dimensions
    adjmatrix = data[:, :, 0, 0]
    # Plot all p values
    fig_matrix = plotting.plot_matrix(adjmatrix, colorbar=True, figure=(40, 40), labels=labels, auto_fit=True).figure
    fig_connectome = plotting.plot_connectome(symmetrize(adjmatrix, mirror_lower=True), coords, colorbar=True, node_size=15, title=title)
    web_connectome = plotting.view_connectome(adjmatrix, coords, node_size = 6, symmetric_cmap=False)

    # Get highest/lowest indices and check if there are any suprathreshold values
    if lower_is_better:
        best_indices = get_extreme_indices(abs(adjmatrix), n_best_values, get_smallest=True)
        adjmatrix[adjmatrix == 0] = np.nan # See get_extreme_indices() for rationale
        nsupthr = len(adjmatrix[abs(adjmatrix) <= alpha])
    else:
        best_indices = get_extreme_indices(abs(adjmatrix), n_best_values)
        nsupthr = len(adjmatrix[abs(adjmatrix) >= alpha])
    msg = "{} best values: {}".format(n_best_values, adjmatrix[best_indices])
    info.append(msg)
    print(msg)
    if nsupthr == 0:
        msg = "\nNo values survive the threshold of {}.".format(alpha)
        info.append(msg)
        print(msg)
        fig_matrix_clean, fig_connectome_clean, web_connectome_clean = None, None, None
        labels_clean = ""
        coords_clean = ""
    else:
        msg = "Number of values to survive the threshold of {}: {}".format(alpha, nsupthr)
        info.append(msg)
        print(msg)
        # Purge matrix, coordinates and labels of subthreshold values
        adjmatrix_clean = adjmatrix.copy()
        if lower_is_better:
            adjmatrix_clean[abs(adjmatrix) > alpha] = np.nan
        else:
            adjmatrix_clean[abs(adjmatrix_clean) < alpha] = np.nan
        supthr_indices = np.argwhere(~np.isnan(adjmatrix_clean))
        labels_clean = ["" if i not in supthr_indices else x for i, x in enumerate(labels)]
        coords_clean = [[np.nan, np.nan, np.nan] if i not in supthr_indices else x for i, x in enumerate(coords)]

        fig_matrix_clean = plotting.plot_matrix(adjmatrix_clean, colorbar=True, figure=(40, 40), labels=labels_clean, auto_fit=True).figure
        fig_connectome_clean = plotting.plot_connectome(symmetrize(np.nan_to_num(adjmatrix_clean), mirror_lower=True), coords_clean, figure=plt.figure(figsize=(10, 5)), colorbar=True, node_size=30, title=title)
        web_connectome_clean = plotting.view_connectome(adjmatrix_clean, coords_clean, node_size = 6, symmetric_cmap=False)

    if output_dir:
        create_dir_if_not_exist(output_dir)
        print("Saving plots to {}...".format(output_dir))
        fig_matrix.savefig(os.path.join(output_dir, "matrix.svg"))
        if fig_matrix_clean:
            fig_matrix_clean.savefig(os.path.join(output_dir, "matrix_clean.svg"))
        fig_connectome.savefig(os.path.join(output_dir, "connectome.svg"))
        if fig_connectome_clean:
            fig_connectome_clean.savefig(os.path.join(output_dir, "connectome_clean.svg"))
        web_connectome.save_as_html(os.path.join(output_dir, "web_connectome.html"))
        if web_connectome_clean:
            web_connectome_clean.save_as_html(os.path.join(output_dir, "web_connectome_clean.html"))

        info.append("")
        info.append("Top 5 connections according to p value:")
        info.append("Node A (label) <-> Node B (label): p value")
        for i in range(5):
            x = best_indices[0][i]
            y = best_indices[1][i]
            msg = "{} ({}) <-> {} ({}): {}".format(x+1, labels[x], y+1, labels[y], adjmatrix[x,y])
            info.append(msg)
        save_list_to_file(info, os.path.join(output_dir, "info.txt"))

        if labels_clean:
            save_list_to_file(labels_clean, os.path.join(output_dir, "labels_clean"))
            np.savetxt(os.path.join(output_dir, "coords_clean"), coords_clean)

        if pconn_dummy:
            save_pconn(adjmatrix, pconn_dummy, os.path.join(output_dir, pconn_fname))

    for entry in info:
        print(entry)
    return fig_matrix, fig_matrix_clean, fig_connectome, fig_connectome_clean, web_connectome, web_connectome_clean, labels_clean, coords_clean


def plot_nodes(node_list, parcellation_img_path, bg_img_path, plot_degree=False, n_slices=5, display_modes=['x', 'y', 'z'], cmap='gist_ncar', normalize=True, incremental_values=False, colorbar=False, **plot_roi_kwargs):
    """
    Takes a list of nodes (or a networkx degrees object) and plots them on
    a specified background image.
    You can disable actual plodding by setting bg_img_path to 'None' if all you want is a NIfTI image.
    Set plot_degree to True to use degrees instead of node number als voxel values (useful to create heatmap-like
    plots, it makes sense to change cmap in that case to something like "hot" or "cool").

    Returns:
        - a dictionary of figures corresponding to the display modes specified
        - an interactive viewer of the plot, which can be saved with viewer.to_html()
        - a NIfTI image with two volumes: first volume's voxels' values are set to the node degree
                                          second volume's voxel's values equal are set to the node number
    """
    import nibabel as nib

    if type(node_list) == pd.core.frame.DataFrame: # Assume DataFrame from get_top_n_nodes()
        degree_list = node_list['Degree'].to_list()
        node_list = node_list['Node'].to_list()

    elif type(node_list[0]) == tuple: # Assume degree list instead of node list and convert
        node_list = [x[0] for x in node_list]
        degree_list = None

    if normalize:
        # The following normalization was only introduced to deal with plotting.view_img()'s
        # weird fixation with symmetricity and therefore coloring inconsistent with
        # plotting_plot_roi()'s, which did just fine before normalization. Anyway, results
        # are identical.
        normalizer = get_mpl_normalizer(node_list)
        value_list = [normalizer(node) for node in node_list]
        vmin = 0 # plotting.plot_roi determines these wrong for some reason
        vmax = 1
    elif incremental_values:
        value_list = [n for n in range(len(node_list))]
        vmin = None
        vmax = None
    else:
        value_list = node_list
        vmin = None
        vmax = None

    parcellation_img = nib.load(parcellation_img_path)
    parcellation_img_data = parcellation_img.get_fdata()

    new_data_plotting = np.empty(parcellation_img.shape) # Additional array in case of normalization
    new_data_plotting[:, :] = np.nan # Replace all values with NaN to achieve transparency in NiFTI viewers, e.g. fsleyes
    new_data_degree = new_data_plotting.copy()
    new_data_index = new_data_plotting.copy()

    for idx in range(len(node_list)):
        # In a parcellation image, parcels = nodes are represented as neighbouring
        # voxels with values corresponding to the parcel/node number, for plot_roi()
        # and view_img() to produce consistent colors we need to use normed parcel
        # values as voxel values (see above comment)
        if degree_list is not None:
            new_data_degree[parcellation_img_data == node_list[idx]] = degree_list[idx]
        if plot_degree is True:
            new_data_plotting = new_data_degree.copy()
            vmax = None
            vmin = None
        else:
            new_data_plotting[parcellation_img_data == node_list[idx]] = value_list[idx]
        new_data_index[parcellation_img_data == node_list[idx]] = node_list[idx]

    new_img_plotting = nib.Nifti1Image(new_data_plotting, parcellation_img.affine, parcellation_img.header)

    if bg_img_path is not None:
        bg_img = nib.load(bg_img_path)
        figures_dict = {}
        for display_mode in display_modes:
            plotting.plot_roi(new_img_plotting, cut_coords=n_slices, display_mode=display_mode, bg_img=bg_img, black_bg=False, threshold=0.0, cmap=cmap, vmin=vmin, vmax=vmax, colorbar=colorbar, **plot_roi_kwargs)
            figures_dict[display_mode] = plt.gcf()
            plt.close()
        viewer = plotting.view_img(new_img_plotting, bg_img=bg_img, black_bg=False, colorbar=colorbar, cmap=cmap, symmetric_cmap=False, resampling_interpolation='nearest', threshold=0.0)
    else:
        figures_dict = None
        viewer = None

    # Create NiFTI image with two volumes
    new_data_degree = np.expand_dims(new_data_degree, 3) # Expand array to hold the volume dimension
    new_data_index = np.expand_dims(new_data_index, 3) # Expand array to hold the volume dimension
    new_data_concat = np.concatenate((new_data_degree, new_data_index), axis=3) # Concat expanded images to create two-volume image

    new_img_volumes = nib.Nifti1Image(new_data_concat, parcellation_img.affine, parcellation_img.header)

    return figures_dict, viewer, new_img_volumes


def get_mpl_normalizer(list_of_numbers):
    """
    Simple wrapper function returning a Matplotlib normalizer for a list
    of numbers, e.g. nodes for consistent colors across plots
    """
    import matplotlib as mpl
    normalizer = mpl.colors.Normalize(vmin=np.min(list_of_numbers), vmax=np.max(list_of_numbers))

    return normalizer


def add_color_to_nodes_table(nodes_df, cmap='gist_ncar', normalize=True, incremental_values=False):
    """
    Input: Dataframe with 'Node' column containing node numbers and Matplotlib color map
           (defaults to 'gist_ncar' which is used by nilearn's plot functions by default)
    Output: Dataframe with added 'Color' column as Pandas styler, which can then be used
            to display and/or save an informative table of nodes, their labels and colors
            corresponding to their plotted version
    """
    import matplotlib as mpl
    colormap = mpl.colormaps[cmap]
    def color_nodes(df):
        def tocolor(node):
            if normalize:
                node = normalizer(node)
            cmap_color = mpl.colors.to_hex(colormap(node))
            return "background-color: {}".format(cmap_color)
### TODO implement if normalize etc.. remove tocolor
        if normalize:
            normalizer = get_mpl_normalizer(df['Node'])
        df['Color'] = df['Node'].apply(tocolor)
        df['Node'] = '' # We need to return something for styler, original content won't work

        return df

    nodes_df_color = nodes_df.copy()
    nodes_df_color['Color'] = '' # Add new column
    styler = nodes_df_color.style.apply(color_nodes, axis=None, subset=["Color", "Node"])

    return styler


def get_extreme_indices(array, n, get_smallest=False):
    """
    Return the indices of the n highest/lowest values from an array.
    Multidimensional arrays are flattened and the indices unravelled.

    From https://stackoverflow.com/a/38884051
    """
    array_flattened = array.flatten()
    if get_smallest:
        # Here we need to remove all 0 values otherwise they tend to be the smallest
        # We will leave negative values alone as some software might use negative p values
        # for e.g. inverse contrasts
        array_flattened[array_flattened == 0] = np.nan
        indices = np.argpartition(array_flattened, n)[:n]
        indices = indices[np.argsort(array_flattened[indices])]
    else:
        indices = np.argpartition(array_flattened, -n)[-n:]
        indices = indices[np.argsort(-array_flattened[indices])]

    return np.unravel_index(indices, array.shape)


def get_top_n_edges(matrix, top_n):
    matrix = abs(matrix) # Let this function work for both pos and neg tails
    top_n_indices = np.unravel_index(np.argpartition(-matrix.flatten(), range(top_n))[:top_n], matrix.shape)
    top_n_values = matrix[top_n_indices]

    return top_n_indices, top_n_values


def get_safe_path(path_base, ext):
    path_n = 0
    if os.path.exists(path := path_base + "{}".format(ext)):
        while os.path.exists(path := path_base + "-{}{}".format(path_n, ext)):
            path_n += 1
            # Check if file is writable. CAVE: Might change during execution but overall should work
    with open(path, 'w') as f: pass # Throws error if path is not writable
    os.remove(path) # File is created with open(..., 'w') -> delete it.
    printv("File to save data to: {}".format(path))
    return path


def safesave(object_to_save, path_base):
    """
    Determines appropriate extension for object to be saved and
    iterates integer suffix if suggested file is already present
    """
    if isinstance(object_to_save, dict) or isinstance(object_to_save, np.ndarray):
        path = get_safe_path(path_base, ".npy")
        np.save(path, object_to_save)
    elif isinstance(object_to_save, pd.core.frame.DataFrame):
        path = get_safe_path(path_base, ".feather")
        from pyarrow import feather
        feather.write_feather(object_to_save, path)
    elif isinstance(object_to_save, plt.matplotlib.axes.Subplot):
        path = get_safe_path(path_base, ".svg")
        object_to_save.get_figure().savefig(path)
        if 'perm_info' in dir(object_to_save): # In case of permutation results, save permutation info as separate files
            import re
            path = re.sub('.svg$', '.txt', path)
            printv("Saving permutation info to {}".format(path))
            with open(path, "w") as info_file:
                print(object_to_save.perm_info, file=info_file)
    elif isinstance(object_to_save, plt.matplotlib.figure.Figure):
        path = get_safe_path(path_base, ".svg")
        object_to_save.savefig(path)
    else:
        print("Type of data object not recognized: {}".format(type(object_to_save)))


def reduce_matrix_by_threshold(matrix, threshold=0, labels=None, coords=None):
    """
    Takes a correlation/connectivity matrix (e.g. r_mat) and returns a matrix
    reduced to nodes with suprathreshold edges (as absolute values). Useful e.g. for
    drawing less-convoluted plots for thresholded matrices.
    If a list of labels is provided, a labels dictionary mapping the new (after reducing)
    node number to the corresponding label is returned. This dictionary can be used
    with set_node_attributes() from networkx.
    If a list of coordinates is provided, a list of coordinates will be returned
    corresponding to the reduced matrix's nodes.
    Example input: Matrix with shape (513, 513) -> returns a dictionary with
    key 'reduced_matrix' with a matrix of shape (56, 56) and a key with 'original_index'
    for access to the list of the original node numbers.
    """

    results_dict = {}

    # Create empty dataframe with reduced nodes
    indices = np.where(matrix > threshold)
    df = pd.DataFrame(index=np.sort(list(set(indices[0]))), columns=np.sort(list(set(indices[1]))))

    indices_2d = np.argwhere(matrix > threshold)

    for idx in range(len(indices_2d)):
        r = indices_2d[idx][0]
        c = indices_2d[idx][1]
        df.loc[r, c] = matrix[r, c]

    a = df.to_numpy(na_value=0, dtype=float)

    results_dict['reduced_matrix'] = a
    results_dict['original_index'] = list(df.index)

    if labels is not None or coords is not None:
        labels_clean = {}
        coords_clean = []
        for n in range(len(df.index)):
            original_index = df.index[n]
            if labels is not None:
                labels_clean[n] = labels[original_index]
            if coords is not None:
                coords_clean.append(coords[original_index])
        if labels is not None:
            results_dict['reduced_labels'] = labels_clean
        if coords is not None:
            results_dict['reduced_coords'] = coords_clean

    return results_dict


def download_hcp_files(subj_ids, data_dir, tasks, aws_access_key_id, aws_secret_access_key, reg="MSMAll", runs=('LR', 'RL')):

    from cloudpathlib import S3Client
    from cloudpathlib.exceptions import CloudPathNotExistsError

    if type(tasks) == str:
        task = tasks
        tasks = [ task ]

    for task in tasks:
        # rsfMRI consists of two separate runs and requires a different prefix
        if task.upper() == "REST":
            tasks[tasks.index(task)] = "REST1" # Replace "rest" or "REST" with REST1
            tasks.append("REST2")
        task = task.upper() # Convert task names to upper-case

    os.makedirs(data_dir, exist_ok=True) # Create data directory if it does not already exist

    # Instantiate client with credentials
    s3client = S3Client(aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)

    s1200 = s3client.CloudPath("s3://hcp-openaccess/HCP_1200") # Open the S1200 release bucket

    missing_subjs = []
    n_subj = 1
    N_subj = len(subj_ids)
    for subj_id in subj_ids:
        n_task = 1
        local_results_path = os.path.join(data_dir, subj_id, 'MNINonLinear/Results')
        remote_results_path = s1200 / subj_id / 'MNINonLinear/Results/'

        print("Downloading files for subject {} ({}/{} | {} %)...".format(subj_id, n_subj, N_subj, round(n_subj/N_subj*100, 2)))
        for task in tasks:
            print("\tDownloading files of {} task ({}/{})...".format(task.lower(), n_task, len(tasks)), end="", flush=True)

            if len(reg) > 0:
                rege = "_{}".format(reg)
            else:
                rege = reg

            if "REST" in task:
                prefix = 'rfMRI'
                rege = "{}_hp2000_clean".format(rege)
            else:
                prefix = 'tfMRI'

            local_task_path = os.path.join(local_results_path, "{}_{}/".format(prefix, task))

            for run in ['LR', 'RL']: # for every task, there are two fMRI runs
                run_dir = "{}_{}_{}".format(prefix, task, run)
                remote_run_path = remote_results_path / run_dir
                ev_dir = remote_run_path / 'EVs'

                fmri_file = remote_run_path / "{}_Atlas{}.dtseries.nii".format(run_dir, rege)
                local_run_path = os.path.join(local_results_path, run_dir)
                os.makedirs(local_run_path, exist_ok=True)

                try:
                    fmri_file.download_to(local_run_path)
                    if prefix == 'tfMRI':
                        ev_dir.download_to(os.path.join(local_run_path, 'EVs'))
                except CloudPathNotExistsError:
                    print(" FAILED for object {}".format(fmri_file))
                    if subj_id not in missing_subjs:
                        missing_subjs.append(subj_id)
                else:
                    if run == runs[-1]:
                        print(" OK") # print "OK" after last run

            n_task += 1
        n_subj += 1

    if len(missing_subjs) > 0:
        print("\nDownload failed for the following subjects due to missing remote files: {}".format(missing_subjs))
    else:
        print("\nDownload complete.")

    return missing_subjs
