"""Short typed graph programs with exact role unbinding."""

from __future__ import annotations

from dataclasses import dataclass

import torch

from strata.modeling.algebra.relations import SparseRelationOperators
from strata.modeling.algebra.role_binding import OrthonormalRoleBinder


@dataclass(frozen=True, slots=True)
class GraphProgram:
    relations: tuple[str, ...] = ()
    read_role: str | None = None

    def __post_init__(self) -> None:
        if len(self.relations) > 3:
            raise ValueError("qualification graph programs are bounded to three hops")


@dataclass(frozen=True, slots=True)
class GraphProgramResult:
    node_state: torch.Tensor
    filler: torch.Tensor | None
    trace: tuple[torch.Tensor, ...]


class GraphProgramExecutor:
    def __init__(
        self,
        relations: SparseRelationOperators,
        role_binder: OrthonormalRoleBinder,
    ) -> None:
        self.relations = relations
        self.role_binder = role_binder

    def execute(
        self,
        anchor_state: torch.Tensor,
        program: GraphProgram,
        *,
        event_memories: torch.Tensor | None = None,
        semiring: str = "max_product",
    ) -> GraphProgramResult:
        state = anchor_state
        trace = [state]
        for relation in program.relations:
            state = self.relations.step(state, relation, semiring=semiring)
            trace.append(state)
        filler = None
        if program.read_role is not None:
            if event_memories is None:
                raise ValueError("event_memories are required for a role read")
            expected = (
                *state.shape[:-1],
                self.relations.node_count,
                self.role_binder.filler_dim,
                self.role_binder.role_dim,
            )
            if event_memories.shape != expected:
                raise ValueError(f"event_memories must have shape {expected}, got {tuple(event_memories.shape)}")
            memory = torch.einsum("...n,...nfd->...fd", state, event_memories)
            filler = self.role_binder.unbind(memory, program.read_role)
        return GraphProgramResult(node_state=state, filler=filler, trace=tuple(trace))


__all__ = ["GraphProgram", "GraphProgramExecutor", "GraphProgramResult"]