import copy
from types import SimpleNamespace

import pytest

from model.wan_kv_mmd import WanKVMMD


def _triple_state_only():
    model = object.__new__(WanKVMMD)
    model.feature_type = "triple_dinov3_vjepa2_videomae_v2_dense_latent"
    model.is_dinov3_dense = False
    model.is_triple_dense = True
    model.uses_fresh_video_gradcache = True
    model.dense_feature_prefix = "triple"
    model.dense_encoder_name = "triple"
    model.dense_log_prefix = "triple_dense"
    model.vjepa2_multilayer_enabled = False
    model.mmd_estimator = "nystrom_exact_repulsion"
    model.triple_teacher_bank_content_hash = "a" * 64
    model.dual_teacher_bank_content_hash = "a" * 64
    model.dual_branch_registry = {
        "dinov3": {
            "reference_path": "/data/dino.pt",
            "reference_hash": "b" * 64,
            "nystrom_landmarks": 256,
        },
        "vjepa2": {
            "reference_path": "/data/vj.pt",
            "reference_hash": "c" * 64,
            "nystrom_landmarks": 1024,
        },
        "videomae_v2": {
            "reference_path": "/data/vm.pt",
            "reference_hash": "d" * 64,
            "nystrom_landmarks": 256,
        },
    }
    model.latent_mmd = SimpleNamespace(
        attraction_ema_enabled=False,
        repulsion_fifo_enabled=False,
    )
    model.distribution_state_pending = lambda: False
    model._fd_metadata = lambda: {
        "feature_type": model.feature_type,
        "teacher_bank_content_hash": model.triple_teacher_bank_content_hash,
    }
    return model


def test_triple_fd_state_round_trip_uses_own_namespace_and_version_10():
    model = _triple_state_only()
    state = model.fd_state_dict()
    assert state["triple_dense_nystrom_mmd_version"] == 10
    assert "vjepa2_dense_nystrom_mmd_version" not in state
    model.load_fd_state_dict(copy.deepcopy(state))
    assert model.can_load_initialization_fd_state(state)


def test_triple_fd_state_rejects_foreign_or_wrong_version():
    model = _triple_state_only()
    state = model.fd_state_dict()
    wrong = copy.deepcopy(state)
    wrong["triple_dense_nystrom_mmd_version"] = 9
    with pytest.raises(ValueError, match="triple_dense_nystrom_mmd_version"):
        model.load_fd_state_dict(wrong)
    foreign = copy.deepcopy(state)
    foreign["vjepa2_dense_nystrom_mmd_version"] = 10
    with pytest.raises(ValueError, match="another dense feature namespace"):
        model.load_fd_state_dict(foreign)