from os.path import join as opj
import pandas as pd  # A library for working with tabular data

# Files from fMRIPrep
data_file = "sub-04570_task-rest_space-scanner_desc-optcom_bold.nii.gz"
mask_file = "sub-04570_task-rest_space-scanner_desc-brain_mask.nii.gz"
confounds_file = "sub-04570_task-rest_desc-confounds_timeseries.tsv"

# Files from tedana (after running on fMRIPrepped data)
mixing_file = "sub-04570_task-rest_space-scanner_desc-ICA_mixing.tsv"
metrics_file = "sub-04570_task-rest_space-scanner_desc-tedana_metrics.tsv"

# Load the mixing matrix
mixing_df = pd.read_table(opj("tedana_manual",mixing_file))  # Shape is time-by-components

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

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

# Combine the rejected components and the fMRIPrep confounds into a single array
regressors = np.hstack((rejected_components, confounds))

masker = NiftiMasker(
    mask_img=opj("func", 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(opj("tedana_manual", 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(
    opj("func", "sub-04570_task-rest_space-scanner_desc-aggrDenoised_bold.nii.gz")
)

import numpy as np  # A library for working with numerical data
from nilearn.masking import apply_mask, unmask  # Functions for (un)masking fMRI data

# Apply the mask to the data image to get a 2d array
data = apply_mask(opj("tedana_manual", data_file), opj("func", mask_file))

# Fit GLM to accepted components, rejected components and nuisance regressors
# (after adding a constant term)
regressors = np.hstack(
    (
        confounds,
        rejected_components,
        accepted_components,
        np.ones((mixing_df.shape[0], 1)),
    ),
)
betas = np.linalg.lstsq(regressors, data, rcond=None)[0][:-1]

# Denoise the data using the betas from just the bad components
confounds_idx = np.arange(confounds.shape[1] + rejected_components.shape[1])
pred_data = np.dot(np.hstack((confounds, rejected_components)), betas[confounds_idx, :])
data_denoised = data - pred_data

# Save to file
denoised_img = unmask(data_denoised, opj("func", mask_file))
denoised_img.to_filename(
    opj("func", "sub-04570_task-rest_space-scanner_desc-nonaggrDenoised_bold.nii.gz")
)

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

# Combine the confounds and rejected components in a single array
bad_timeseries = np.hstack((rejected_components, confounds))

# Regress the good components out of the bad time series to get "pure evil" regressors
betas = np.linalg.lstsq(accepted_components, bad_timeseries, rcond=None)[0]
pred_bad_timeseries = np.dot(accepted_components, betas)
orth_bad_timeseries = bad_timeseries - pred_bad_timeseries

# Once you have these "pure evil" components, you can denoise the data
masker = NiftiMasker(
    mask_img=opj("func", 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(opj("tedana_manual", data_file), confounds=orth_bad_timeseries)

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

# Save to file
denoised_img_4d.to_filename(
    opj("func", "sub-04570_task-rest_space-scanner_desc-orthAggrDenoised_bold.nii.gz")
)