# models/backbones/video_mae.py

from transformers import VideoMAEForVideoClassification, VideoMAEImageProcessor
import os

def build_video_mae_backbone(config, device, ckpt_path=None):
    """
    Build HuggingFace VideoMAE model for video classification.
    This includes both backbone and classification head.
    """
    num_classes = config.MODEL.num_classes
    pretrained_model_name = ckpt_path if ckpt_path else config.MODEL.pretrained_model
    processor = VideoMAEImageProcessor.from_pretrained(config.MODEL.pretrained_model)
    model = VideoMAEForVideoClassification.from_pretrained(
        pretrained_model_name,
        num_labels=num_classes,
        ignore_mismatched_sizes=True,
        trust_remote_code=True,
        device_map=device
    )

    # Freeze backbone if specified in config
    if config.MODEL.freeze_backbone:
        for param in model.videomae.parameters():
            param.requires_grad = False

    return model, processor


def load_video_mae_checkpoint(config, device, ckpt_path):
    """
    Load fine-tuned VideoMAE checkpoint from a HuggingFace-style directory.

    Supports:
      - model.safetensors
      - pytorch_model.bin
      - config.json
    """
    return build_video_mae_backbone(config, device, ckpt_path)
