"""PRISM configuration.

Defaults reproduce the released checkpoint (SigLIP2-so400m / Qwen3-Embedding-0.6B).
"""

from __future__ import annotations

from transformers import PretrainedConfig


class PRISMConfig(PretrainedConfig):
    model_type = "prism"

    def __init__(
        self,
        # ---- frozen backbones (loaded from the Hub by name) ----
        vision_backbone_name: str = "google/siglip2-so400m-patch14-384",
        text_backbone_name: str = "Qwen/Qwen3-Embedding-0.6B",
        # ---- architecture ----
        d_z: int = 512,
        qformer_depth: int = 4,
        temporal_depth: int = 12,
        predictor_depth: int = 4,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        max_frames: int = 128,
        logit_scale_init: float = 2.6592,
        # ---- training objective ----
        lambda_decomp: float = 1.0,
        lambda_temp: float = 0.5,
        infonce_all_gather: bool = True,
        sliding_shift_aug: bool = True,
        # ---- EMA target encoder ----
        use_ema: bool = True,
        ema_decay: float = 0.998,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vision_backbone_name = vision_backbone_name
        self.text_backbone_name = text_backbone_name
        self.d_z = d_z
        self.qformer_depth = qformer_depth
        self.temporal_depth = temporal_depth
        self.predictor_depth = predictor_depth
        self.num_heads = num_heads
        self.mlp_ratio = mlp_ratio
        self.max_frames = max_frames
        self.logit_scale_init = logit_scale_init
        self.lambda_decomp = lambda_decomp
        self.lambda_temp = lambda_temp
        self.infonce_all_gather = infonce_all_gather
        self.sliding_shift_aug = sliding_shift_aug
        self.use_ema = use_ema
        self.ema_decay = ema_decay