import numpy as np
from scipy.stats import pearsonr
from laion_fmri.subject import load_subject

sub = load_subject("sub-03")

# Load betas for PPA and FFA voxels (check sub.get_available_rois() for exact names)
betas_ppa = sub.get_betas(session="ses-01", roi="PPA")  # (n_images, n_ppa)
betas_ffa = sub.get_betas(session="ses-01", roi="FFA")  # (n_images, n_ffa)
betas = np.concatenate([betas_ppa, betas_ffa], axis=1)  # (n_images, n_voxels)

# model_predictions: your encoding model's predicted responses
# shape (n_images, n_voxels), same image ordering as betas

def mean_pearson_r(pred, resp):
    return np.mean([pearsonr(pred[:, v], resp[:, v])[0]
                    for v in range(resp.shape[1])])

# Step 1: Observed test statistic
t_obs = mean_pearson_r(model_predictions, betas)

# Steps 2-3: Permute image indices in model predictions -> null distribution
#   This breaks the correspondence between which image the model is
#   predicting and which brain response it is compared against.
rng = np.random.default_rng(seed=0)
t_perm = np.array([
    mean_pearson_r(model_predictions[rng.permutation(len(betas))], betas)
    for _ in range(1000)
])

# Step 4: One-sided p-value
p = (1 + np.sum(t_perm >= t_obs)) / (1 + len(t_perm))
print(f"Observed mean r = {t_obs:.4f},  p = {p:.4f}")

import numpy as np
from scipy.stats import pearsonr
from laion_fmri.subject import load_subject

sub = load_subject("sub-03")

betas_ppa = sub.get_betas(session="ses-01", roi="PPA")  # (n_images, n_ppa)
betas_ffa = sub.get_betas(session="ses-01", roi="FFA")  # (n_images, n_ffa)
n_ppa, n_ffa = betas_ppa.shape[1], betas_ffa.shape[1]

betas_all = np.concatenate([betas_ppa, betas_ffa], axis=1)
# model_preds_all: predictions for all PPA+FFA voxels, shape (n_images, n_ppa+n_ffa)

# Compute per-voxel Pearson r once (image labels and betas stay fixed throughout)
r_all = np.array([pearsonr(model_preds_all[:, v], betas_all[:, v])[0]
                  for v in range(betas_all.shape[1])])

# ROI membership: 0 = PPA, 1 = FFA
roi_labels = np.array([0] * n_ppa + [1] * n_ffa)

# Step 1: Observed test statistic (mean r in PPA minus mean r in FFA)
t_obs = r_all[roi_labels == 0].mean() - r_all[roi_labels == 1].mean()

# Steps 2-3: Permute ROI labels across voxels -> null distribution
#   Image labels and beta values are NOT touched. Only the assignment of
#   voxels to PPA vs FFA is shuffled, breaking any regional difference.
rng = np.random.default_rng(seed=0)
t_perm = []
for _ in range(1000):
    shuffled = rng.permutation(roi_labels)
    t_perm.append(r_all[shuffled == 0].mean() - r_all[shuffled == 1].mean())
t_perm = np.array(t_perm)

# Step 4: Two-sided p-value
p = (1 + np.sum(np.abs(t_perm) >= np.abs(t_obs))) / (1 + len(t_perm))
print(f"Observed r-diff (PPA - FFA) = {t_obs:.4f},  p = {p:.4f}")