"""Constraint-membership embeddings with no board geometry assumptions."""

from collections.abc import Sequence

import torch
from torch import nn


def membership_embeddings(
    tables: Sequence[nn.Module], memberships: Sequence[torch.Tensor]
) -> torch.Tensor:
    """Sum family embeddings in their declared order (also preserves legacy order)."""
    if not tables or len(tables) != len(memberships):
        raise ValueError("Provide one membership index tensor per embedding table")
    result = tables[0](memberships[0])
    for table, indices in zip(tables[1:], memberships[1:], strict=True):
        result = result + table(indices)
    return result


class GroupMembershipEmbedding(nn.Module):
    """Embed static memberships; input shape is (families, positions)."""

    memberships: torch.Tensor

    def __init__(
        self, group_counts: Sequence[int], memberships: torch.Tensor, d_model: int
    ) -> None:
        super().__init__()
        if memberships.ndim != 2 or len(group_counts) != memberships.shape[0]:
            raise ValueError("Memberships must have shape (families, positions)")
        if not group_counts or memberships.dtype != torch.long:
            raise ValueError("At least one family with integer memberships is required")
        for count, indices in zip(group_counts, memberships, strict=True):
            if count <= 0 or torch.any((indices < 0) | (indices >= count)):
                raise ValueError("Membership index outside its family")
        self.tables = nn.ModuleList([nn.Embedding(count, d_model) for count in group_counts])
        self.register_buffer("memberships", memberships.clone())

    def forward(self) -> torch.Tensor:
        return membership_embeddings(list(self.tables), list(self.memberships.unbind()))