#!/usr/bin/env python
# coding: utf-8


# # GLM

# In[ ]:


"""
First-level surface GLM analysis 

This script fits a surface-based GLM for each subject/run using nilearn,
restricting analysis to trials without behavioral errors. Sequence identity
is modeled at the pair level, and nuisance regressors are included as confounds.

Outputs:
- Z-maps for each sequence × pair contrast (fsaverage5 surface)

Important:
- 

Author: Theo Marins
"""

import os
import numpy as np
import pandas as pd

from nilearn.glm.first_level import FirstLevelModel
from nilearn.surface import PolyData, SurfaceImage
from nilearn.datasets import load_fsaverage, load_fsaverage_data
from nilearn.plotting import plot_surf_stat_map, show
from pathlib import Path

# ============================================================================
# Paths and parameters
# ============================================================================

# Base directories 
BASE_DIR = Path("/PATH/TO/PROJECT")
RAW_DIR = BASE_DIR / "raw"
DERIV_DIR = RAW_DIR / "derivatives"
OUT_DIR = BASE_DIR / "z_scores"

# Make output dir if needed
OUT_DIR.mkdir(parents=True, exist_ok=True)


TR = 1.0

subjects = [f"{i:02d}" for i in range(1, 27)]
sessions = ["01"]
runs = [f"{i:02d}" for i in range(1, 9)]

fsaverage = load_fsaverage()
fsaverage_data = load_fsaverage_data(data_type="sulcal")

# ============================================================================
# Sequence remapping (group-specific counterbalancing)
# ============================================================================

GROUP2_SUBJECTS = [
    "02", "04", "07", "09", "11", "13",
    "15", "17", "19", "21", "23", "25"
]

def remap_seq_num(seq_num, subject):
    """
    Swap sequence labels (1–6 ↔ 7–12) for group 2 subjects.
    """
    if subject in GROUP2_SUBJECTS:
        if 1 <= seq_num <= 6:
            return seq_num + 6
        elif 7 <= seq_num <= 12:
            return seq_num - 6
    return seq_num

# ============================================================================
# Main loop
# ============================================================================

for subject in subjects:
    try:
        for ses in sessions:
            print(f"\nStarting sub-{subject}, ses-{ses}")

            for run in runs:
                print(f"  Processing run-{run}")

                # ------------------------------------------------------------------
                # Load surface fMRI data
                # ------------------------------------------------------------------
                fmri_img_L = (
                    f"{DERIV_DIR}/sub-{subject}/ses-{ses}/func/"
                    f"sub-{subject}_ses-{ses}_task-motorseq_run-{run}_hemi-l_space-fsaverage5_bold.func.gii"
                )
                fmri_img_R = (
                    f"{DERIV_DIR}/sub-{subject}/ses-{ses}/func/"
                    f"sub-{subject}_ses-{ses}_task-motorseq_run-{run}_hemi-r_space-fsaverage5_bold.func.gii"
                )

                surface_image = SurfaceImage(
                    mesh=fsaverage["pial"],
                    data=PolyData(left=fmri_img_L, right=fmri_img_R)
                )

                # ------------------------------------------------------------------
                # Load and prepare events
                # ------------------------------------------------------------------
                events_file = (
                    f"{RAW_DIR}/sub-{subject}/ses-{ses}/func/"
                    f"sub-{subject}_ses-{ses}_task-motorseq_run-{run}_events.tsv"
                )
                events = pd.read_table(events_file)

                events["duration"] = 3.5  # fixed duration (Berlot, personal comm.)
                events["sequence_type"] = events["sequence_type"].replace(
                    {1: "trained", 2: "untrained"}
                )

                events["sequence_num"] = events["sequence_num"].astype(int)
                events["sequence_num"] = events["sequence_num"].apply(
                    lambda x: remap_seq_num(x, subject)
                )

                events["sequence"] = events["sequence_num"].astype(str)

                # Pair repetitions (2 trials per pair)
                events["pair"] = (events.groupby("sequence").cumcount() // 2) + 1
                events["condition"] = events["sequence"] + "_pair" + events["pair"].astype(str)

                # ------------------------------------------------------------------
                # Remove error trials (pair-wise exclusion)
                # ------------------------------------------------------------------
                valid_conditions = (
                    events.groupby("condition")["is_error"].sum()
                )
                valid_conditions = valid_conditions[valid_conditions == 0].index

                events_no_error = events[events["condition"].isin(valid_conditions)]

                events_glm = (
                    events_no_error[["onset", "duration", "condition"]]
                    .rename(columns={"condition": "trial_type"})
                    .replace(np.nan, 0)
                )

                # ------------------------------------------------------------------
                # Load confounds
                # ------------------------------------------------------------------
                confounds_file = (
                    f"{DERIV_DIR}/sub-{subject}/ses-{ses}/func/"
                    f"sub-{subject}_ses-{ses}_task-motorseq_run-{run}_desc-confounds_timeseries.tsv"
                )
                confounds = pd.read_csv(confounds_file, sep="\t")

                confounds_glm = confounds[
                    [
                        "white_matter", "global_signal", "framewise_displacement",
                        "trans_x", "trans_y", "trans_z",
                        "rot_x", "rot_y", "rot_z",
                    ]
                ].replace(np.nan, 0)

                # ------------------------------------------------------------------
                # Fit surface GLM
                # ------------------------------------------------------------------
                glm = FirstLevelModel(
                    t_r=TR,
                    hrf_model="spm",
                ).fit(
                    run_imgs=surface_image,
                    events=events_glm,
                    confounds=confounds_glm,
                )

                # ------------------------------------------------------------------
                # Define contrasts (task regressors only)
                # ------------------------------------------------------------------
                design_matrix = glm.design_matrices_[0]

                task_columns = [
                    col for col in design_matrix.columns
                    if not col.startswith(
                        ("white_matter", "global_signal", "framewise", "trans_", "rot_")
                    )
                    and col != "constant"
                ]

                # ------------------------------------------------------------------
                # Compute and visualize contrasts
                # ------------------------------------------------------------------
                for contrast_id in task_columns:
                    contrast_vector = np.zeros(len(design_matrix.columns))
                    contrast_vector[design_matrix.columns.get_loc(contrast_id)] = 1

                    z_map = glm.compute_contrast(
                        contrast_vector,
                        output_type="z_score"
                    )
                    out_file = (
                        OUT_DIR / f"sub-{subject}_ses-{ses}_run-{run}_contrast-{contrast_id}_zmap.func.gii"
                    )
                    z_map.to_filename(out_file)
                    print(f"    Saved z-map for contrast '{contrast_id}' to {out_file}")

    except Exception as e:
        print(f"Error computing contrasts for sub-{subject}: {e}")