# -*- coding: utf-8 -*-

"""
Functions for pruning or expanding the sample space of a distribution.

This can be important when calculating meet and join random variables. It
is also important for the calculations of various PID quantities.

"""

def pruned_samplespace(d, sample_space=None):
    """
    Returns a new distribution with pruned sample space.

    The pruning is such that zero probability outcomes are removed.

    Parameters
    ----------
    d : distribution
        The distribution used to create the pruned distribution.

    sample_space : set
        A list of outcomes with zero probability that should be kept in the
        sample space. If `None`, then all outcomes with zero probability
        will be removed.

    Returns
    -------
    pd : distribution
        The distribution with a pruned sample space.

    """
    if sample_space is None:
        sample_space = []

    keep = set(sample_space)
    outcomes = []
    pmf = []
    for o, p in d.zipped(mode='atoms'):
        if not d.ops.is_null_exact(p) or o in keep:
            outcomes.append(o)
            pmf.append(p)

    pd = d.__class__(outcomes, pmf, base=d.get_base())
    return pd

def expanded_samplespace(d, alphabets=None):
    """
    Returns a new distribution with an expanded sample space.

    Expand the sample space so that it is the Cartesian product of the
    alphabets for each random variable. Note, only the effective alphabet of
    each random variable is used. So if one index in an outcome only has the
    value 1, then its alphabet is [1], and not [0,1] for example.

    Parameters
    ----------
    d : distribution
        The distribution used to create the pruned distribution.

    alphabets : list
        A list of alphabets, with length equal to the outcome length in `d`.
        Each alphabet specifies the alphabet to be used for a single index
        random variable. The sample space of the new distribution will be the
        Cartesian product of these alphabets.

    Returns
    -------
    ed : distribution
        The distribution with an expanded sample space.

    """
    if alphabets is None:
        # Note, we sort the alphabets now, so we are possibly changing the
        # order of the original sample space.
        alphabets = map(sorted, d.alphabet)
    else:
        L = d.outcome_length()
        if len(alphabets) != L:
            raise Exception("You need to provide {0} alphabets".format(L))


    outcomes = list(d._product(*alphabets))
    pmf = [0] * len(outcomes)
    for o, p in d.zipped(mode='patoms'):
        idx = outcomes.index(o)
        pmf[idx] = p

    ed = d.__class__(outcomes, pmf, base=d.get_base())
    return ed

# -*- coding: utf-8 -*-

"""
Functions for pruning or expanding the sample space of a distribution.

This can be important when calculating meet and join random variables. It
is also important for the calculations of various PID quantities.

"""

def pruned_samplespace(d, sample_space=None):
    """
    Returns a new distribution with pruned sample space.

    The pruning is such that zero probability outcomes are removed.

    Parameters
    ----------
    d : distribution
        The distribution used to create the pruned distribution.

    sample_space : set
        A list of outcomes with zero probability that should be kept in the
        sample space. If `None`, then all outcomes with zero probability
        will be removed.

    Returns
    -------
    pd : distribution
        The distribution with a pruned sample space.

    """
    if sample_space is None:
        sample_space = []

    keep = set(sample_space)
    outcomes = []
    pmf = []
    for o, p in d.zipped(mode='atoms'):
        if not d.ops.is_null_exact(p) or o in keep:
            outcomes.append(o)
            pmf.append(p)

    pd = d.__class__(outcomes, pmf, base=d.get_base())
    return pd

def expanded_samplespace(d, alphabets=None):
    """
    Returns a new distribution with an expanded sample space.

    Expand the sample space so that it is the Cartesian product of the
    alphabets for each random variable. Note, only the effective alphabet of
    each random variable is used. So if one index in an outcome only has the
    value 1, then its alphabet is [1], and not [0,1] for example.

    Parameters
    ----------
    d : distribution
        The distribution used to create the pruned distribution.

    alphabets : list
        A list of alphabets, with length equal to the outcome length in `d`.
        Each alphabet specifies the alphabet to be used for a single index
        random variable. The sample space of the new distribution will be the
        Cartesian product of these alphabets.

    Returns
    -------
    ed : distribution
        The distribution with an expanded sample space.

    """
    if alphabets is None:
        # Note, we sort the alphabets now, so we are possibly changing the
        # order of the original sample space.
        alphabets = map(sorted, d.alphabet)
    else:
        L = d.outcome_length()
        if len(alphabets) != L:
            raise Exception("You need to provide {0} alphabets".format(L))


    outcomes = list(d._product(*alphabets))
    pmf = [d.ops.zero] * len(outcomes)
    for o, p in d.zipped(mode='patoms'):
        idx = outcomes.index(o)
        pmf[idx] = p

    ed = d.__class__(outcomes, pmf, base=d.get_base())
    return ed

Make prune_expand work with log distributions.