import torch

from data.schema import FACT_TYPE, RULE_TYPE
from losses.alignment_loss import cross_modal_alignment_loss
from losses.matching_loss import hungarian_matching_loss
from losses.prototype_loss import canonical_prototype_loss
from models.fact_prototype import CanonicalFactPrototype
from models.slot_extractor import SlotExtractor


def test_slot_extractor_and_matching_loss_backward():
    extractor = SlotExtractor(
        input_dim=32,
        state_dim=64,
        num_slots=4,
        num_layers=2,
        num_heads=4,
        dropout=0.0,
        num_subjects=12,
        num_predicates=7,
        num_objects=12,
        num_state_types=3,
    )
    tokens = torch.randn(2, 10, 32)
    mask = torch.ones(2, 10, dtype=torch.bool)
    predictions = extractor(tokens, mask)
    targets = [
        torch.tensor(
            [
                [FACT_TYPE, 0, 0, 4, -1, -1],
                [RULE_TYPE, 1, 1, 5, 6, 7],
            ]
        ),
        torch.tensor([[FACT_TYPE, 1, 0, 6, -1, -1]]),
    ]
    result = hungarian_matching_loss(predictions, targets)
    assert torch.isfinite(result.loss)
    result.loss.backward()
    assert extractor.slot_queries.grad is not None


def test_cross_modal_alignment_and_prototype_losses():
    assignments = [(torch.tensor([0]), torch.tensor([0]))]
    targets = [torch.tensor([[FACT_TYPE, 0, 0, 4, -1, -1]])]
    visual = torch.randn(1, 2, 16, requires_grad=True)
    text = torch.randn(1, 2, 16, requires_grad=True)
    alignment = cross_modal_alignment_loss(
        visual, targets, assignments, text, targets, assignments
    )
    prototype = CanonicalFactPrototype(16, 3, 12, 7, 12)
    prototype_value = canonical_prototype_loss(visual, targets, assignments, prototype)
    loss = alignment + prototype_value
    assert torch.isfinite(loss)
    loss.backward()
    assert visual.grad is not None