"""Typed relation-sequence composition utilities used by path reasoning."""

from __future__ import annotations

from collections.abc import Mapping, Sequence

import torch
from torch import Tensor


def compose_relation_sequence(
    relation_sequence: Sequence[str],
    relation_states: Mapping[str, Tensor],
    query_state: Tensor,
) -> Tensor:
    """Compose relation states with position-aware gated products.

    The operation is deliberately inductive: relation IDs select supplied
    graph-derived states, not a closed-world entity table. PAT remains the
    learned path encoder; this utility gives the reasoning core an explicit
    compositional state for supervision and diagnostics.
    """
    if query_state.ndim == 2:
        query_state = query_state[0]
    if not relation_sequence:
        return torch.zeros_like(query_state)
    states = [
        relation_states.get(relation, torch.zeros_like(query_state))
        for relation in relation_sequence
    ]
    result = states[0]
    for index, state in enumerate(states[1:], start=1):
        gate = 1.0 / float(index + 1)
        result = torch.tanh(
            result * (1.0 - gate) + state * gate + query_state * (gate / 2.0)
        )
    return result


def relation_sequence_score(
    relation_sequence: Sequence[str],
    relation_states: Mapping[str, Tensor],
    query_state: Tensor,
) -> Tensor:
    """Return query compatibility for a composed relation sequence."""
    composed = compose_relation_sequence(
        relation_sequence, relation_states, query_state
    )
    if query_state.ndim == 2:
        query_state = query_state[0]
    return torch.nn.functional.cosine_similarity(
        composed.unsqueeze(0), query_state.unsqueeze(0), dim=-1
    ).squeeze(0)