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


# # Decoding model
# 

# ### Between-subject classification

# In[ ]:


"""
Permutation-based searchlight decoding on hyperaligned/anatomically aligned data

This script loads hyperaligned/anatomically aligned training and test data, aggregates runs,
and performs manual permutation testing using a Linear SVM decoder
within searchlights on the cortical surface.

Outputs:
- Decoding accuracy per vertex
- Permutation distributions
- P-values (nonparametric)

Author: Theo Marins
"""

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

import os
import gc
import time
import pickle
import warnings
import psutil

import numpy as np
import neuroboros as nb
from tqdm import tqdm
from sklearn.svm import LinearSVC
from sklearn.exceptions import ConvergenceWarning
from sklearn.base import clone

# Silence convergence warnings
warnings.filterwarnings("ignore", category=ConvergenceWarning)

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


PERM_DIR = os.path.join(BASE_DIR, "permutations_HA") # or "permutations_AA" for anatomical alignment
os.makedirs(OUT_DIR, exist_ok=True)

test_data_file = os.path.join(PERM_DIR, "test_data.pkl")
train_data_file = os.path.join(PERM_DIR, "train_data.pkl")
test_label_file = os.path.join(PERM_DIR, "test_labels.pkl")
train_label_file = os.path.join(PERM_DIR, "train_labels.pkl")

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

with open(test_data_file, "rb") as f:
    test_data_dict = pickle.load(f)

with open(train_data_file, "rb") as f:
    train_data_dict = pickle.load(f)

with open(test_label_file, "rb") as f:
    test_label_dict = pickle.load(f)

with open(train_label_file, "rb") as f:
    train_label_dict = pickle.load(f)

subjects = list(test_data_dict["l"].keys())
print("Loaded subjects:", subjects)

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

radius = 10
batch_size = 60
n_permutations = 100
n_jobs = -1

decoder = LinearSVC(
    penalty="l2",
    dual="auto",
    max_iter=5000
)

# Memory guard (important for cluster runs)
memory_limit_gb = 200
memory_limit_bytes = memory_limit_gb * 1024**3

# ============================================================================
# Utilities
# ============================================================================

def check_memory_usage():
    """Abort if memory usage exceeds threshold."""
    process = psutil.Process(os.getpid())
    memory_usage = process.memory_info().rss
    if memory_usage > memory_limit_bytes:
        raise MemoryError(
            f"Memory usage {memory_usage / 1024**3:.2f} GB "
            f"exceeds limit {memory_limit_gb} GB"
        )

# ============================================================================
# Permutation test (manual, searchlight-wise)
# ============================================================================

def permutation_test_searchlight(
    batch_indices,
    sls,
    train_data,
    test_data,
    train_labels,
    test_labels,
    decoder,
    n_permutations
):
    """
    Perform permutation testing for a batch of searchlights.
    """

    rng = np.random.default_rng(42)
    results = []

    # concatenate runs for training
    train_data = np.concatenate(train_data, axis=0)

    for idx in batch_indices:
        sl = sls[idx]

        X_train = train_data[:, sl]
        X_test = test_data[:, sl]

        # True score
        clf = clone(decoder)
        clf.fit(X_train, train_labels)
        score = clf.score(X_test, test_labels)

        # Permutations
        perm_scores = []
        for _ in range(n_permutations):
            y_perm = rng.permutation(train_labels)
            clf = clone(decoder)
            clf.fit(X_train, y_perm)
            perm_scores.append(clf.score(X_test, test_labels))

        perm_scores = np.array(perm_scores)
        pval = (np.sum(perm_scores >= score) + 1) / (n_permutations + 1)

        results.append((idx, score, perm_scores, pval))

    return results

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

scores = {}
permutation_scores = {}
pvalues = {}

for hemi in ["l", "r"]:

    print("\n" + "=" * 60)
    print(f"Processing hemisphere: {hemi}")
    print("=" * 60)

    sls = nb.sls(hemi, radius, space="fsavg-ico32", mask=True)
    valid_indices = list(range(len(sls)))

    scores[hemi] = {}
    permutation_scores[hemi] = {}
    pvalues[hemi] = {}

    for subject in tqdm(subjects, desc=f"{hemi} hemisphere"):

        print(f"Subject {subject}")

        # Output files (resume-safe)
        score_file = os.path.join(PERM_DIR, f"score--{hemi}--sub{subject}.pkl")
        perm_file = os.path.join(PERM_DIR, f"perm_scores--{hemi}--sub{subject}.pkl")
        pval_file = os.path.join(PERM_DIR, f"pvals_perm--{hemi}--sub{subject}.pkl")

        if all(os.path.exists(f) for f in [score_file, perm_file, pval_file]):
            print("  Skipping (already processed)")
            continue

        scores[hemi][subject] = np.zeros(len(sls))
        permutation_scores[hemi][subject] = np.zeros((len(sls), n_permutations))
        pvalues[hemi][subject] = np.zeros(len(sls))

        train_data = train_data_dict[hemi][subject]
        test_data = test_data_dict[hemi][subject]
        train_labels = train_label_dict[hemi][subject]
        test_labels = test_label_dict[hemi][subject]

        # ------------------------------------------------------------
        # Process in batches (memory-safe)
        # ------------------------------------------------------------
        for i in range(0, len(valid_indices), batch_size):

            batch_indices = valid_indices[i:i + batch_size]
            print(f"Batch {i}–{i + len(batch_indices)}")

            start = time.time()
            batch_results = permutation_test_searchlight(
                batch_indices,
                sls,
                train_data,
                test_data,
                train_labels,
                test_labels,
                decoder,
                n_permutations
            )

            for idx, score, perm, pval in batch_results:
                scores[hemi][subject][idx] = score 
                permutation_scores[hemi][subject][idx] = perm
                pvalues[hemi][subject][idx] = pval

            print(f"  Done in {time.time() - start:.2f}s")

        check_memory_usage()
        gc.collect()

        # ------------------------------------------------------------
        # Save subject-level outputs
        # ------------------------------------------------------------
        with open(score_file, "wb") as f:
            pickle.dump(scores[hemi][subject], f)

        with open(perm_file, "wb") as f:
            pickle.dump(permutation_scores[hemi][subject], f)

        with open(pval_file, "wb") as f:
            pickle.dump(pvalues[hemi][subject], f)

print("\nAll processing completed successfully.")