# AUTOGENERATED! DO NOT EDIT! File to edit: ../discrete.ipynb.

# %% auto 0
__all__ = ['DiscreteKey', 'DiscreteVariable', 'Variables']

# %% ../discrete.ipynb 3
from dataclasses import dataclass
from typing import Callable, Iterable
import gtsam

# %% ../discrete.ipynb 11
DiscreteKey = tuple[int, int]


@dataclass
class DiscreteVariable:
    name: str
    domain: list[str]


class Variables:
    """A problem domain for discrete variables."""

    BINARY = ["false", "true"]

    def __init__(self):
        """Default constructor"""
        self._variables: dict[DiscreteVariable] = {}

    def size(self):
        """Return number of variables defined."""
        return len(self._variables)

    def discrete(self, name: str, domain: list[str]) -> DiscreteKey:
        """Create a variable with given name and discrete domain of named values.

        Args:
            name (str): name of the variable.
            domain (list[str]): names for the different values.

        Returns:
            DiscreteKey, i.e., (gtsam.Key, cardinality)
        """
        discreteKey = len(self._variables), len(domain)
        self._variables[discreteKey[0]] = DiscreteVariable(name, domain)
        return discreteKey

    def discrete_series(
        self: "Variables",
        character: str,
        indices: Iterable[int],
        domain: list[str],
    ) -> dict[int, DiscreteKey]:
        """Create several discrete variables with Symbol names.

        Args:
            character (str): a single character.
            indices: (Iterable[int]): a set of integer indices.
            domain (list[str]): names for the different values.

        Returns:
            dict[int, DiscreteKey], i.e., [(gtsam.Key, cardinality)]
        """
        assert len(character) == 1, "discrete_series: requires character only"
        n = len(domain)

        def make_key(index):
            symbol = gtsam.Symbol(character, index)
            key = symbol.key()
            name = symbol.string()
            self._variables[key] = DiscreteVariable(name, domain)
            return (key, n)

        return {index: make_key(index) for index in indices}

    def binary(self: "Variables", name: str) -> DiscreteKey:
        """Create a binary variable with given name.

        Args:
            name (str): name of the variable.

        Returns:
            DiscreteKey, i.e., (gtsam.Key, cardinality)
        """
        return self.discrete(name, self.BINARY)

    def name(self, discreteKey: DiscreteKey) -> str:
        """Return name of variable with given discreteKey.

        Args:
            discreteKey (DiscreteKey): (gtsam.Key, cardinality)

        Returns:
            str: name of the variable.
        """
        return self._variables[discreteKey[0]].name

    def domain(self, discreteKey: DiscreteKey):
        """Return domain of variable with given discreteKey.

        Args:
            discreteKey (DiscreteKey): (gtsam.Key, cardinality)

        Returns:
            str: domain of the variable.
        """
        return self._variables[discreteKey[0]].domain

    def assignment(
        self: "Variables", map: dict[DiscreteKey, str]
    ) -> gtsam.DiscreteValues:
        """Create a GTSAM assignment of keys to values.

        Args:
            map (dict[DiscreteKey, str]): map from discrete keys to values.

        Returns:
            gtsam.DiscreteValues: the GTSAM equivalent.
        """
        values = gtsam.DiscreteValues()
        for discreteKey, value in map.items():
            domain = self.domain(discreteKey)
            assert (
                value in domain
            ), f"Specified value '{value}' not found in domain of '{self.name(discreteKey)}'. Allowed values are {domain}."
            values[discreteKey[0]] = domain.index(value)
        return values

    def keyFormatter(self: "Variables") -> Callable:
        """Return a lambda that can be used as KeyFormatter in GTSAM"""
        return lambda key: self._variables[key].name

    def names(self: "Variables") -> dict[int, list[str]]:
        """Return a names dictionary that is used by GTSAM markdown methods"""
        return {key: var.domain for (key, var) in self._variables.items()}

    def values_markdown(
        self: "Variables", assignment: gtsam.DiscreteValues
    ) -> str:
        """Render a DiscreteValues instance as markdown.

        Args:
                assignment (gtsam.DiscreteValues): the values to render.

        Returns:
                str: a markdown string.
        """
        return gtsam.markdown(assignment, self.keyFormatter(), self.names())

    def values_html(self: "Variables", assignment: gtsam.DiscreteValues) -> str:
        """Render a DiscreteValues instance as html.

        Args:
                assignment (gtsam.DiscreteValues): the values to render.

        Returns:
                str: a html string.
        """
        return gtsam.html(assignment, self.keyFormatter(), self.names())

    def _repr_html_(self: "Variables"):
        """An html representation"""
        # Print out preamble.
        ss = "<div>\n<table class='Variables'>\n  <thead>\n"

        # Print out header row.
        ss += "    <tr><th>Variable</th><th>Domain</th></tr>\n"

        # Finish header and start body.
        ss += "  </thead>\n  <tbody>\n"

        # Print out all rows.
        for variable in self._variables.values():
            ss += "    <tr>"
            ss += "<th>" + variable.name + "</th><td>"
            for value in variable.domain[:-1]:
                ss += value + ", "
            ss += variable.domain[-1]
            ss += "</td>"
            ss += "</tr>\n"
        ss += "  </tbody>\n</table>\n</div>"
        return ss

