"""Implementation parity against the reference `facebookresearch/vjepa2`.

Reference modules are built directly from `app/vjepa_2_1/models/`, their weights
are copied into the port, and the outputs are compared. Reduced hidden sizes keep
the suite fast: these checks establish that the *code paths* are equivalent,
which is width-independent. Parity on the published weights at full resolution is
a separate concern, covered by `verify_vjepa21_port.py`.

    VJEPA2_REPO=/path/to/vjepa2 python -m pytest test_parity_official.py -s -q

Note on tolerances. The encoder is compared eager-vs-eager and should be
bit-exact. The predictor is not: `VisionTransformerPredictor.__init__` has no
`use_sdpa` parameter, so it is swallowed by `**kwargs` and the reference
predictor blocks always run SDPA, while the port runs eager here. That plus the
gather-vs-stack reordering puts the residual around 1e-06, three orders of
magnitude inside the 1e-3 tolerance Meta uses for its own ports.
"""

from __future__ import annotations

import importlib
import os
import sys
import types
from functools import partial

import pytest
import torch
import torch.nn as nn

PORT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
VJEPA2_REPO = os.environ.get("VJEPA2_REPO", "")

pytestmark = pytest.mark.skipif(
    not VJEPA2_REPO or not os.path.isdir(os.path.join(VJEPA2_REPO, "app", "vjepa_2_1")),
    reason="set VJEPA2_REPO to a clone of facebookresearch/vjepa2",
)


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")
    return cfg.VJEPA21Config, mdl.VJEPA21Model, mdl.apply_masks


VJEPA21Config, VJEPA21Model, apply_masks = _load_port()


def _reference_classes():
    repo = os.path.abspath(VJEPA2_REPO)
    if repo not in sys.path:
        sys.path.insert(0, repo)
    from app.vjepa_2_1.models.predictor import VisionTransformerPredictor
    from app.vjepa_2_1.models.vision_transformer import VisionTransformer

    return VisionTransformer, VisionTransformerPredictor


# Small but structurally faithful: 64x64 x 4 frames, patch 16, tubelet 2
# -> 2 x 4 x 4 = 32 tokens. head_dim = 96 // 6 = 16, so the RoPE split is
# d=h=w=4 with 4 unrotated dimensions, exercising the same branch as the real
# checkpoints (head_dim 64, split 20/20/20 + 4).
DIM, HEADS, DEPTH = 96, 6, 12
PRED_DIM, PRED_HEADS, PRED_DEPTH = 48, 6, 12
IMG, PATCH, FRAMES, TUBELET = 64, 16, 4, 2
NORM = partial(nn.LayerNorm, eps=1e-6)


# ---------------------------------------------------------------------------
# weight mapping
# ---------------------------------------------------------------------------


def _map_block(prefix, idx, sub, tensor, hidden, out):
    if sub.startswith("attn.qkv."):
        kind = sub.rsplit(".", 1)[-1]
        q, k, v = tensor.split(hidden, dim=0)
        out[f"{prefix}.layer.{idx}.attention.query.{kind}"] = q
        out[f"{prefix}.layer.{idx}.attention.key.{kind}"] = k
        out[f"{prefix}.layer.{idx}.attention.value.{kind}"] = v
    elif sub.startswith("attn.proj."):
        out[f"{prefix}.layer.{idx}.attention.proj." + sub.rsplit(".", 1)[-1]] = tensor
    else:
        out[f"{prefix}.layer.{idx}.{sub}"] = tensor


def reference_to_port(enc_sd, pred_sd, hidden, pred_hidden):
    """Rename reference tensors to port names. Splits the fused QKV projection."""
    out = {}
    for k, v in enc_sd.items():
        if k in ("img_mod_embed", "video_mod_embed"):
            out[f"encoder.embeddings.{k}"] = v
        elif k.startswith("patch_embed_img."):
            out["encoder.embeddings.patch_embeddings_img." + k[len("patch_embed_img."):]] = v
        elif k.startswith("patch_embed."):
            out["encoder.embeddings.patch_embeddings." + k[len("patch_embed."):]] = v
        elif k.startswith("norms_block."):
            out["encoder." + k] = v
        elif k.startswith("blocks."):
            idx, sub = k[len("blocks."):].split(".", 1)
            _map_block("encoder", idx, sub, v, hidden, out)
        elif k == "pos_embed":
            continue  # unused under RoPE
        else:
            raise AssertionError(f"unmapped reference encoder key: {k}")

    for k, v in pred_sd.items():
        if k in ("img_mod_embed", "video_mod_embed"):
            out[f"predictor.embeddings.{k}"] = v
        elif k.startswith("predictor_embed."):
            out["predictor.embeddings.predictor_embed." + k[len("predictor_embed."):]] = v
        elif k.startswith("mask_tokens."):
            out["predictor.embeddings." + k] = v
        elif k.startswith("predictor_norm."):
            out["predictor.layernorm." + k[len("predictor_norm."):]] = v
        elif k.startswith("predictor_proj_context."):
            out["predictor.proj_context." + k[len("predictor_proj_context."):]] = v
        elif k.startswith("predictor_proj."):
            out["predictor.proj." + k[len("predictor_proj."):]] = v
        elif k.startswith("predictor_blocks."):
            idx, sub = k[len("predictor_blocks."):].split(".", 1)
            _map_block("predictor", idx, sub, v, pred_hidden, out)
        elif k == "predictor_pos_embed":
            continue
        else:
            raise AssertionError(f"unmapped reference predictor key: {k}")
    return out


# ---------------------------------------------------------------------------
# builders
# ---------------------------------------------------------------------------


def build_reference(n_distill=1, out_layers=None, teacher_embed_dim=None):
    VisionTransformer, VisionTransformerPredictor = _reference_classes()
    torch.manual_seed(0)
    encoder = VisionTransformer(
        img_size=(IMG, IMG), patch_size=PATCH, num_frames=FRAMES, tubelet_size=TUBELET,
        in_chans=3, embed_dim=DIM, depth=DEPTH, num_heads=HEADS, mlp_ratio=4.0,
        qkv_bias=True, norm_layer=NORM, use_rope=True, use_sdpa=False,
        uniform_power=False, handle_nonsquare_inputs=True, img_temporal_dim_size=1,
        interpolate_rope=True, modality_embedding=True,
        n_output_distillation=n_distill, out_layers=out_layers,
    ).eval()
    predictor = VisionTransformerPredictor(
        img_size=(IMG, IMG), patch_size=PATCH, num_frames=FRAMES, tubelet_size=TUBELET,
        embed_dim=DIM, predictor_embed_dim=PRED_DIM, depth=PRED_DEPTH,
        num_heads=PRED_HEADS, mlp_ratio=4.0, qkv_bias=True, norm_layer=NORM,
        use_mask_tokens=True, num_mask_tokens=8, zero_init_mask_tokens=True,
        use_silu=False, wide_silu=True, use_rope=True, interpolate_rope=True,
        modality_embedding=True, img_temporal_dim_size=1, uniform_power=False,
        return_all_tokens=True, teacher_embed_dim=teacher_embed_dim,
        n_output_distillation=n_distill,
    ).eval()
    return encoder, predictor


def build_port(n_distill=1, teacher_embed_dim=None):
    cfg = VJEPA21Config(
        patch_size=PATCH, crop_size=IMG, frames_per_clip=FRAMES, tubelet_size=TUBELET,
        hidden_size=DIM, num_attention_heads=HEADS, num_hidden_layers=DEPTH,
        mlp_ratio=4.0, num_pooler_heads=HEADS,
        pred_hidden_size=PRED_DIM, pred_num_attention_heads=PRED_HEADS,
        pred_num_hidden_layers=PRED_DEPTH, pred_num_mask_tokens=8,
        pred_teacher_embed_dim=teacher_embed_dim, pred_return_all_tokens=True,
        n_output_distillation=n_distill, img_temporal_dim_size=1,
        interpolate_rope=True, modality_embedding=True,
        attn_implementation="eager",
    )
    return VJEPA21Model(cfg).eval()


def make_pair(n_distill=1, out_layers=None, teacher_embed_dim=None):
    """Reference pair plus a port loaded with exactly those weights."""
    ref_enc, ref_pred = build_reference(n_distill, out_layers, teacher_embed_dim)
    port = build_port(n_distill, teacher_embed_dim)
    state = reference_to_port(ref_enc.state_dict(), ref_pred.state_dict(), DIM, PRED_DIM)
    missing, unexpected = port.load_state_dict(state, strict=False)
    # This is the check that matters: a parameter with no counterpart stays at its
    # random initialisation and nothing else in the suite would notice.
    assert not missing, f"port parameters with no reference origin: {sorted(missing)[:10]}"
    assert not unexpected, f"reference tensors the port has no home for: {sorted(unexpected)[:10]}"
    return ref_enc, ref_pred, port


def maxdiff(a, b):
    return (a - b).abs().max().item()


@pytest.fixture(scope="module")
def video():
    torch.manual_seed(1)
    return torch.randn(2, 3, FRAMES, IMG, IMG)


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


@torch.no_grad()
@pytest.mark.parametrize("n_distill", [1, 4])
def test_encoder_video_branch(video, n_distill):
    ref_enc, _, port = make_pair(n_distill=n_distill)
    expected = ref_enc(video)
    got = port(pixel_values_videos=video, skip_predictor=True).last_hidden_state
    d = maxdiff(expected, got)
    print(f"\n[encoder n_distill={n_distill}] max|Δ| = {d:.3e}")
    assert d < 1e-5


@torch.no_grad()
def test_encoder_hierarchical_output(video):
    ref_enc, _, port = make_pair(n_distill=4)
    expected = ref_enc(video, training=True)          # concatenated levels
    got = port(
        pixel_values_videos=video, skip_predictor=True, return_hierarchical=True
    ).hierarchical_hidden_state
    assert expected.shape == got.shape == (2, 32, DIM * 4)
    d = maxdiff(expected, got)
    print(f"[encoder hierarchical] max|Δ| = {d:.3e}")
    assert d < 1e-5


@torch.no_grad()
def test_encoder_image_branch():
    ref_enc, _, port = make_pair()
    torch.manual_seed(2)
    image = torch.randn(2, 3, 1, IMG, IMG)
    expected = ref_enc(image)
    got = port(pixel_values_videos=image, skip_predictor=True).last_hidden_state
    assert got.shape == (2, 16, DIM)
    d = maxdiff(expected, got)
    print(f"[encoder image branch] max|Δ| = {d:.3e}")
    assert d < 1e-5


@torch.no_grad()
def test_encoder_non_square_input():
    """Exercises the RoPE interpolation on both spatial axes independently."""
    ref_enc, _, port = make_pair()
    torch.manual_seed(3)
    x = torch.randn(1, 3, FRAMES, 96, 128)
    expected = ref_enc(x)
    got = port(pixel_values_videos=x, skip_predictor=True).last_hidden_state
    assert got.shape == (1, 2 * 6 * 8, DIM)
    d = maxdiff(expected, got)
    print(f"[encoder 96x128] max|Δ| = {d:.3e}")
    assert d < 1e-5


@torch.no_grad()
@pytest.mark.parametrize("n_distill", [1, 4])
def test_encoder_out_layers(video, n_distill):
    """`out_layers` is a constructor argument on the reference and a forward
    argument on the port; the per-level norms must still agree."""
    layers = [2, 5, 8, 11]
    ref_enc, _, port = make_pair(n_distill=n_distill, out_layers=layers)
    expected = ref_enc(video)                          # list, one tensor per level
    got = port(
        pixel_values_videos=video, skip_predictor=True, out_layers=layers
    ).multilevel_hidden_states
    assert len(expected) == len(got) == 4
    for i, (a, b) in enumerate(zip(expected, got)):
        d = maxdiff(a, b)
        print(f"[out_layers n_distill={n_distill}] level {layers[i]}: max|Δ| = {d:.3e}")
        assert d < 1e-5


@torch.no_grad()
def test_masked_encoder_forward(video):
    """The JEPA training forward: tokens are dropped before the blocks and RoPE
    receives their original indices. Without this the predictor can only ever be
    fed representations computed with attention over tokens that were masked
    during pre-training.
    """
    ref_enc, _, port = make_pair()
    torch.manual_seed(4)
    idx = torch.stack([torch.randperm(32)[:14] for _ in range(2)])
    expected = ref_enc(video, masks=[idx])
    got = port.encoder(video, masks=[idx]).last_hidden_state
    assert got.shape == (2, 14, DIM)
    d = maxdiff(expected, got)
    print(f"[masked encoder] max|Δ| = {d:.3e}")
    assert d < 1e-5


@torch.no_grad()
def test_masked_encoder_is_not_equivalent_to_masking_the_output(video):
    """Control: gathering after the blocks gives a different answer, so the test
    above is checking something real."""
    ref_enc, _, port = make_pair()
    idx = torch.arange(0, 14).unsqueeze(0).expand(2, -1)
    inside = port.encoder(video, masks=[idx]).last_hidden_state
    outside = apply_masks(
        port(pixel_values_videos=video, skip_predictor=True).last_hidden_state, [idx]
    )
    assert not torch.allclose(inside, outside, atol=1e-4)


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


def _masks(batch=2, n_tokens=32, n_ctx=20):
    ctx = torch.arange(0, n_ctx).unsqueeze(0).expand(batch, -1)
    tgt = torch.arange(n_ctx, n_tokens).unsqueeze(0).expand(batch, -1)
    return ctx, tgt


@torch.no_grad()
@pytest.mark.parametrize(
    "n_distill,teacher", [(1, DIM), (1, None), (4, None), (4, DIM * 4)]
)
def test_predictor(video, n_distill, teacher):
    ref_enc, ref_pred, port = make_pair(n_distill=n_distill, teacher_embed_dim=teacher)
    ctx, tgt = _masks()

    # The predictor input is the concatenated levels when the encoder distils
    # more than one, and the last hidden state otherwise.
    z = ref_enc(video, training=True) if n_distill > 1 else ref_enc(video)

    exp_pred, exp_ctx = ref_pred(apply_masks(z, [ctx]), [ctx], [tgt], mod="video")
    got = port.predictor(z, [ctx], [tgt], mode="video")

    dp, dc = maxdiff(exp_pred, got.last_hidden_state), maxdiff(exp_ctx, got.context_hidden_state)
    print(f"\n[predictor n_distill={n_distill} teacher={teacher}] "
          f"target max|Δ| = {dp:.3e}   context max|Δ| = {dc:.3e}")
    assert dp < 1e-3 and dc < 1e-3


@torch.no_grad()
def test_predictor_with_trained_mask_tokens(video):
    """The shipped checkpoints have `zero_init_mask_tokens=True`, so every mask
    token is exactly zero and any comparison of the mask-token path passes even
    if the lookup is broken. Give them distinct values first, then compare.
    """
    ref_enc, ref_pred, port = make_pair()
    torch.manual_seed(5)
    for i, token in enumerate(ref_pred.mask_tokens):
        nn.init.normal_(token, std=0.1 * (i + 1))
    # copy the same values into the port
    port_tokens = port.predictor.embeddings.mask_tokens
    for dst, src in zip(port_tokens, ref_pred.mask_tokens):
        dst.copy_(src)
    assert max(t.abs().max().item() for t in ref_pred.mask_tokens) > 0

    ctx, tgt = _masks()
    z = ref_enc(video)
    for mask_index in (0, 1, 3, 7, 9):  # 9 exercises the modulo wrap
        exp_pred, exp_ctx = ref_pred(
            apply_masks(z, [ctx]), [ctx], [tgt], mod="video", mask_index=mask_index
        )
        got = port.predictor(z, [ctx], [tgt], mode="video", mask_index=mask_index)
        dp = maxdiff(exp_pred, got.last_hidden_state)
        dc = maxdiff(exp_ctx, got.context_hidden_state)
        print(f"[mask_index={mask_index}] target max|Δ| = {dp:.3e}  context max|Δ| = {dc:.3e}")
        assert dp < 1e-3 and dc < 1e-3

    # and the indices must actually differ from one another
    a = port.predictor(z, [ctx], [tgt], mask_index=0).last_hidden_state
    b = port.predictor(z, [ctx], [tgt], mask_index=1).last_hidden_state
    assert not torch.allclose(a, b, atol=1e-4)


@torch.no_grad()
def test_predictor_image_modality(video):
    """The reference spells the video/image switch `mod="image"`; the port uses
    `mode="img"` and accepts "image" as an alias. A mismatch here silently adds
    the video modality embedding to an image."""
    ref_enc, ref_pred, port = make_pair()
    torch.manual_seed(6)
    image = torch.randn(2, 3, 1, IMG, IMG)
    z = ref_enc(image)
    ctx = torch.arange(0, 10).unsqueeze(0).expand(2, -1)
    tgt = torch.arange(10, 16).unsqueeze(0).expand(2, -1)

    exp_pred, _ = ref_pred(apply_masks(z, [ctx]), [ctx], [tgt], mod="image")
    for alias in ("img", "image"):
        got = port.predictor(z, [ctx], [tgt], mode=alias)
        assert maxdiff(exp_pred, got.last_hidden_state) < 1e-3, alias

    # and the video path must give something different
    video_out = port.predictor(z, [ctx], [tgt], mode="video")
    assert not torch.allclose(exp_pred, video_out.last_hidden_state, atol=1e-4)


@torch.no_grad()
def test_full_jepa_forward_matches_reference(video):
    """Encoder-with-masks followed by the predictor, i.e. the pre-training
    forward end to end."""
    ref_enc, ref_pred, port = make_pair()
    ctx, tgt = _masks()

    z_ref = ref_enc(video, masks=[ctx])                # already gathered
    exp_pred, exp_ctx = ref_pred(z_ref, [ctx], [tgt], mod="video")

    out = port(
        pixel_values_videos=video, masks=[ctx], context_mask=[ctx], target_mask=[tgt]
    )
    dp = maxdiff(exp_pred, out.predictor_output.last_hidden_state)
    dc = maxdiff(exp_ctx, out.predictor_output.context_hidden_state)
    print(f"\n[full JEPA forward] target max|Δ| = {dp:.3e}   context max|Δ| = {dc:.3e}")
    assert dp < 1e-3 and dc < 1e-3