from pathlib import Path
import numpy as np
import pandas as pd
from nilearn.datasets import func
from nilearn.glm.first_level import FirstLevelModel
from nilearn.glm.contrasts import compute_fixed_effects


output_dir = Path.cwd() / "results" / "plot_two_runs_model"
output_dir.mkdir(exist_ok=True, parents=True)
print(f"Output will be saved to: {output_dir}")

data = func.fetch_fiac_first_level()
fmri_imgs = [data["func1"], data["func2"]]

from nilearn.image import mean_img

mean_img_ = mean_img(fmri_imgs[0])

design_files = [data["design_matrix1"], data["design_matrix2"]]
design_matrices = [pd.DataFrame(np.load(df)["X"]) for df in design_files]

fmri_glm = FirstLevelModel(
    mask_img=data["mask"],
    smoothing_fwhm=5,
    minimize_memory=True,
)

contrast_val = [[-1, -1, 1, 1]]
fmri_glm_run_1 = fmri_glm.fit(fmri_imgs[0], design_matrices=design_matrices[0])
summary_statistics_run_1 = fmri_glm_run_1.compute_contrast(
    contrast_val,
    output_type="all",
)

fmri_glm_run_2 = fmri_glm.fit(fmri_imgs[1], design_matrices=design_matrices[1])
contrast_val = np.array([[-1, -1, 1, 1]])
summary_statistics_run_2 = fmri_glm_run_2.compute_contrast(
    contrast_val,
    output_type="all",
)

# Run fixed effects analysis

contrast_imgs = [
    summary_statistics_run_1["effect_size"],
    summary_statistics_run_2["effect_size"],
]
variance_imgs = [
    summary_statistics_run_1["effect_variance"],
    summary_statistics_run_2["effect_variance"],
]

fixed_fx_contrast, fixed_fx_variance, fixed_fx_stat = compute_fixed_effects(
    contrast_imgs,
    variance_imgs,
    data["mask"],
    precision_weighted=True
)

print('This should be a 3D image and it is not')
print(fixed_fx_contrast.shape)