"""Sparse typed relation operators for predicate-hypergraph traversal."""

from __future__ import annotations

from collections.abc import Mapping, Sequence

import torch
from torch import nn


class SparseRelationOperators(nn.Module):
    """A fixed set of sparse source-to-target adjacency operators.

    Every edge is ``(relation, source, target, weight)``.  A traversal applies
    ``A_relation.T @ state``.  Relation names choose distinct operators rather
    than modulating a shared learned message function.
    """

    def __init__(
        self,
        node_count: int,
        relation_names: Sequence[str],
        edges: Mapping[str, Sequence[tuple[int, int] | tuple[int, int, float]]],
    ) -> None:
        super().__init__()
        if node_count <= 0:
            raise ValueError("node_count must be positive")
        names = tuple(str(name) for name in relation_names)
        if not names or len(set(names)) != len(names):
            raise ValueError("relation_names must be non-empty and unique")
        relation_ids: list[int] = []
        sources: list[int] = []
        targets: list[int] = []
        weights: list[float] = []
        for relation_id, name in enumerate(names):
            for edge in edges.get(name, ()):
                if len(edge) == 2:
                    source, target = edge
                    weight = 1.0
                elif len(edge) == 3:
                    source, target, weight = edge
                else:
                    raise ValueError("edges must be (source, target) or (source, target, weight)")
                source = int(source)
                target = int(target)
                if not 0 <= source < node_count or not 0 <= target < node_count:
                    raise IndexError(f"edge ({source}, {target}) is outside {node_count} nodes")
                if float(weight) < 0:
                    raise ValueError("relation weights must be non-negative")
                relation_ids.append(relation_id)
                sources.append(source)
                targets.append(target)
                weights.append(float(weight))
        self.node_count = int(node_count)
        self.relation_names = names
        self.register_buffer("edge_relations", torch.tensor(relation_ids, dtype=torch.long))
        self.register_buffer("edge_sources", torch.tensor(sources, dtype=torch.long))
        self.register_buffer("edge_targets", torch.tensor(targets, dtype=torch.long))
        self.register_buffer("edge_weights", torch.tensor(weights, dtype=torch.float32))

    @property
    def relation_count(self) -> int:
        return len(self.relation_names)

    def relation_id(self, relation: str | int) -> int:
        if isinstance(relation, int):
            if not 0 <= relation < self.relation_count:
                raise IndexError("relation index out of range")
            return relation
        try:
            return self.relation_names.index(relation)
        except ValueError as exc:
            raise KeyError(f"unknown relation {relation!r}") from exc

    def sparse_operator(self, relation: str | int, *, transpose: bool = False) -> torch.Tensor:
        relation_id = self.relation_id(relation)
        selected = self.edge_relations == relation_id
        source = self.edge_sources[selected]
        target = self.edge_targets[selected]
        if transpose:
            source, target = target, source
        indices = torch.stack([source, target]) if source.numel() else torch.empty(
            2, 0, dtype=torch.long, device=self.edge_sources.device
        )
        values = self.edge_weights[selected]
        return torch.sparse_coo_tensor(
            indices,
            values,
            (self.node_count, self.node_count),
            device=values.device,
        ).coalesce()

    def step(
        self,
        state: torch.Tensor,
        relation: str | int,
        *,
        semiring: str = "sum_product",
    ) -> torch.Tensor:
        if state.shape[-1] != self.node_count:
            raise ValueError(f"state must end with {self.node_count} nodes")
        relation_id = self.relation_id(relation)
        selected = self.edge_relations == relation_id
        source = self.edge_sources[selected]
        target = self.edge_targets[selected]
        weight = self.edge_weights[selected].to(dtype=state.dtype)
        flat = state.reshape(-1, self.node_count)
        if semiring == "sum_product":
            out = flat.new_zeros(flat.shape)
            if source.numel():
                contribution = flat[:, source] * weight.unsqueeze(0)
                out.scatter_add_(1, target.unsqueeze(0).expand(flat.shape[0], -1), contribution)
        elif semiring == "max_product":
            out = flat.new_zeros(flat.shape)
            if source.numel():
                contribution = flat[:, source] * weight.unsqueeze(0)
                out.scatter_reduce_(
                    1,
                    target.unsqueeze(0).expand(flat.shape[0], -1),
                    contribution,
                    reduce="amax",
                    include_self=True,
                )
        elif semiring == "logsumexp":
            out = flat.new_full(flat.shape, -torch.inf)
            if source.numel():
                log_weight = weight.clamp_min(torch.finfo(weight.dtype).tiny).log()
                contribution = flat[:, source] + log_weight.unsqueeze(0)
                index = target.unsqueeze(0).expand(flat.shape[0], -1)
                maximum = flat.new_full(flat.shape, -torch.inf)
                maximum.scatter_reduce_(1, index, contribution, reduce="amax", include_self=True)
                gathered_max = maximum.gather(1, index)
                shifted = torch.where(
                    torch.isfinite(gathered_max),
                    (contribution - gathered_max).exp(),
                    torch.zeros_like(contribution),
                )
                total = flat.new_zeros(flat.shape)
                total.scatter_add_(1, index, shifted)
                finite = torch.isfinite(maximum) & (total > 0)
                out = torch.where(finite, maximum + total.clamp_min(torch.finfo(total.dtype).tiny).log(), out)
        else:
            raise ValueError("semiring must be sum_product, max_product, or logsumexp")
        return out.reshape(state.shape)


__all__ = ["SparseRelationOperators"]