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


"""
Post-GLM data aggregation (no-error trials)

This script collects surface-level GLM contrast maps produced in the
first-level analysis and aggregates them into structured dictionaries
for later multivariate analyses (e.g., decoding, hyperalignment).

For each subject/run/sequence, only the first valid (no-error) pair is used.

Outputs (optional):
- Pickled dictionaries containing z-maps, vertex data arrays, and labels

Author: Theo Marins
"""

import os
import glob
import warnings
import pickle
import numpy as np
import nibabel as nib

from nilearn.surface import SurfaceImage

# ============================================================================
# Configuration
# ============================================================================

warnings.filterwarnings("ignore", category=UserWarning, module="nilearn")

OUT_DIR = "/PATH/TO/RESULTS"

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

# Subjects with insufficient valid repetitions (excluded)
INVALID_SUBJECTS = [
    "03", "06", "09", "10", "12", "13", "14",
    "15", "17", "18", "19", "21", "24", "25"
]

VALID_SUBJECTS = [s for s in subjects if s not in INVALID_SUBJECTS]

# ============================================================================
# Containers
# ============================================================================

subject_data = {}
sequence_label = {}
run_label = {}
subject_label = {}
session_label = {}
darrays_data = {}

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

for ses in sessions:
    try:
        subject_data[ses] = {}
        sequence_label[ses] = {}
        run_label[ses] = {}
        subject_label[ses] = {}
        session_label[ses] = {}
        darrays_data[ses] = {}

        for subject in VALID_SUBJECTS:
            print(f"Running sub-{subject}, ses-{ses}")

            subject_data[ses][subject] = {}
            sequence_label[ses][subject] = {}
            run_label[ses][subject] = {}
            subject_label[ses][subject] = {}
            session_label[ses][subject] = {}
            darrays_data[ses][subject] = {}

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

                z_maps = []
                seq_l, rl, subl, sesl = [], [], [], []
                left_darrays, right_darrays = [], []

                # ------------------------------------------------------------------
                # Loop over sequences and pairs (keep first valid only)
                # ------------------------------------------------------------------
                for seq in sequences:
                    collected = 0

                    for pair in range(1, 4):
                        if collected >= 1: #only processes the first valid pair
                            break

                        pattern = (
                            f"sub-{subject}_ses-{ses}_task-motorseq_run-{run}_"
                            f"desc-sequence_{seq}_pair{pair}"
                        )

                        files_data_L = glob.glob(
                            os.path.join(OUT_DIR, f"{pattern}_data_hemi-L.gii")
                        )
                        files_mesh_L = glob.glob(
                            os.path.join(OUT_DIR, f"{pattern}_mesh_hemi-L.gii")
                        )
                        files_data_R = glob.glob(
                            os.path.join(OUT_DIR, f"{pattern}_data_hemi-R.gii")
                        )
                        files_mesh_R = glob.glob(
                            os.path.join(OUT_DIR, f"{pattern}_mesh_hemi-R.gii")
                        )

                        for dL, mL, dR, mR in zip(
                            files_data_L, files_mesh_L, files_data_R, files_mesh_R
                        ):
                            z_maps.append(
                                SurfaceImage(
                                    data={"left": dL, "right": dR},
                                    mesh={"left": mL, "right": mR},
                                )
                            )

                            left_darrays.append(
                                nib.load(dL).darrays[0].data
                            )
                            right_darrays.append(
                                nib.load(dR).darrays[0].data
                            )

                            seq_l.append(seq)
                            rl.append(run)
                            subl.append(subject)
                            sesl.append(ses)

                            collected += 1

                        if collected >= 1:
                            break

                # ------------------------------------------------------------------
                # Concatenate vertex data
                # ------------------------------------------------------------------
                if left_darrays and right_darrays:
                    left_concatenated = np.vstack(left_darrays)
                    right_concatenated = np.vstack(right_darrays)

                    darrays_data[ses][subject][run] = {
                        "left": left_concatenated,
                        "right": right_concatenated,
                    }
                else:
                    darrays_data[ses][subject][run] = {
                        "left": None,
                        "right": None,
                    }

                # ------------------------------------------------------------------
                # Store labels and maps
                # ------------------------------------------------------------------
                subject_data[ses][subject][run] = z_maps
                sequence_label[ses][subject][run] = seq_l
                run_label[ses][subject][run] = rl
                subject_label[ses][subject][run] = subl
                session_label[ses][subject][run] = sesl

            print(f"Finished sub-{subject}, ses-{ses}")

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

# ============================================================================
# Save outputs as pickle files
# ============================================================================

PKL_DIR = BASE_DIR / "pkl"
os.makedirs(PKL_DIR, exist_ok=True)

def save_pickle(obj, name):
    path = os.path.join(PKL_DIR, name)
    with open(path, "wb") as f:
        pickle.dump(obj, f)
    print(f"Saved {path}")

save_pickle(subject_data, "subject_data.pkl")
save_pickle(sequence_label, "sequence_label.pkl")
save_pickle(run_label, "run_label.pkl")
save_pickle(subject_label, "subject_label.pkl")
save_pickle(session_label, "session_label.pkl")
save_pickle(darrays_data, "darrays.pkl")