import torch.nn as nn
from transformers import Dinov2Config, Dinov2Model, DINOv3ViTConfig, DINOv3ViTModel, VJEPA2Config, VJEPA2Model


def _hf_config(config_class, config):
    fields = config.to_dict()
    fields.pop("model_type", None)
    return config_class(**fields)


def build_vision_encoder(config):
    if config.model_type == "dinov2":
        model = Dinov2Model(_hf_config(Dinov2Config, config))
        nn.init.trunc_normal_(model.embeddings.mask_token, std=0.02)

    elif config.model_type == "dinov3_vit":
        model = DINOv3ViTModel(_hf_config(DINOv3ViTConfig, config))
        nn.init.trunc_normal_(model.embeddings.mask_token, std=0.02)

    elif config.model_type == "vjepa2":
        model = VJEPA2Model(_hf_config(VJEPA2Config, config))

    elif config.model_type == "vjepa2.1":
        from .vjepa21_encoder import VJEPA21Encoder

        model = VJEPA21Encoder(config)

    else:
        raise NotImplementedError()

    return model