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

__all__ = ['P', 'DiscretePrior', 'V']

# Cell
import gtsam


# Cell
def P(*args):
    """ Create a DiscreteKeys instances from a variable number of DiscreteKey pairs.

    It is currently very annoying to create a set of discrete keys.
    This is probably because we don't understand pybind11 fully yet.

    Discrete keys are represented as pairs in Python, with the first element the key and
    the second element the cardinality of the discrete variable. This function accepts a
    variable number of such pairs and creates a gtsam.DiscreteKeys instance.


    Returns:
        DiscreteKeys: an instance of DiscreteKeys.
    """
    #TODO: We can make life easier by providing variable argument functions in C++ itself.
    dks = gtsam.DiscreteKeys()
    for key in args:
        dks.push_back(key)
    return dks

# Cell
def DiscretePrior(key: int, pmf: str):
    """ A prior on a discrete variable.

    This really just creates a gtsam.DiscreteConditional without parents.
    """
    # TODO: add to GTSAM
    return gtsam.DiscreteConditional(key, P(), pmf)


# Cell
def V(values):
    """Convert a dictionary into a DiscreteValues instance.

    Args:
        values (dict): A map from keys to values.

    Returns:
        gtsam.DiscreteValues: an assignment to discrete variables usable by gtsam.
    """
    # TODO: add variable argument operator that takes values in keys order.
    # TODO: allow conversion from dictionary to DiscreteValues.
    a = gtsam.DiscreteValues()
    for key, value in values.items():
        a[key] = value
    return a
