import sys
import argparse
import json
import time
from pathlib import Path
from typing import Dict, Any, Optional

import numpy as np
import torch
import torch.nn as nn
from omegaconf import DictConfig, OmegaConf
from torch.utils.data import DataLoader
from timm.utils import AverageMeter, random_seed


from data import (
    Algonauts2025Dataset,
    load_algonauts2025_friends_fmri,
    load_algonauts2025_movie10_fmri,
    load_sharded_features,
    episode_filter,
)
from flow_matching.src.stage1.medarc_architecture import MultiSubjectConvLinearEncoder
from matcha_architecture import CFM
from metric import pearsonr_score

# DEFAULT_DATA_DIR = ROOT.parent / "algonauts2025/datasets" # Adjust based on workspace
DEFAULT_DATA_DIR = Path("/raid/lttung05/fmri_encoder/data/fmri")
SUBJECTS = (1, 2, 3, 5)