# AUTOGENERATED! DO NOT EDIT! File to edit: discrete.ipynb (unless otherwise specified).

__all__ = ['Variables', 'DiscreteKey']

# Cell
import gtsam

from typing import List, Tuple, Callable, Dict, Iterable


# Cell

DiscreteKey = Tuple[int, int]


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

    BINARY = ["false", "true"]

    def __init__(self):
        """Default constructor"""
        self._variables = dict()

    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]] = name, domain
        return discreteKey

    def discrete_series(self, 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] = name, domain
            return (key, n)

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

    def binary(self, 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]][0]

    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]][1]

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

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

    def assignment(self, 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 values_markdown(self, 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, 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):
        """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 name, domain in self._variables.values():
            ss += "    <tr>"
            ss += "<th>" + name + "</th><td>"
            for value in domain[:-1]:
                ss += value + ", "
            ss += domain[-1]
            ss += "</td>"
            ss += "</tr>\n"
        ss += "  </tbody>\n</table>\n</div>"
        return ss
