"""Self-contained test suite for the V-JEPA 2.1 HuggingFace port.

Does not require the reference repository nor the published weights: it builds a
small randomly-initialised model and exercises every public code path.

    python -m pytest test_vjepa21.py -q

Run it against both transformers majors before publishing: the processor and
attention APIs differ between 4.x and 5.x and a regression on one is invisible
from the other.
"""

from __future__ import annotations

import importlib
import json
import os
import shutil
import sys
import tempfile
import types

import numpy as np
import pytest
import torch

PORT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _load_port():
    name = "_vjepa21_port_under_test"
    if name not in sys.modules:
        pkg = types.ModuleType(name)
        pkg.__path__ = [PORT_DIR]
        sys.modules[name] = pkg
    cfg = importlib.import_module(f"{name}.configuration_vjepa21")
    mdl = importlib.import_module(f"{name}.modeling_vjepa21")
    vpr = importlib.import_module(f"{name}.video_processing_vjepa21")
    return cfg, mdl, vpr


CFG_MOD, MDL_MOD, VPR_MOD = _load_port()
VJEPA21Config = CFG_MOD.VJEPA21Config
VJEPA21Model = MDL_MOD.VJEPA21Model
VJEPA21ForVideoClassification = MDL_MOD.VJEPA21ForVideoClassification
VJEPA21VideoProcessor = VPR_MOD.VJEPA21VideoProcessor
apply_masks = MDL_MOD.apply_masks


def tiny_config(**kw):
    base = dict(
        patch_size=16,
        crop_size=64,
        tubelet_size=2,
        hidden_size=96,
        num_attention_heads=6,
        num_hidden_layers=12,
        num_pooler_heads=6,     # 96 // 16 = 6 would also work; pin it for clarity
        pred_hidden_size=48,
        pred_num_attention_heads=6,
        pred_num_hidden_layers=12,
        pred_teacher_embed_dim=96,
        pred_return_all_tokens=True,
        n_output_distillation=1,
        img_temporal_dim_size=1,
        interpolate_rope=True,
        num_labels=5,
    )
    base.update(kw)
    return VJEPA21Config(**base)


@pytest.fixture(scope="module")
def model():
    torch.manual_seed(0)
    return VJEPA21Model(tiny_config()).eval()


VIDEO = torch.randn(2, 3, 4, 64, 64)  # (B, C, T, H, W), 2 tubelets x 4 x 4 patches = 32 tokens


# --- configuration -----------------------------------------------------------


def test_hierarchical_layer_properties():
    cfg = tiny_config()
    assert cfg.encoder_hierarchical_layers == [2, 5, 8, 11]
    assert cfg.encoder_distillation_layers == [11]
    assert cfg.predictor_hierarchical_layers == [11]
    assert cfg.pretrained_grid_size == 16


def test_predictor_layer_map_differs_from_the_encoder_at_depth_24():
    """The reference predictor table is [4, 11, 17, 23] at depth 24 while the
    encoder table is [5, 11, 17, 23]. Sharing one table is a latent bug."""
    cfg = VJEPA21Config(
        hidden_size=96, num_attention_heads=6, num_hidden_layers=24,
        pred_hidden_size=48, pred_num_attention_heads=6, pred_num_hidden_layers=24,
        n_output_distillation=4, num_pooler_heads=6,
    )
    assert cfg.encoder_hierarchical_layers == [5, 11, 17, 23]
    assert cfg.predictor_hierarchical_layers == [4, 11, 17, 23]


def test_config_rejects_inconsistent_values():
    with pytest.raises(ValueError):
        tiny_config(n_output_distillation=9)
    with pytest.raises(ValueError):
        tiny_config(hidden_size=100, num_attention_heads=6)
    with pytest.raises(ValueError):
        VJEPA21Config(num_hidden_layers=13)
    with pytest.raises(ValueError):  # predictor depth not in the reference table
        tiny_config(pred_num_hidden_layers=48)
    with pytest.raises(ValueError):  # teacher dim not divisible by the level count
        tiny_config(
            num_hidden_layers=40, pred_num_hidden_layers=24,
            n_output_distillation=3, pred_teacher_embed_dim=100,
        )


def test_num_pooler_heads_defaults_to_sixteen():
    """Every frozen-probe config in `configs/eval_2_1/` uses classifier.num_heads: 16,
    for all four model sizes. Inheriting `num_attention_heads` would give 22 on
    ViT-g and 26 on ViT-G."""
    cfg = VJEPA21Config(hidden_size=1408, num_attention_heads=22, num_hidden_layers=40,
                        n_output_distillation=4, pred_num_hidden_layers=24)
    assert cfg.num_pooler_heads == 16
    assert cfg.num_pooler_layers == 3  # + 1 cross-attention block = num_probe_blocks: 4
    assert VJEPA21Config(num_pooler_heads=8).num_pooler_heads == 8


# --- encoder -----------------------------------------------------------------


@torch.no_grad()
def test_forward_shapes(model):
    out = model(pixel_values_videos=VIDEO, skip_predictor=True)
    assert out.last_hidden_state.shape == (2, 32, 96)
    assert out.predictor_output is None


@torch.no_grad()
@pytest.mark.parametrize("layout", ["BCTHW", "BTCHW", "BTHWC"])
def test_input_layouts_are_equivalent(model, layout):
    x = {
        "BCTHW": VIDEO,
        "BTCHW": VIDEO.permute(0, 2, 1, 3, 4).contiguous(),
        "BTHWC": VIDEO.permute(0, 2, 3, 4, 1).contiguous(),
    }[layout]
    ref = model(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
    got = model(pixel_values_videos=x, skip_predictor=True).last_hidden_state
    assert torch.equal(ref, got)


def test_unknown_layout_raises(model):
    with pytest.raises(ValueError, match="channel axis"):
        model(pixel_values_videos=torch.randn(2, 5, 4, 64, 64), skip_predictor=True)


def test_ambiguous_layout_warns(recwarn, capsys):
    """A 3-frame clip in (B, T, C, H, W) also matches (B, C, T, H, W). Deciding
    silently is the failure mode; the channels-first reading plus a warning is
    the contract.

    `warning_once` is memoised per message, so the warning is captured by calling
    `normalize_video_layout` directly with a fresh cache rather than through the
    module-scoped model, whose first call may already have spent it.
    """
    messages = []
    original = MDL_MOD._warn_once
    try:
        MDL_MOD._warn_once = messages.append
        x = torch.randn(2, 3, 3, 64, 64)
        out = MDL_MOD.normalize_video_layout(x, in_chans=3)
    finally:
        MDL_MOD._warn_once = original

    assert out.shape == (2, 3, 3, 64, 64)          # channels-first reading wins
    assert torch.equal(out, x)                     # and it is not transposed
    assert any("ambiguous video layout" in m for m in messages), messages


@torch.no_grad()
def test_ambiguous_layout_still_runs(model):
    """The warning must not become an exception: a 3-frame clip is legal input."""
    out = model(pixel_values_videos=torch.randn(2, 3, 3, 64, 64), skip_predictor=True)
    assert out.last_hidden_state.shape[0] == 2


@torch.no_grad()
def test_output_hidden_states(model):
    out = model(pixel_values_videos=VIDEO, skip_predictor=True, output_hidden_states=True)
    # embeddings + one entry per layer
    assert len(out.hidden_states) == model.config.num_hidden_layers + 1
    assert all(h.shape == (2, 32, 96) for h in out.hidden_states)
    # the documented relation between the last hidden state and the final norm
    expected = model.encoder.norms_block[-1](out.hidden_states[-1])
    assert torch.allclose(expected, out.last_hidden_state, atol=1e-6)


@torch.no_grad()
def test_output_attentions(model):
    out = model(pixel_values_videos=VIDEO, skip_predictor=True, output_attentions=True)
    assert len(out.attentions) == model.config.num_hidden_layers
    attn = out.attentions[0]
    assert attn.shape == (2, 6, 32, 32)
    assert torch.allclose(attn.sum(-1), torch.ones_like(attn.sum(-1)), atol=1e-5)


@torch.no_grad()
def test_multilevel_and_hierarchical_outputs(model):
    hier = model.config.encoder_hierarchical_layers
    out = model(
        pixel_values_videos=VIDEO,
        skip_predictor=True,
        out_layers=hier,
        return_hierarchical=True,
    )
    assert len(out.multilevel_hidden_states) == 4
    # each level uses its own norm, per the reference implementation
    assert torch.allclose(out.multilevel_hidden_states[-1], out.last_hidden_state, atol=1e-6)
    # distillation levels concatenated on the channel axis
    n_distill = len(model.config.encoder_distillation_layers)
    assert out.hierarchical_hidden_state.shape == (2, 32, 96 * n_distill)


@torch.no_grad()
def test_multilevel_output_order_follows_the_network(model):
    """Levels come back in network order regardless of how they were requested,
    so a probe that zips them with `encoder_hierarchical_layers` is correct."""
    hier = model.config.encoder_hierarchical_layers
    forward = model(pixel_values_videos=VIDEO, skip_predictor=True,
                    out_layers=hier).multilevel_hidden_states
    shuffled = model(pixel_values_videos=VIDEO, skip_predictor=True,
                     out_layers=list(reversed(hier))).multilevel_hidden_states
    for a, b in zip(forward, shuffled):
        assert torch.equal(a, b)


def test_out_layers_validation(model):
    with pytest.raises(ValueError, match="hierarchical layers"):
        model(pixel_values_videos=VIDEO, skip_predictor=True, out_layers=[3])


@torch.no_grad()
def test_image_branch_uses_image_patch_embedding(model):
    image = torch.randn(2, 3, 1, 64, 64)
    out = model(pixel_values_videos=image, skip_predictor=True)
    assert out.last_hidden_state.shape == (2, 16, 96)  # 1 frame x 4 x 4 patches
    assert model._detect_mode(image) == "img"
    assert model._detect_mode(VIDEO) == "video"


@torch.no_grad()
def test_variable_resolution(model):
    out = model(pixel_values_videos=torch.randn(1, 3, 4, 96, 128), skip_predictor=True)
    assert out.last_hidden_state.shape == (1, 2 * 6 * 8, 96)


@torch.no_grad()
def test_sdpa_matches_eager():
    cfg = tiny_config(attn_implementation="eager")
    torch.manual_seed(0)
    eager = VJEPA21Model(cfg).eval()
    sdpa = VJEPA21Model(tiny_config(attn_implementation="sdpa")).eval()
    sdpa.load_state_dict(eager.state_dict())
    a = eager(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
    b = sdpa(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
    assert torch.allclose(a, b, atol=1e-5)


@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_low_precision_weights(dtype):
    """Regression: RoPE frequencies must be cast back to the input dtype. Building
    them in `x.dtype` promotes q/k to float32 under bf16/fp16 weights and the fused
    attention kernels reject the mismatch against v."""
    if dtype is torch.float16 and not torch.cuda.is_available():
        pytest.skip("fp16 matmul is not supported on CPU")
    model = VJEPA21Model(tiny_config()).eval().to(dtype)
    out = model(pixel_values_videos=VIDEO.to(dtype), skip_predictor=True)
    assert out.last_hidden_state.dtype is dtype
    assert torch.isfinite(out.last_hidden_state).all()


@torch.no_grad()
def test_unexpected_kwargs_are_reported_not_swallowed(model, caplog):
    """`out_layer=[11]` (no s) used to be a silent no-op."""
    import logging

    with caplog.at_level(logging.WARNING, logger=MDL_MOD.logger.name):
        model(pixel_values_videos=VIDEO, skip_predictor=True, out_layer=[11])
    # warning_once may have been consumed by another test in the same process,
    # so only assert that the call did not raise and produced no multilevel output
    out = model(pixel_values_videos=VIDEO, skip_predictor=True, out_layer=[11])
    assert out.multilevel_hidden_states is None


# --- masked (JEPA) encoder forward -------------------------------------------


@torch.no_grad()
def test_encoder_masks_reduce_the_sequence(model):
    """`masks` drops tokens before the layers, as `encoder(clips, masks_enc)` does
    in the reference; attention then only sees the context."""
    idx = torch.stack([torch.randperm(32)[:12].sort().values for _ in range(2)])
    out = model.encoder(VIDEO, masks=[idx])
    assert out.last_hidden_state.shape == (2, 12, 96)


@torch.no_grad()
def test_encoder_masks_carry_true_positions(model):
    """Masked tokens must keep their original RoPE ids. Feeding the same tokens
    with a different index set has to change the output, otherwise positions are
    being renumbered 0..K-1."""
    keep = torch.arange(0, 12).unsqueeze(0).expand(2, -1)
    shifted = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    a = model.encoder(VIDEO, masks=[keep]).last_hidden_state
    b = model.encoder(VIDEO, masks=[shifted]).last_hidden_state
    assert not torch.allclose(a, b, atol=1e-4)


def test_encoder_rejects_out_of_range_masks(model):
    with pytest.raises(ValueError, match="out of range"):
        model.encoder(VIDEO, masks=[torch.tensor([[0, 999]]).expand(2, -1)])


@torch.no_grad()
def test_masked_model_forward_runs_the_predictor(model):
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    out = model(pixel_values_videos=VIDEO, masks=[ctx], context_mask=[ctx], target_mask=[tgt])
    assert out.last_hidden_state.shape == (2, 20, 96)
    assert out.predictor_output.last_hidden_state.shape == (2, 12, 96)
    assert out.masked_hidden_state.shape == (2, 20, 96)


def test_masked_model_forward_requires_target_mask(model):
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    with pytest.raises(ValueError, match="target_mask"):
        model(pixel_values_videos=VIDEO, masks=[ctx])


# --- predictor ---------------------------------------------------------------


@torch.no_grad()
def test_predictor_default_masks(model):
    out = model(pixel_values_videos=VIDEO)
    pred = out.predictor_output
    assert pred.last_hidden_state.shape == (2, 32, 96)  # pred_teacher_embed_dim
    assert pred.context_hidden_state.shape == (2, 32, 96)
    assert out.masked_hidden_state.shape == (2, 32, 96)


@torch.no_grad()
def test_predictor_explicit_masks(model):
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    out = model(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
    assert out.predictor_output.last_hidden_state.shape == (2, 12, 96)
    assert out.predictor_output.context_hidden_state.shape == (2, 20, 96)


@torch.no_grad()
def test_predictor_accepts_bare_tensors(model):
    """The reference wraps a bare tensor in a list; so do we."""
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    a = model(pixel_values_videos=VIDEO, context_mask=ctx, target_mask=tgt)
    b = model(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
    assert torch.equal(a.predictor_output.last_hidden_state,
                       b.predictor_output.last_hidden_state)


@torch.no_grad()
def test_mask_index_selects_the_intended_token():
    """With zero-initialised mask tokens every index behaves identically, which
    makes any test that relies on them vacuous. Give them distinct values first."""
    torch.manual_seed(0)
    m = VJEPA21Model(tiny_config()).eval()
    for i, tok in enumerate(m.predictor.embeddings.mask_tokens):
        torch.nn.init.constant_(tok, 0.1 * (i + 1))

    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    a = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=1)
    b = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=2)
    assert not torch.allclose(a.predictor_output.last_hidden_state,
                              b.predictor_output.last_hidden_state, atol=1e-5)
    # the index wraps modulo the number of mask tokens, as in the reference
    n = m.config.pred_num_mask_tokens
    c = m(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt], mask_index=1 + n)
    assert torch.allclose(a.predictor_output.last_hidden_state,
                          c.predictor_output.last_hidden_state, atol=1e-6)


@torch.no_grad()
def test_predictor_context_is_masked_flag(model):
    """Passing an already-gathered encoder output must give the same answer as
    passing the full sequence with the same context indices."""
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    z = model(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
    full = model.predictor(z, [ctx], [tgt])
    pre = model.predictor(apply_masks(z, [ctx]), [ctx], [tgt], context_is_masked=True)
    assert torch.allclose(full.last_hidden_state, pre.last_hidden_state, atol=1e-6)


def test_predictor_rejects_multiple_mask_pairs(model):
    m = torch.arange(0, 16).unsqueeze(0).expand(2, -1)
    with pytest.raises(NotImplementedError):
        model(pixel_values_videos=VIDEO, context_mask=[m, m], target_mask=[m, m])


def test_predictor_rejects_inconsistent_masked_context(model):
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    z = model(pixel_values_videos=VIDEO, skip_predictor=True).last_hidden_state
    with pytest.raises(ValueError, match="context_is_masked"):
        model.predictor(z, [ctx], [tgt], context_is_masked=True)


# --- classification head -----------------------------------------------------


def test_classification_loss_and_backward():
    torch.manual_seed(0)
    clf = VJEPA21ForVideoClassification(tiny_config())
    out = clf(pixel_values_videos=VIDEO, labels=torch.tensor([1, 3]))
    assert out.logits.shape == (2, 5)
    assert torch.isfinite(out.loss)
    out.loss.backward()
    assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in clf.parameters())


def test_multi_label_classification_loss():
    """XD-Violence is multi-label; float targets must route to BCE."""
    torch.manual_seed(0)
    cfg = tiny_config()
    cfg.problem_type = "multi_label_classification"
    clf = VJEPA21ForVideoClassification(cfg)
    labels = torch.zeros(2, 5)
    labels[0, 1] = labels[0, 4] = labels[1, 2] = 1.0
    out = clf(pixel_values_videos=VIDEO, labels=labels)
    assert torch.isfinite(out.loss)
    out.loss.backward()


def test_trainer_style_kwargs_do_not_reach_the_encoder():
    """The Trainer injects `num_items_in_batch`; it must not blow up or be
    forwarded down to the encoder."""
    clf = VJEPA21ForVideoClassification(tiny_config()).eval()
    with torch.no_grad():
        out = clf(pixel_values_videos=VIDEO, num_items_in_batch=2)
    assert out.logits.shape == (2, 5)


def test_gradient_checkpointing_matches_plain_forward():
    torch.manual_seed(0)
    a = VJEPA21ForVideoClassification(tiny_config())
    b = VJEPA21ForVideoClassification(tiny_config())
    b.load_state_dict(a.state_dict())
    a.train()
    b.train()
    b.gradient_checkpointing_enable()
    assert b.vjepa21.encoder.gradient_checkpointing
    assert torch.allclose(
        a(pixel_values_videos=VIDEO).logits, b(pixel_values_videos=VIDEO).logits, atol=1e-5
    )


def test_gradient_checkpointing_matches_on_the_predictor_path():
    """The classification head never runs the predictor, so its checkpointing
    path is otherwise untested."""
    torch.manual_seed(0)
    a = VJEPA21Model(tiny_config())
    b = VJEPA21Model(tiny_config())
    b.load_state_dict(a.state_dict())
    a.train()
    b.train()
    b.gradient_checkpointing_enable()
    assert b.predictor.gradient_checkpointing
    ctx = torch.arange(0, 20).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(20, 32).unsqueeze(0).expand(2, -1)
    ta = a(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
    tb = b(pixel_values_videos=VIDEO, context_mask=[ctx], target_mask=[tgt])
    assert torch.allclose(
        ta.predictor_output.last_hidden_state,
        tb.predictor_output.last_hidden_state,
        atol=1e-5,
    )


@torch.no_grad()
def test_classification_propagates_hidden_states():
    clf = VJEPA21ForVideoClassification(tiny_config()).eval()
    out = clf(pixel_values_videos=VIDEO, output_hidden_states=True, output_attentions=True)
    assert len(out.hidden_states) == 13
    assert len(out.attentions) == 12


# --- video processor ---------------------------------------------------------


def test_video_processor_matches_reference_transform():
    vp = VJEPA21VideoProcessor()
    assert vp.size == {"shortest_edge": 384}
    assert vp.crop_size == {"height": 384, "width": 384}
    assert list(vp.image_mean) == [0.485, 0.456, 0.406]
    assert list(vp.image_std) == [0.229, 0.224, 0.225]
    assert vp.rescale_factor == pytest.approx(1 / 255)

    clip = [np.full((240, 320, 3), 255, dtype=np.uint8) for _ in range(4)]
    pv = vp([clip], return_tensors="pt")["pixel_values_videos"]
    assert pv.shape == (1, 4, 3, 384, 384)  # (B, T, C, H, W)
    expected = (1.0 - np.array(vp.image_mean)) / np.array(vp.image_std)
    for c in range(3):
        assert pv[0, 0, c].mean().item() == pytest.approx(expected[c], abs=1e-4)


def test_video_processor_does_not_subclass_videos_kwargs():
    """Regression: an empty `VideosKwargs` subclass loses the TypedDict defaults on
    transformers 5, and `preprocess` then raises StrictDataclassFieldValidationError."""
    from transformers.processing_utils import VideosKwargs

    assert VJEPA21VideoProcessor.valid_kwargs is VideosKwargs


def test_video_processor_crop_size_is_configurable():
    vp = VJEPA21VideoProcessor(crop_size=256)
    assert vp.crop_size == {"height": 256, "width": 256}
    assert vp.size == {"shortest_edge": 256}


def test_processor_output_feeds_the_model_without_a_transpose():
    """The processor emits (B, T, C, H, W); the model must accept it as-is."""
    vp = VJEPA21VideoProcessor(crop_size=64)
    clip = [np.random.randint(0, 255, (80, 100, 3), dtype=np.uint8) for _ in range(4)]
    pv = vp([clip], return_tensors="pt")["pixel_values_videos"]
    assert pv.shape == (1, 4, 3, 64, 64)
    m = VJEPA21Model(tiny_config()).eval()
    with torch.no_grad():
        assert m(pixel_values_videos=pv, skip_predictor=True).last_hidden_state.shape == (
            1, 2 * 4 * 4, 96
        )


# --- serialisation round-trip ------------------------------------------------


def test_auto_classes_round_trip():
    from transformers import (
        AutoConfig,
        AutoModel,
        AutoModelForVideoClassification,
        AutoVideoProcessor,
    )

    tmp = tempfile.mkdtemp()
    try:
        for f in (
            "configuration_vjepa21.py",
            "modeling_vjepa21.py",
            "video_processing_vjepa21.py",
        ):
            shutil.copy(os.path.join(PORT_DIR, f), tmp)

        VJEPA21Config.register_for_auto_class()
        VJEPA21Model.register_for_auto_class("AutoModel")
        VJEPA21ForVideoClassification.register_for_auto_class("AutoModelForVideoClassification")
        VJEPA21VideoProcessor.register_for_auto_class("AutoVideoProcessor")

        VJEPA21Model(tiny_config()).save_pretrained(tmp)
        VJEPA21VideoProcessor(crop_size=64).save_pretrained(tmp)

        cfg_path = os.path.join(tmp, "config.json")
        cfg_json = json.load(open(cfg_path))
        cfg_json["auto_map"]["AutoModelForVideoClassification"] = (
            "modeling_vjepa21.VJEPA21ForVideoClassification"
        )
        cfg_json["auto_map"]["AutoVideoProcessor"] = (
            "video_processing_vjepa21.VJEPA21VideoProcessor"
        )
        json.dump(cfg_json, open(cfg_path, "w"), indent=2)

        assert type(AutoConfig.from_pretrained(tmp, trust_remote_code=True)).__name__ == (
            "VJEPA21Config"
        )
        assert type(AutoModel.from_pretrained(tmp, trust_remote_code=True)).__name__ == (
            "VJEPA21Model"
        )
        clf = AutoModelForVideoClassification.from_pretrained(
            tmp, trust_remote_code=True, num_labels=7
        )
        assert type(clf).__name__ == "VJEPA21ForVideoClassification"
        assert clf.config.num_labels == 7

        vp = AutoVideoProcessor.from_pretrained(tmp, trust_remote_code=True)
        clip = [np.random.randint(0, 255, (100, 140, 3), dtype=np.uint8) for _ in range(4)]
        pv = vp([clip], return_tensors="pt")["pixel_values_videos"]
        with torch.no_grad():
            assert clf(pixel_values_videos=pv).logits.shape == (1, 7)
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def test_backbone_weights_survive_the_classification_wrapper():
    """`base_model_prefix` must let the encoder weights load into the head model."""
    tmp = tempfile.mkdtemp()
    try:
        base = VJEPA21Model(tiny_config())
        base.save_pretrained(tmp)
        clf = VJEPA21ForVideoClassification.from_pretrained(tmp)
        for k, v in base.encoder.state_dict().items():
            assert torch.equal(v, clf.vjepa21.encoder.state_dict()[k]), k
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def test_saved_config_round_trips_the_pooler_heads():
    """A regression here would silently change the probe architecture."""
    tmp = tempfile.mkdtemp()
    try:
        cfg = tiny_config(num_pooler_heads=6)
        VJEPA21Model(cfg).save_pretrained(tmp)
        reloaded = VJEPA21Config.from_pretrained(tmp)
        assert reloaded.num_pooler_heads == 6
        assert reloaded.num_pooler_layers == cfg.num_pooler_layers
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


# --- multi-level distillation variants (ViT-g / ViT-G style) ------------------


def giant_like_config(**kw):
    base = dict(
        patch_size=16, crop_size=64, hidden_size=96, num_attention_heads=6,
        num_hidden_layers=40, mlp_ratio=48 / 11, n_output_distillation=4,
        num_pooler_heads=6,
        pred_hidden_size=48, pred_num_attention_heads=6, pred_num_hidden_layers=24,
        pred_num_mask_tokens=8, pred_teacher_embed_dim=None, pred_return_all_tokens=True,
    )
    base.update(kw)
    return VJEPA21Config(**base)


@torch.no_grad()
def test_predictor_consumes_hierarchical_features_when_required():
    """n_output_distillation > 1: the predictor input is the concatenated levels."""
    cfg = giant_like_config()
    assert cfg.encoder_distillation_layers == [9, 19, 29, 39]
    assert len(cfg.predictor_hierarchical_layers) == 4

    model = VJEPA21Model(cfg).eval()
    x = torch.randn(1, 3, 4, 64, 64)
    out = model(pixel_values_videos=x)  # would raise a shape error before the fix
    assert out.predictor_output.last_hidden_state.shape == (1, 32, 96 * 4)
    assert out.predictor_output.context_hidden_state.shape == (1, 32, 96 * 4)
    # masked_hidden_state carries the hierarchical width in this regime
    assert out.masked_hidden_state.shape == (1, 32, 96 * 4)


@torch.no_grad()
def test_hierarchical_output_is_not_leaked():
    """The hierarchical tensor is only exposed when explicitly requested."""
    model = VJEPA21Model(giant_like_config()).eval()
    x = torch.randn(1, 3, 4, 64, 64)
    assert model(pixel_values_videos=x).hierarchical_hidden_state is None
    out = model(pixel_values_videos=x, skip_predictor=True, return_hierarchical=True)
    assert out.hierarchical_hidden_state.shape == (1, 32, 96 * 4)