"""Robustness checks relevant to downstream experiments.

These are not port-correctness tests (see `test_parity_official.py` for those).
They check properties the experimental pipeline depends on: that cached features
equal on-the-fly features, that reduced precision does not silently change the
representation, and that a merged PEFT adapter is exactly the baseline model.

Runs on a tiny random model by default, so it can be part of CI:

    python -m pytest test_thesis_robustness.py -s -q

Re-run it against a published checkpoint before trusting a feature store:

    VJEPA21_CKPT=apiantonio/vjepa2.1-vit-base-384 python -m pytest test_thesis_robustness.py -s -q
"""

from __future__ import annotations

import copy
import importlib
import os
import sys
import types

import pytest
import torch

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


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.VJEPA21ForVideoClassification


VJEPA21Config, VJEPA21Model, VJEPA21ForVideoClassification = _load_port()

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def tiny_config(**kw):
    base = dict(
        patch_size=16, crop_size=64, hidden_size=96, num_attention_heads=6,
        num_hidden_layers=12, num_pooler_heads=6, pred_hidden_size=48,
        pred_num_attention_heads=6, pred_num_hidden_layers=12,
        pred_teacher_embed_dim=96, n_output_distillation=1, num_labels=5,
        attn_implementation="sdpa",
    )
    base.update(kw)
    return VJEPA21Config(**base)


@pytest.fixture(scope="module")
def model():
    if CKPT:
        from transformers import AutoModel

        return AutoModel.from_pretrained(CKPT, trust_remote_code=True).eval().to(DEVICE)
    torch.manual_seed(0)
    return VJEPA21Model(tiny_config()).eval().to(DEVICE)


@pytest.fixture(scope="module")
def clip(model):
    res = model.config.crop_size
    torch.manual_seed(1)
    return torch.randn(8, 3, 4, res, res, device=DEVICE)


def _rel(a: torch.Tensor, b: torch.Tensor) -> float:
    return ((a - b).abs().mean() / a.abs().mean()).item()


def _cos(a: torch.Tensor, b: torch.Tensor) -> float:
    return torch.nn.functional.cosine_similarity(a.flatten(0, 1), b.flatten(0, 1)).min().item()


# --- 0. the fixture must stay pristine ---------------------------------------


@torch.no_grad()
def test_module_fixture_is_not_mutated(model, clip):
    """Guard for the bug this file used to contain.

    `test_reduced_precision_gap` called `model.to(bfloat16)` on a module-scoped
    fixture and then `model.to(float32)` to undo it. Casting back does not
    restore the mantissa: every subsequent test saw bf16-rounded weights stored
    in float32, and the "fp16 vs fp32" number it reported was really "fp16 vs
    bf16-rounded", which is far smaller than the truth. Fingerprinting the
    weights here makes any repeat of that failure loud.
    """
    fingerprint = torch.stack(
        [p.detach().float().flatten()[:8].cpu() for p in list(model.parameters())[:5]]
    )
    torch.testing.assert_close(
        fingerprint,
        torch.stack(
            [p.detach().float().flatten()[:8].cpu() for p in list(model.parameters())[:5]]
        ),
    )
    assert all(p.dtype is torch.float32 for p in model.parameters()), (
        "the shared model fixture is no longer float32; some test cast it in place"
    )


# --- 1. cached features must equal on-the-fly features ------------------------


@torch.no_grad()
def test_batch_size_invariance(model, clip):
    """A feature store is precomputed in batches and consumed one clip at a time.

    This is not guaranteed to be bit-exact on any device: BLAS and the fused
    attention kernels pick tiling and reduction order as a function of the input
    shape, so the same clip can accumulate in a different order at batch 1 and at
    batch 8. The assertion is therefore on relative error and cosine similarity.
    `test_determinism_across_calls` covers the same-shape case, which must be
    exact everywhere.
    """
    full = model(pixel_values_videos=clip, skip_predictor=True).last_hidden_state
    single = torch.cat(
        [model(pixel_values_videos=clip[i : i + 1], skip_predictor=True).last_hidden_state
         for i in range(clip.shape[0])]
    )
    max_abs = (full - single).abs().max().item()
    rel, cos = _rel(full, single), _cos(full, single)
    print(f"\n[batch invariance]   max|Δ| = {max_abs:.3e}   rel = {rel:.3e}   "
          f"min cos-sim = {cos:.6f}   device = {DEVICE}   bit-exact = {max_abs == 0.0}")
    assert rel < 1e-4
    assert cos > 0.9999


@torch.no_grad()
def test_determinism_across_calls(model, clip):
    """Two identical forwards in eval mode must be bit-identical, otherwise a
    cached feature is not reproducible from the same input."""
    a = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
    b = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
    print(f"[determinism]        bit-identical = {torch.equal(a, b)}")
    assert torch.equal(a, b)


# --- 2. reduced precision -----------------------------------------------------


@torch.no_grad()
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
def test_reduced_precision_gap(model, clip, dtype):
    """Training runs in bf16 or fp16. Measure the gap against the fp32 reference
    once and record it, rather than assuming it is negligible.

    The low-precision copy is a `deepcopy`: `.to(dtype)` is in-place on the
    module, and casting back to float32 afterwards does *not* recover the
    discarded mantissa bits. Mutating the shared fixture would corrupt every
    later test and, worse, make the second parametrisation measure the gap
    against an already-degraded reference.
    """
    if dtype is torch.float16 and DEVICE == "cpu":
        pytest.skip("fp16 matmul is not supported on CPU")

    ref = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
    assert ref.dtype is torch.float32, "the reference forward is not float32"

    low = copy.deepcopy(model).to(dtype)
    try:
        got = low(pixel_values_videos=clip[:2].to(dtype), skip_predictor=True).last_hidden_state
        assert got.dtype is dtype, f"output dtype drifted to {got.dtype}"
        got = got.float()
    finally:
        del low
        if DEVICE == "cuda":
            torch.cuda.empty_cache()

    rel, cos = _rel(ref, got), _cos(ref, got)
    print(f"[{str(dtype).split('.')[-1]:>9}]        rel = {rel:.3e}   min cos-sim = {cos:.6f}")
    assert torch.isfinite(got).all()
    assert cos > 0.99


@torch.no_grad()
def test_feature_store_fp16_roundtrip(model, clip):
    """Storing fp32 features as fp16 halves the cache. Quantify what it costs."""
    ref = model(pixel_values_videos=clip[:2], skip_predictor=True).last_hidden_state
    back = ref.half().float()
    print(f"[fp16 storage]       rel = {_rel(ref, back):.3e}   min cos-sim = {_cos(ref, back):.6f}")
    assert _cos(ref, back) > 0.9999


# --- 3. the central PEFT claim of the thesis ----------------------------------
# These run on a small randomly initialised model regardless of VJEPA21_CKPT:
# exact merging is a property of the reparameterisation, not of the weights, and
# holding a published checkpoint plus a second classification copy does not fit
# in a notebook runtime for the larger variants.

TARGET_MODULES = r".*vjepa21\.encoder\.layer\.\d+\.attention\.(query|key|value|proj)$"


def _peft_model(config, r=8, use_dora=False):
    from peft import LoraConfig, get_peft_model

    torch.manual_seed(0)
    base = VJEPA21ForVideoClassification(config).eval()
    adapter = LoraConfig(
        r=r, lora_alpha=2 * r, lora_dropout=0.0, use_dora=use_dora,
        target_modules=TARGET_MODULES,
        modules_to_save=["classifier", "pooler"],
    )
    peft_model = get_peft_model(base, adapter)
    # Give the adapters non-zero B matrices, otherwise the merge is trivially
    # exact: peft initialises lora_B to zero, so the adapter is the identity.
    for name, param in peft_model.named_parameters():
        if "lora_B" in name:
            torch.nn.init.normal_(param, std=0.02)
    return peft_model


@torch.no_grad()
@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
def test_merge_is_exact(use_dora):
    """`merge_and_unload` must not change the function computed. This underpins
    the claim that reparameterisation methods have zero inference cost.

    DoRA is the case worth testing: it decomposes the update into direction and
    magnitude, so merging has to fold the magnitude vector back in as well. The
    model card claimed DoRA was verified while only LoRA was covered.
    """
    pytest.importorskip("peft")
    cfg = tiny_config()
    peft_model = _peft_model(cfg, use_dora=use_dora).eval()
    x = torch.randn(2, 3, 4, cfg.crop_size, cfg.crop_size)

    before = peft_model(pixel_values_videos=x).logits
    merged = peft_model.merge_and_unload().eval()
    after = merged(pixel_values_videos=x).logits

    max_abs = (before - after).abs().max().item()
    label = "DoRA" if use_dora else "LoRA"
    print(f"\n[{label} merge, tiny]   max|Δ| = {max_abs:.3e}   rel = {_rel(before, after):.3e}")
    assert max_abs < 1e-4


@torch.no_grad()
@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
def test_merged_model_has_baseline_parameter_count(use_dora):
    """After merging there must be no residual adapter parameters: the merged
    model is parameter-identical to the untouched baseline. This is the
    zero-inference-cost claim in its structural form, and it is what prompt-based
    methods cannot satisfy."""
    pytest.importorskip("peft")
    cfg = tiny_config()
    baseline = VJEPA21ForVideoClassification(cfg)
    merged = _peft_model(cfg, use_dora=use_dora).merge_and_unload()

    n_base = sum(p.numel() for p in baseline.parameters())
    n_merged = sum(p.numel() for p in merged.parameters())
    leftover = [n for n, _ in merged.named_parameters() if "lora" in n.lower()]
    label = "DoRA" if use_dora else "LoRA"
    print(f"[{label} merged params]  baseline {n_base:,} vs merged {n_merged:,}   "
          f"residual adapter tensors: {len(leftover)}")
    assert n_merged == n_base
    assert not leftover


@torch.no_grad()
@pytest.mark.parametrize("use_dora", [False, True], ids=["lora", "dora"])
def test_unmerged_adapter_adds_parameters(use_dora):
    """Control for the test above: before merging the adapters really are extra
    parameters, so the equality afterwards is meaningful."""
    pytest.importorskip("peft")
    cfg = tiny_config()
    peft_model = _peft_model(cfg, use_dora=use_dora)
    n_base = sum(p.numel() for p in VJEPA21ForVideoClassification(cfg).parameters())
    n_peft = sum(p.numel() for p in peft_model.parameters())
    label = "DoRA" if use_dora else "LoRA"
    print(f"[{label} unmerged]       baseline {n_base:,} vs adapted {n_peft:,} "
          f"(+{n_peft - n_base:,})")
    assert n_peft > n_base


def test_target_modules_regex_does_not_touch_the_predictor():
    """`VJEPA21ForVideoClassification` never runs the predictor, so adapting it
    would spend trainable parameters on dead weight. Matching on bare names such
    as ["query", "key", "value"] does exactly that."""
    pytest.importorskip("peft")
    import re

    cfg = tiny_config()
    model = VJEPA21ForVideoClassification(cfg)
    pattern = re.compile(TARGET_MODULES)
    matched = [n for n, _ in model.named_modules() if pattern.match(n)]
    assert matched, "the target_modules regex matches nothing"
    assert not any(".predictor." in n for n in matched)
    assert len(matched) == cfg.num_hidden_layers * 4