"""Transformers config for the V-JEPA2 fMRI encoder."""

from __future__ import annotations

from transformers import PretrainedConfig


class VJEPA2FMRIEncoderConfig(PretrainedConfig):
    model_type = "vjepa2_fmri_encoder"

    def __init__(
        self,
        checkpoint_filename: str = "vjepa2_offline_encoder.pth",
        output_dim: int = 20484,
        input_duration_seconds: float = 3.0,
        input_format: str = "video_tensor_b_t_c_h_w",
        output_description: str = "z_scored_fmri_betas_no_time_dimension",
        backbone_filename: str = "vitl.pt",
        vjepa_size: str = "large",
        load_vjepa: bool = True,
        image_size: int = 224,
        normalize_input: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.checkpoint_filename = checkpoint_filename
        self.output_dim = int(output_dim)
        self.input_duration_seconds = float(input_duration_seconds)
        self.input_format = input_format
        self.output_description = output_description
        self.backbone_filename = backbone_filename
        self.vjepa_size = vjepa_size
        self.load_vjepa = bool(load_vjepa)
        self.image_size = int(image_size)
        self.normalize_input = bool(normalize_input)