import pytest

torch = pytest.importorskip("torch")

from hippotorch.encoders.dual import DualEncoder
from hippotorch.memory.replay import HippocampalReplayBuffer


def test_deserialize_accepts_dataparallel_state_dict():
    # Build a simple encoder and grab its serialized config/state
    enc = DualEncoder(input_dim=6, embed_dim=8, momentum=0.9, refresh_interval=7)
    payload = {
        "type": "DualEncoder",
        "config": {
            "input_dim": 6,
            "embed_dim": 8,
            "momentum": 0.9,
            "refresh_interval": 7,
        },
        "state_dict": {},
    }

    # Simulate DataParallel by prefixing '.module.' under online/target scopes
    sd = enc.state_dict()
    dp_sd = {}
    for k, v in sd.items():
        if k.startswith("online."):
            dp_sd[k.replace("online.", "online.module.", 1)] = v
        elif k.startswith("target."):
            dp_sd[k.replace("target.", "target.module.", 1)] = v
        else:
            dp_sd[k] = v
    payload["state_dict"] = dp_sd

    restored = HippocampalReplayBuffer._deserialize_encoder(payload, device="cpu")

    x = torch.randn(2, 3, 6)
    out_orig = enc.encode_query(x)
    out_restored = restored.encode_query(x)
    assert out_orig.shape == out_restored.shape
