# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Weight-free compatibility tests for Transformers 4 and 5.

Load standalone remote-code modules without model weights. V5-specific
conversion tests are skipped when the conversion API is unavailable.
"""

import importlib.util
import sys
from pathlib import Path

import pytest
import torch

ROOT = Path(__file__).resolve().parent.parent


def _load(name: str):
    spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


modeling_vit = _load("modeling_vit")
modeling_qformer = _load("modeling_qformer")


def test_qformer_initialization_and_attention_masks() -> None:
    from transformers import BertConfig

    config = BertConfig(
        hidden_size=16, num_hidden_layers=1, num_attention_heads=2,
        intermediate_size=32, vocab_size=64,
    )
    model = modeling_qformer.BertModel(config).eval()
    if modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API:
        assert hasattr(model, "all_tied_weights_keys")
    assert model.get_head_mask(None, 1) == [None]
    mask = model.invert_attention_mask(torch.tensor([[1, 0]]))
    assert mask.shape == (1, 1, 1, 2)
    assert mask[0, 0, 0, 0] == 0
    assert mask[0, 0, 0, 1] < 0


def test_pruning_helper_is_vendored() -> None:
    heads, index = modeling_qformer.find_pruneable_heads_and_indices([0, 1], 4, 2, set())
    assert heads == {0, 1}
    assert index.tolist() == [4, 5, 6, 7]


def test_bicubic_interpolation_matches_reference() -> None:
    pos_embed = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
    out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
    assert out.shape == (1, 577, 8)

    # Independent reference: same math as the historical interpolate_pos_embed.
    extra = pos_embed[:, :1]
    tokens = pos_embed[:, 1:].reshape(-1, 16, 16, 8).permute(0, 3, 1, 2)
    tokens = torch.nn.functional.interpolate(tokens, size=(24, 24), mode="bicubic", align_corners=False)
    tokens = tokens.permute(0, 2, 3, 1).flatten(1, 2)
    reference = torch.cat((extra, tokens), dim=1)
    assert torch.equal(out, reference)
    assert torch.equal(out[:, :1], pos_embed[:, :1])  # cls token untouched


def test_bicubic_interpolation_is_noop_on_matching_grid() -> None:
    pos_embed = torch.rand(1, 577, 8)
    out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
    assert torch.equal(out, pos_embed)


def test_interpolation_rejects_non_square_source_and_resizes_equal_area_rectangle() -> None:
    with pytest.raises(ValueError, match="square patch grid"):
        modeling_vit._bicubic_interpolate_pos_embed(torch.rand(1, 258, 8), 1, 24, 24)
    source = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
    resized = modeling_vit._bicubic_interpolate_pos_embed(source, 1, 8, 32)
    assert resized.shape == source.shape
    assert not torch.equal(resized[:, 1:], source[:, 1:])
    state_dict = {"visual_encoder.pos_embed": source.clone()}
    modeling_vit.interpolate_pos_embed(
        "visual_encoder.pos_embed", 256, 257, state_dict, target_h=8, target_w=32
    )
    assert torch.equal(state_dict["visual_encoder.pos_embed"], resized)


def test_interpolate_pos_embed_writes_only_on_mismatch() -> None:
    state_dict = {"visual_encoder.pos_embed": torch.rand(1, 257, 8)}
    modeling_vit.interpolate_pos_embed(
        "visual_encoder.pos_embed", 256, 257, state_dict, target_h=16, target_w=16
    )
    assert state_dict["visual_encoder.pos_embed"].shape == (1, 257, 8)  # no-op


@pytest.mark.skipif(
    not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
    reason="Transformers 5 conversion API not available",
)
def test_v5_conversion_mapping_registered() -> None:
    from transformers.conversion_mapping import get_checkpoint_conversion_mapping

    modeling_vit.register_transformers5_pos_embed_conversion()
    converters = get_checkpoint_conversion_mapping("cosmos-embed1")
    assert converters and len(converters) == 1
    assert converters[0].source_patterns == ["visual_encoder.pos_embed"]


@pytest.mark.skipif(
    not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
    reason="Transformers 5 conversion API not available",
)
def test_v5_operation_interpolates_257_to_577() -> None:
    class _PatchEmbed:
        patch_shape = (24, 24)

    class _Visual:
        patch_embed = _PatchEmbed()

    class _Model:
        visual_encoder = _Visual()

    op = modeling_vit.InterpolatePosEmbed()
    tensor = torch.arange(1 * 257 * 4, dtype=torch.bfloat16).reshape(1, 257, 4)
    out = op.convert(
        {"visual_encoder.pos_embed": [tensor]},
        ["visual_encoder.pos_embed"],
        ["visual_encoder.pos_embed"],
        model=_Model(),
    )
    result = out["visual_encoder.pos_embed"][0]
    assert result.shape == (1, 577, 4)
    assert result.dtype == torch.bfloat16