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


# # Hyperalignment

# In[ ]:


"""
Whole-brain searchlight hyperalignment

This script performs searchlight-based hyperalignment using Neuroboros
and Procrustes alignment on surface data. For each subject, data from all
other subjects are used to build the common space, and the left-out subject is
aligned to this common space (leave-one-subject-out scheme).

Outputs:
- Aligned training data
- Aligned test data
- Corresponding labels (sequences)

Author: Theo Marins
"""

# ============================================================================
# Imports
# ============================================================================

import os
import pickle
import warnings
from time import time

import numpy as np
import neuroboros as nb
import nibabel as nib
from tqdm import tqdm
from sklearn.exceptions import ConvergenceWarning

from hyperalignment.searchlight import (
    searchlight_procrustes,
    searchlight_template
)

# Silence convergence warnings (expected in ridge/procrustes)
warnings.filterwarnings("ignore", category=ConvergenceWarning)

# ============================================================================
# Paths
# ============================================================================


subject_data_file = os.path.join(PKL_DIR, "darrays.pkl")
run_label_file = os.path.join(PKL_DIR, "run_label.pkl")
subject_label_file = os.path.join(PKL_DIR, "subject_label.pkl")
session_label_file = os.path.join(PKL_DIR, "session_label.pkl")
sequence_label_file = os.path.join(PKL_DIR, "sequence_label.pkl")

# ============================================================================
# Load data
# ============================================================================

with open(subject_data_file, "rb") as f:
    subject_data = pickle.load(f)

with open(run_label_file, "rb") as f:
    run_label = pickle.load(f)

with open(subject_label_file, "rb") as f:
    subject_label = pickle.load(f)

with open(session_label_file, "rb") as f:
    session_label = pickle.load(f)

with open(sequence_label_file, "rb") as f:
    sequence_label = pickle.load(f)

print("Loaded data for sessions:", subject_data.keys())

# ============================================================================
# Parameters
# ============================================================================

sessions = ["01"]
subjects = list(subject_data["01"].keys())
runs = list(subject_data["01"]["01"].keys())

radius = 20
space = "fsavg-ico32"

# ============================================================================
# Masks
# ============================================================================

mask_l = nb.mask("l", space=space, mask=True)
mask_r = nb.mask("r", space=space, mask=True)

# ============================================================================
# Output containers
# ============================================================================

train_data_dict = {}
test_data_dict = {}
train_label_dict = {}
test_label_dict = {}

# ============================================================================
# Hyperalignment helper
# ============================================================================

def hyperalign_train_data(train_data, sls, dists, radius=20):
    """
    Perform searchlight hyperalignment on training subjects.

    Parameters
    ----------
    train_data : array
        Shape (n_subjects, n_samples, n_vertices)
    sls, dists : searchlight definitions
    radius : int
        Searchlight radius

    Returns
    -------
    template : array
        Group template
    aligned_train : array
        Hyperaligned training data
    """

    print("  Building template...")
    template = searchlight_template(
        dss=train_data,
        sls=sls,
        dists=dists,
        radius=radius,
        tmpl_kind="pca",
        n_jobs=1
    )

    aligned_train = []

    for i, subj_data in enumerate(train_data):
        print(f"  Aligning training subject {i}")

        W = searchlight_procrustes(
            subj_data,
            template,
            sls,
            dists,
            radius=radius
        )

        aligned_train.append(subj_data @ W)

    aligned_train = np.stack(aligned_train, axis=0)
    return template, aligned_train

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

start_time = time()

for hemi_name, hemi_key in zip(["left", "right"], ["l", "r"]):

    print("\n" + "=" * 60)
    print(f"Processing {hemi_name.upper()} hemisphere")
    print("=" * 60)

    mask = mask_l if hemi_key == "l" else mask_r

    # Define searchlights
    sls, dists = nb.sls(
        hemi_key,
        radius,
        space=space,
        mask=True,
        return_dists=True
    )

    train_data_dict[hemi_key] = {}
    test_data_dict[hemi_key] = {}
    train_label_dict[hemi_key] = {}
    test_label_dict[hemi_key] = {}

    # Leave-one-subject-out
    for test_subject in subjects:

        print(f"\nTest subject: {test_subject}")

        train_data = []
        train_labels = []
        test_labels = []

        # ------------------------------------------------------------
        # Assemble data
        # ------------------------------------------------------------
        for ses in sessions:
            for subject in subjects:

                subj_runs = []
                subj_labels = []

                for run in runs:
                    X = subject_data[ses][subject][run][hemi_name][:, mask] 
                    y = sequence_label[ses][subject][run]

                    subj_runs.append(X) #
                    subj_labels.extend(y)

                subj_data = np.concatenate(subj_runs, axis=0)

                if subject == test_subject:
                    test_data = subj_data
                    test_labels = subj_labels
                else:
                    train_data.append(subj_data)
                    train_labels.extend(subj_labels)

        train_data = np.stack(train_data, axis=0)

        # ------------------------------------------------------------
        # Hyperalignment (train subjects)
        # ------------------------------------------------------------
        print("Running hyperalignment...")
        template, train_data_ha = hyperalign_train_data(
            train_data, sls, dists, radius
        )

        # ------------------------------------------------------------
        # Align test subject
        # ------------------------------------------------------------
        W_test = searchlight_procrustes(
            test_data,
            template,
            sls,
            dists,
            radius=radius
        )

        test_data_ha = test_data @ W_test

        # ------------------------------------------------------------
        # Store
        # ------------------------------------------------------------
        train_data_dict[hemi_key][test_subject] = train_data_ha
        test_data_dict[hemi_key][test_subject] = test_data_ha
        train_label_dict[hemi_key][test_subject] = train_labels
        test_label_dict[hemi_key][test_subject] = test_labels

print(f"\nCompleted in {time() - start_time:.2f} seconds.")

# ============================================================================
# Save results 
# ============================================================================

with open(os.path.join(PKL_DIR, "train_data.pkl"), "wb") as f:
    pickle.dump(train_data_dict, f)

with open(os.path.join(PKL_DIR, "test_data.pkl"), "wb") as f:
    pickle.dump(test_data_dict, f)

with open(os.path.join(PKL_DIR, "train_labels.pkl"), "wb") as f:
    pickle.dump(train_label_dict, f)

with open(os.path.join(PKL_DIR, "test_labels.pkl"), "wb") as f:
    pickle.dump(test_label_dict, f)

print(f"\nSaved hyperaligned data to: {PKL_DIR}")