"""The exclusivity operator over explicit groups, independent of board geometry."""

from typing import NamedTuple

import torch
from torch import nn


class OperatorEvaluation(NamedTuple):
    probabilities: torch.Tensor
    consequences: torch.Tensor


def exclusivity_consequence(
    probabilities: torch.Tensor,
    log_complement: torch.Tensor,
    member_dims: int | tuple[int, ...],
) -> torch.Tensor:
    """Exclusion product on group-member axes of a supplied tensor layout.

    Index-based groups and structured partition views share this kernel. Views
    let compatibility adapters preserve reduction order and autograd accumulation.
    """
    excluded = torch.exp(log_complement.sum(dim=member_dims, keepdim=True) - log_complement)
    return probabilities * excluded


class GroupExclusivityOperator(nn.Module):
    """One rectangular family of groups; different families may have different sizes.

    This is the exclusivity family of the probabilistic lift, not an interpreter
    for arbitrary ASP programs. Indices are derived configuration, not weights.
    """

    groups: torch.Tensor

    def __init__(self, groups: torch.Tensor, num_positions: int, eps: float = 1e-7) -> None:
        super().__init__()
        if groups.ndim != 2 or groups.numel() == 0 or groups.dtype != torch.long:
            raise ValueError("Groups must be a nonempty integer (groups, members) tensor")
        if num_positions < 1 or torch.any((groups < 0) | (groups >= num_positions)):
            raise ValueError("Group member outside the declared position range")
        if any(len(set(row)) != len(row) for row in groups.tolist()):
            raise ValueError("A group cannot contain the same position twice")
        if not 0 < eps < 1:
            raise ValueError("eps must lie strictly between zero and one")
        self.num_positions = num_positions
        self.eps = eps
        self.register_buffer("groups", groups.clone(), persistent=False)

    def forward(self, probabilities: torch.Tensor) -> OperatorEvaluation:
        if probabilities.ndim != 3 or probabilities.shape[1] != self.num_positions:
            raise ValueError("Expected probabilities shaped (batch, declared positions, symbols)")
        grouped = probabilities[:, self.groups, :]
        log_complement = torch.log1p(-grouped.clamp(max=1.0 - self.eps))
        return OperatorEvaluation(grouped, exclusivity_consequence(grouped, log_complement, -2))