import torch

from hippotorch.encoders.dual import DualEncoder


def test_dual_encoder_target_updates_via_ema() -> None:
    encoder = DualEncoder(input_dim=8, embed_dim=16, momentum=0.5, refresh_interval=1)
    # Target params start frozen
    assert all(not p.requires_grad for p in encoder.target.parameters())

    batch = torch.randn(2, 5, 8)
    optimizer = torch.optim.SGD(encoder.online.parameters(), lr=0.1)
    optimizer.zero_grad()

    baseline_online = [p.clone().detach() for p in encoder.online.parameters()]
    baseline_target = [p.clone().detach() for p in encoder.target.parameters()]

    loss = encoder.encode_query(batch).sum()
    loss.backward()
    optimizer.step()

    updated_online = [p.clone().detach() for p in encoder.online.parameters()]
    assert any(
        not torch.allclose(before, after)
        for before, after in zip(baseline_online, updated_online)
    )
    # Target does not change until EMA update is invoked.
    assert all(
        torch.allclose(before, after)
        for before, after in zip(baseline_target, encoder.target.parameters())
    )

    encoder.update_target()
    assert any(
        not torch.allclose(before, after)
        for before, after in zip(baseline_target, encoder.target.parameters())
    )
