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


"""
Group-level permutation inference for searchlight decoding accuracy


This script aggregates subject-level permutation results, computes
group-level statistics using a permutation-based null distribution,
and applies FDR correction separately for each hemisphere.

Author: Theo Marins
"""

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

import os
import pickle
import numpy as np
from statsmodels.stats.multitest import multipletests

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

PERM_DIR = os.path.join(BASE_DIR, "permutations_HA")  # or "permutations_AA" or "permutations_ws"
STATS_DIR = PERM_DIR / "stats_HA"  # or "stats_AA"
os.makedirs(STATS_DIR, exist_ok=True)

subjects = ["01", "02", "04", "05", "07", "08", "11", "16", "20", "22", "23", "26"]
hemispheres = ["l", "r"]

# ============================================================================
# Group-level aggregation
# ============================================================================

group_results = {}

for hemi in hemispheres:

    print(f"\nProcessing hemisphere: {hemi}")

    subject_scores = []
    subject_perms = []
    included_subjects = []

    # ------------------------------------------------------------------------
    # Load subject-level data
    # ------------------------------------------------------------------------
    for sub in subjects:

        scores_file = os.path.join(PERM_DIR, f"score--{hemi}--sub{sub}.pkl")
        perm_file = os.path.join(PERM_DIR, f"perm_scores--{hemi}--sub{sub}.pkl")

        if not (os.path.exists(scores_file) and os.path.exists(perm_file)):
            print(f"  Skipping sub-{sub} ({hemi}): missing files")
            continue

        with open(scores_file, "rb") as f:
            scores = np.asarray(pickle.load(f), dtype=float)

        with open(perm_file, "rb") as f:
            perms = np.asarray(pickle.load(f), dtype=float)

        if perms.ndim != 2:
            raise ValueError(
                f"Permutation scores must be 2D (V × P), got {perms.shape} for sub-{sub}"
            )

        if scores.shape[0] != perms.shape[0]:
            raise ValueError(
                f"Vertex mismatch for sub-{sub}: scores={scores.shape[0]}, perms={perms.shape[0]}"
            )

        subject_scores.append(scores)
        subject_perms.append(perms)
        included_subjects.append(sub)

    if len(included_subjects) < 2:
        raise RuntimeError(f"Not enough subjects for group inference ({hemi})")

    # ------------------------------------------------------------------------
    # Stack arrays
    # ------------------------------------------------------------------------
    X_obs = np.stack(subject_scores, axis=0)      # (S, V)
    X_perm = np.stack(subject_perms, axis=0)      # (S, V, P)

    S, V = X_obs.shape
    P = X_perm.shape[2]

    # ------------------------------------------------------------------------
    # Group statistic (mean across subjects)
    # ------------------------------------------------------------------------
    mean_obs = X_obs.mean(axis=0)                 # (V,)
    mean_null = X_perm.mean(axis=0)               # (V, P)

    # One-sided permutation p-values with +1 correction
    C = (mean_null >= mean_obs[:, None]).sum(axis=1)
    pvals = (C + 1) / (P + 1)

    # Standard error across subjects (for plotting)
    se_obs = X_obs.std(axis=0, ddof=1) / np.sqrt(S)

    # ------------------------------------------------------------------------
    # FDR correction (Benjamini–Hochberg)
    # ------------------------------------------------------------------------
    reject, pvals_fdr, _, _ = multipletests(pvals, alpha=0.05, method="fdr_bh")

    masked_accuracy = np.where(reject, mean_obs, 0.0)

    # ------------------------------------------------------------------------
    # Store results
    # ------------------------------------------------------------------------
    group_results[hemi] = {
        "subjects_used": included_subjects,
        "mean_accuracy": mean_obs,
        "se_accuracy": se_obs,
        "pvals_uncorrected": pvals,
        "pvals_fdr": pvals_fdr,
        "reject_fdr": reject,
        "masked_accuracy_fdr": masked_accuracy,
        "mean_null": mean_null,  # (V, P)
    }

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

with open(STATS_DIR, "wb") as f:
    pickle.dump(group_results, f)