import numpy as np  # A library for working with numerical data
from nilearn.maskers import NiftiMasker  # A class for masking and denoising fMRI data
import pandas as pd  # A library for working with tabular data

# Files from fMRIPrep
data_file = "/lustre04/scratch/javan/tedana_test/fmriprep/sub-MTL0002/ses-FU84/func/sub-MTL0002_ses-FU84_task-rest_run-01_space-MNI152NLin2009cAsym_res-2_trimmed_desc-preproc_bold.nii.gz"
mask_file = "/lustre04/scratch/javan/tedana_test/fmriprep/sub-MTL0002/ses-FU84/func/sub-MTL0002_ses-FU84_task-rest_run-01_space-MNI152NLin2009cAsym_res-2_desc-brain_mask.nii.gz"
#confounds_file = "sub-01_task-rest_desc-confounds_timeseries.tsv"

# Files from tedana (after running on fMRIPrepped data)
mixing_file = "/lustre04/scratch/javan/tedana_test/fmriprep/tedana/sub-MTL0002_ses-FU84/sub-MTL0002_ses-FU84_task-_space-Native_desc-ICA_mixing.tsv"
metrics_file = "/lustre04/scratch/javan/tedana_test/fmriprep/tedana/sub-MTL0002_ses-FU84/sub-MTL0002_ses-FU84_task-_space-Native_desc-tedana_metrics.tsv"
confounds_file = "/lustre04/scratch/javan/tedana_test/fmriprep/sub-MTL0002/ses-FU84/func/sub-MTL0002_ses-FU84_task-rest_run-01_desc-confounds_timeseries.tsv"
# Load the mixing matrix
mixing_df = pd.read_table(mixing_file)  # Shape is time-by-components

# Load the component table
metrics_df = pd.read_table(metrics_file)
rejected_columns = metrics_df.loc[metrics_df["classification"] == "rejected", "Component"]
accepted_columns = metrics_df.loc[metrics_df["classification"] == "accepted", "Component"]

# Select "bad" components from the mixing matrix
rejected_components = mixing_df[rejected_columns].to_numpy()
accepted_components = mixing_df[accepted_columns].to_numpy()

# Combine the rejected components and the fMRIPrep confounds into a single array
regressors = rejected_components

masker = NiftiMasker(
    mask_img=mask_file,
    standardize_confounds=True,
    standardize=False,
    smoothing_fwhm=None,
    detrend=False,
    low_pass=None,
    high_pass=None,
    t_r=None,  # This shouldn't be necessary since we aren't bandpass filtering
    reports=False,
)
# Denoise the data by fitting and transforming the data file using the masker
denoised_img_2d = masker.fit_transform(data_file,confounds=regressors)

# Transform denoised data back into 4D space
denoised_img_4d = masker.inverse_transform(denoised_img_2d)

# Save to file
denoised_img_4d.to_filename(
    "/lustre04/scratch/javan/tedana_test/fmriprep/tedana/sub-MTL0002_ses-FU84sub-MTL0002_ses-FU84_task-rest_run-01_space-MNI152NLin2009cAsym_res-2_trimmed_desc-aggrDenoised_bold.nii.gz"
)