from dstz.core.atom import Element
from dstz.core.distribution import Evidence
from dstz.element.combination import simple_space, powerset
from dstz.element.permutation import permutation_set


def max_entropy_distribution(n, simple_generator, event_generator,
                             component_generator=set, condition_func=None):
    """
    Computes the maximum entropy distribution over events generated from a simple space.

    Args:
        - n (int): Size parameter influencing the generation of the simple space.
        - simple_generator (function): Function to generate the simple space from `n`.
        - event_generator (function): Transforms the simple space into a set of events.
        - component_generator (function, optional): Method to create subsets of events. Defaults to `set`.
        - condition_func (function, optional): A conditional function applied to event generations.

    Returns:
        - Evidence: An Evidence instance representing the calculated maximum entropy distribution.

    Description:
        This function calculates a distribution where the entropy is maximized given a set of events derived from a simple space.
        It iterates through each event, calculating its weight based on the number of outcomes from applying `condition_func`
        (or directly the event generation if no condition function is provided) to subsets generated by `component_generator`.
        After determining weights, it normalizes these values to ensure the distribution sums to unity.

    Example Usage:
        >>> dist = max_entropy_distribution(3, simple_space, powerset)
    """

    res = Evidence()
    ss = simple_generator(n)
    es = event_generator(ss)

    # Compute weights for each event in the distribution
    for element in es:
        count = len(condition_func(event_generator(component_generator(element)))) \
            if condition_func else len(event_generator(component_generator(element)))
        res[Element(element)] = float(count)

    # Normalize the distribution
    states = sum(res.values())
    for element in es:
        res[Element(element)] /= states

    return res


def max_deng_entropy_distribution(n):
    """
    Calculates the maximum Deng entropy distribution for a given size `n`.

    Args:
        - n (int): Determines the size of the simple space used for generating the distribution.

    Returns:
        - Evidence: The maximum Deng entropy distribution computed using `simple_space` and `powerset`.

    Description:
        Utilizes the `max_entropy_distribution` function with specific generators (`simple_space` and `powerset`)
        to derive the Deng entropy distribution, a particular case of maximum entropy distributions in the context of
        evidence theory.

    Example Usage:
        >>> deng_dist = max_deng_entropy_distribution(4)
    """
    return max_entropy_distribution(n, simple_space, powerset)


def max_rps_entropy_distribution(n):
    """
    Computes the maximum entropy distribution based on permutations of a simple space for a given size `n`.

    Args:
        - n (int): Size of the simple space, which influences the generation of permutations.

    Returns:
        - Evidence: An Evidence instance reflecting the maximum entropy distribution derived from all possible permutations of the simple space elements.

    Description:
        This function leverages the principle of maximum entropy to determine a probability distribution across all permutations of a simple space of size `n`.
        It employs the `max_entropy_distribution` function, specifying `simple_space` as the generator for the base space and `permutation_set` to create the set of events
        as all possible permutations of the simple space elements. The resulting distribution assigns probabilities such that the entropy is maximized.

    Example Usage:
        >>> rps_dist = max_rps_entropy_distribution(3)
    """
    return max_entropy_distribution(n, simple_space, permutation_set)
