#! /usr/bin/env python

import os
import tempfile
import shutil
import subprocess
from rabies.utils import generate_token_data
import sys
import argparse

'''PARAMETERS NOT TESTED
preprocess:
    --bids_filter: requires creating a JSON
    --log_transform: makes the token data crash

confound_correction:
    --highpass/lowpass/edge_cutoff; since filtering doesn't work with 3 timepoints

analysis:
    --ROI_csv: no such file for token data
    --ROI_type parcellated: no parcellation for token data
    --seed_prior_list: not available for token data
    --outlier_threshold: not necessary
'''


def get_parser():
    """Build parser object"""
    parser = argparse.ArgumentParser(
        description=
            "Parser to handle testing using token data.",
        formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "--complete", dest='complete', action='store_true',
        help=
            "Run a complete testing of the pipeline."
        )
    parser.add_argument(
        "--test_registration", dest='test_registration', action='store_true',
        help=
            "Test whether registration steps are functional. This is not part of the main series of tests, since this takes much longer."
        )
    parser.add_argument(
        '--output_dir', action='store', type=str,
        help=
            "Provide an output directory instead of using a temporary directory.\n"
            "This prevents the deletion of outputs.\n"
        )
    parser.add_argument(
        "--custom", dest='custom', type=str, default=None,
        help=
            "Provide a custom command to run, without specifying input/output directories.\n "
            "The path to token data and output folder is appended at the end of --custom, \n"
            "then the command is run, and the script is exited. The command must be provided \n"
            "as a string between ''.\n"
            "If the preprocessing stage is run, the following arguments are automatically \n"
            "provided to ensure compatibility with token data:\n"
            "   --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \ \n"
            "   --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz  \ \n"
            "   --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz  \ \n"
            "   --vascular_mask {tmppath}/inputs/token_mask.nii.gz \ \n"
            "   --bold2anat_coreg registration=no_reg,masking=false,brain_extraction=false,keep_mask_after_extract=false,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 \ \n"
            "   --commonspace_reg masking=false,brain_extraction=false,keep_mask_after_extract=false,fast_commonspace=true,template_registration=no_reg,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 \n"
        )
    
    return parser

parser = get_parser()
opts = parser.parse_args()

if opts.output_dir is None:
    tmppath = tempfile.mkdtemp()
else:
    tmppath = opts.output_dir

generate_token_data(tmppath, number_scans=3)

if not opts.custom is None:
    minimal_preproc = f"rabies --interpolation BSpline3 --inclusion_ids {tmppath}/inputs/sub-token1_bold.nii.gz --verbose 1 --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
        --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
        --bold2anat_coreg registration=no_reg,masking=false,brain_extraction=false,keep_mask_after_extract=false,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 --commonspace_reg masking=false,brain_extraction=false,keep_mask_after_extract=false,fast_commonspace=true,template_registration=no_reg,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995"
    minimal_cc = f"rabies --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs"

    command = opts.custom
    if 'preprocess' in command:
        command += f" --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
    --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
    --bold2anat_coreg registration=no_reg,masking=false,brain_extraction=false,keep_mask_after_extract=false,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 --commonspace_reg masking=false,brain_extraction=false,keep_mask_after_extract=false,fast_commonspace=true,template_registration=no_reg,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995"
        command += f" {tmppath}/inputs {tmppath}/outputs"

    if 'confound_correction' in command or 'analysis' in command:
        if not os.path.isfile(f'{tmppath}/outputs/rabies_preprocess_workflow.pkl'):
            # provide preprocess outputs to run cc stage
            process = subprocess.run(
                minimal_preproc,
                check=True,
                shell=True,
                )
        if ' analysis ' in command:
            if not os.path.isfile(f'{tmppath}/outputs/rabies_confound_correction_workflow.pkl'):
                # provide cc outputs to run analysis stage
                process = subprocess.run(
                    minimal_cc,
                    check=True,
                    shell=True,
                    )
            command += f" --prior_bold_idx 0 1 --prior_confound_idx 0 1 --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --ROI_labels_file {tmppath}/inputs/token_labels.nii.gz"
        command += f" {tmppath}/outputs {tmppath}/outputs"

    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )
    sys.exit()


if opts.test_registration:
    ####TESTING REGISTRATION####

    # testing all registration/modelbuild steps on 2 scans
    command = f"rabies -p MultiProc --exclusion_ids {tmppath}/inputs/sub-token3_bold.nii.gz --force --verbose 1 --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
        --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
        --commonspace_reg stages=rigid,fast_commonspace=false,template_registration=Rigid --bold2anat_coreg registration=Rigid --nativespace_resampling 1x1x1 --commonspace_resampling 1x1x1 --anatomical_resampling 0.4x0.4x0.4 \
        --anat_robust_inho_cor apply=true,stages=rigid,template_registration=Rigid --bold_robust_inho_cor apply=true,stages=rigid,template_registration=Rigid"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # testing only inhomogeneity correction
    command = f"rabies -p MultiProc --exclusion_ids {tmppath}/inputs/sub-token2_bold.nii.gz {tmppath}/inputs/sub-token3_bold.nii.gz --force --verbose 1 --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=Rigid,otsu_thresh=2,multiotsu=false --bold_inho_cor method=Rigid,otsu_thresh=2,multiotsu=false \
        --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
        --commonspace_reg fast_commonspace=true,template_registration=Rigid --bold2anat_coreg registration=no_reg --nativespace_resampling 1x1x1 --commonspace_resampling 1x1x1 --anatomical_resampling 0.4x0.4x0.4"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    if opts.output_dir is None:
        shutil.rmtree(tmppath)
    
    sys.exit()


# this function is to repeat certain operations that are subject to random Singular matrix errors
def repeat_attempts(command, number_attempts=3):
    for attempt in range(1, number_attempts + 1):
        try:
            process = subprocess.run(
                command,
                check=True,
                shell=True,
                )
            break  # Success! Exit the loop
        except Exception as e:
            print(f"Error: {e}")
            if attempt == number_attempts:
                raise ValueError("All retries failed. Crashing now.")
            

command = f"rabies --inclusion_ids {tmppath}/inputs/sub-token1_bold.nii.gz --verbose 1 --force --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
    --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
    --bold2anat_coreg registration=no_reg,masking=true,brain_extraction=true,keep_mask_after_extract=false,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 --commonspace_reg masking=true,brain_extraction=true,keep_mask_after_extract=false,fast_commonspace=true,template_registration=no_reg,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 \
    --HMC_level 2 --no_HMC --apply_despiking --anat_autobox --bold_autobox --oblique2card affine --resampling_space both --hmc_qc_report apply=true,fps=5"
process = subprocess.run(
    command,
    check=True,
    shell=True,
    )

command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs --nuisance_regressors aCompCor_5 --nativespace_analysis --resample_to_commonspace"
process = subprocess.run(
    command,
    check=True,
    shell=True,
    )

command = f"rabies --force --verbose 1 --data_type int16 analysis {tmppath}/outputs {tmppath}/outputs --data_diagnosis --seed_list ACA_seed {tmppath}/inputs/token_mask_half.nii.gz --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --prior_bold_idx 0 1 --prior_confound_idx 0 1 --group_ica apply=true,dim=0,random_seed=1"
process = subprocess.run(
    command,
    check=True,
    shell=True,
    )

if opts.complete:
    command = f"rabies --exclusion_ids {tmppath}/inputs/sub-token2_bold.nii.gz {tmppath}/inputs/sub-token3_bold.nii.gz --force --verbose 1 --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
        --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask.nii.gz --CSF_mask {tmppath}/inputs/token_mask.nii.gz --vascular_mask {tmppath}/inputs/token_mask.nii.gz \
        --commonspace_reg fast_commonspace=true,template_registration=no_reg --bold_only --detect_dummy \
        --apply_STC --args_3dTshift='-linear -tpattern alt-z' --nativespace_resampling 1x1x1 --commonspace_resampling 1x1x1 --anatomical_resampling 1x1x1 --oblique2card 3dWarp --resampling_space both --bold_nativespace"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs --nativespace_analysis --resample_to_commonspace --slicewise_correction_direction AP --smoothing_filter 0.3"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    ####CONFOUND CORRECTION####
    command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs \
        --generate_CR_null --TR 1 --scale_variance_voxelwise \
        --smoothing_filter 0.3 --detrending order=2,time_interval=3-10 --image_scaling global_variance "
    repeat_attempts(command, number_attempts=3)

    # testing censoring on its own, since it removes all scans and prevent further testing
    command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs --frame_censoring FD_censoring=true,FD_threshold=0.05,DVARS_censoring=true,minimum_timepoint=3 --timeseries_interval 2-12"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # testing AROMA on its own to retain degrees of freedom
    command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs --ica_aroma apply=true,dim=2,random_seed=1"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # testing --nuisance_regressors on its own to retain degrees of freedom
    command = f"rabies --force --verbose 1 --data_type int16 confound_correction --read_datasink {tmppath}/outputs {tmppath}/outputs \
        --nuisance_regressors mot_24 aCompCor_percent global_signal"
    repeat_attempts(command, number_attempts=3)

    ####ANALYSIS####
    command = f"rabies --force --verbose 1 analysis {tmppath}/outputs {tmppath}/outputs --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --prior_bold_idx 0 1 --prior_confound_idx 0 1 \
        --network_weighting relative --optimize_NPR apply=true,window_size=2,min_prior_corr=0.5,diff_thresh=0.03,max_iter=5,compute_max=false \
        --FC_matrix --ROI_type parcellated --ROI_labels_file {tmppath}/inputs/token_labels.nii.gz"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    ####GROUP LEVEL, RUNNING ALL 3 SCANS####
    command = f"rabies --force --verbose 1 --data_type int16 preprocess {tmppath}/inputs {tmppath}/outputs --anat_inho_cor method=disable,otsu_thresh=2,multiotsu=false --bold_inho_cor method=disable,otsu_thresh=2,multiotsu=false \
        --anat_template {tmppath}/inputs/sub-token1_T1w.nii.gz --brain_mask {tmppath}/inputs/token_mask.nii.gz --WM_mask {tmppath}/inputs/token_mask_half.nii.gz --CSF_mask {tmppath}/inputs/token_mask_half.nii.gz --vascular_mask {tmppath}/inputs/token_mask_half.nii.gz \
        --bold2anat_coreg registration=no_reg,masking=false,brain_extraction=false,keep_mask_after_extract=false,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 --commonspace_reg masking=false,brain_extraction=false,keep_mask_after_extract=false,fast_commonspace=true,template_registration=no_reg,winsorize_lower_bound=0.005,winsorize_upper_bound=0.995 \
        --HMC_level 1 --resampling_space common_only"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    command = f"rabies --force --verbose 1 --data_type int16 confound_correction {tmppath}/outputs {tmppath}/outputs --nuisance_regressors mot_6"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # testing group level --data_diagnosis
    command = f"rabies --force --verbose 1 analysis {tmppath}/outputs {tmppath}/outputs --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --prior_bold_idx 0 1 --prior_confound_idx 0 1 --NPR_temporal_comp 1 \
        --data_diagnosis --group_avg_prior --extended_QC --DR_ICA --seed_list {tmppath}/inputs/token_mask_half.nii.gz"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # test for the scan QC thresholds
    scan_QC="'{DR:{Dice:[0.5],Conf:[0.1],Amp:true},SBC:{Dice:[0.3]}}'"
    command = f"rabies --force --verbose 1 --data_type int16 analysis {tmppath}/outputs {tmppath}/outputs --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --prior_bold_idx 0 --prior_confound_idx 0 1 \
        --data_diagnosis --extended_QC --scan_QC_thresholds {scan_QC}"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )

    # test group ICA
    command = f"rabies --force --verbose 1 --data_type int16 analysis {tmppath}/outputs {tmppath}/outputs --prior_maps {tmppath}/inputs/melodic_networks.nii.gz --prior_bold_idx 0 1 --prior_confound_idx 0 1 --group_ica apply=true,dim=0,random_seed=1"
    process = subprocess.run(
        command,
        check=True,
        shell=True,
        )


if opts.output_dir is None:
    shutil.rmtree(tmppath)
