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


# ### Within-subject classification


"""
Within-subject searchlight classification with manual permutation testing
(no-error trials, leave-one-run-out CV)

This script performs within-subject decoding of sequence identity using
surface-based searchlight SVM classification. A leave-one-run-out
cross-validation scheme is used, with permutation-based p-values
computed per searchlight.

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

warnings.filterwarnings("ignore", category=ConvergenceWarning)

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

PERM_DIR = os.path.join(BASE_DIR, "permutations_ws")
os.makedirs(PERM_DIR, exist_ok=True)

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 sessions:", subject_data.keys())

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

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

radius = 10
batch_size = 60
n_permutations = 100

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 (within-subject)
# ============================================================================

def permutation_test_searchlight_within(
    batch_indices,
    sls,
    train_data,
    test_data,
    train_labels,
    test_labels,
    decoder,
    n_permutations
):
    """
    Manual permutation test for a batch of searchlights (within-subject).
    """

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

    # concatenate training runs
    X_train_all = np.concatenate(train_data, axis=0)
    y_train_all = train_labels

    for idx in batch_indices:
        sl = sls[idx]

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

        clf = clone(decoder)
        clf.fit(X_train, y_train_all)
        score = clf.score(X_test, test_labels)

        perm_scores = []
        for _ in range(n_permutations):
            y_perm = rng.permutation(y_train_all)
            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
# ============================================================================

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

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

for hemi in ["right", "left"]:

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

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

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

    for session in subject_data.keys():
        for subject in subjects:

            print(f"\nSubject {subject} | Session {session}")

            # ------------------------------------------------------------
            # Load runs
            # ------------------------------------------------------------
            run_data = []
            run_labels = []

            for run in subject_data[session][subject].keys():
                run_data.append(subject_data[session][subject][run][hemi])
                run_labels.append(sequence_label[session][subject][run])

            n_runs = len(run_data)

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

            # ------------------------------------------------------------
            # Leave-one-run-out cross-validation
            # ------------------------------------------------------------
            for test_run in range(n_runs):

                print(f"  Test run {test_run + 1}/{n_runs}")

                train_data = []
                train_labels = []

                for run_id, (data, labels) in enumerate(zip(run_data, run_labels)):
                    if run_id == test_run:
                        test_data = data
                        test_labels = labels
                    else:
                        train_data.append(data)
                        train_labels.append(labels)

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

                # --------------------------------------------------------
                # Searchlight batches
                # --------------------------------------------------------
                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_within(
                        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")

                gc.collect()
                check_memory_usage()

            # ------------------------------------------------------------
            # Average across runs
            # ------------------------------------------------------------
            scores[hemi][subject] /= n_runs
            permutation_scores[hemi][subject] /= n_runs
            pvalues[hemi][subject] /= n_runs

            # ------------------------------------------------------------
            # Save
            # ------------------------------------------------------------
            os.makedirs(PERM_DIR, exist_ok=True)

            with open(os.path.join(PERM_DIR, f"scores--{hemi}--{subject}.pkl"), "wb") as f:
                pickle.dump(scores[hemi][subject], f)

            with open(os.path.join(PERM_DIR, f"perm_scores--{hemi}--{subject}.pkl"), "wb") as f:
                pickle.dump(permutation_scores[hemi][subject], f)

            with open(os.path.join(PERM_DIR, f"pvals--{hemi}--{subject}.pkl"), "wb") as f:
                pickle.dump(pvalues[hemi][subject], f)

            print(f"Saved results for subject {subject}")

print("\nAll within-subject processing completed successfully.")