from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import nn


class CanonicalFactPrototype(nn.Module):
    """Compose a typed symbolic state into a learned canonical prototype."""

    def __init__(
        self,
        state_dim: int,
        num_state_types: int,
        num_subjects: int,
        num_predicates: int,
        num_objects: int,
    ):
        super().__init__()
        self.type_embedding = nn.Embedding(num_state_types, state_dim)
        self.subject_embedding = nn.Embedding(num_subjects + 1, state_dim, padding_idx=0)
        self.predicate_embedding = nn.Embedding(num_predicates + 1, state_dim, padding_idx=0)
        self.object_embedding = nn.Embedding(num_objects + 1, state_dim, padding_idx=0)
        self.consequent_predicate_embedding = nn.Embedding(
            num_predicates + 1, state_dim, padding_idx=0
        )
        self.consequent_object_embedding = nn.Embedding(
            num_objects + 1, state_dim, padding_idx=0
        )
        self.projection = nn.Sequential(
            nn.LayerNorm(state_dim),
            nn.Linear(state_dim, state_dim),
            nn.GELU(),
            nn.Linear(state_dim, state_dim),
        )

    def forward(self, targets: torch.Tensor) -> torch.Tensor:
        shifted = targets[:, 1:].clamp_min(-1) + 1
        value = self.type_embedding(targets[:, 0])
        value = value + self.subject_embedding(shifted[:, 0])
        value = value + self.predicate_embedding(shifted[:, 1])
        value = value + self.object_embedding(shifted[:, 2])
        value = value + self.consequent_predicate_embedding(shifted[:, 3])
        value = value + self.consequent_object_embedding(shifted[:, 4])
        return F.normalize(self.projection(value), dim=-1)