import numpy as np
import random
from fractions import Fraction


def initialize_random_seeds(n=0):
    """Initialize the random seeds.

    :param n: the desired random seed (default: 0).
    """
    random.seed(n)
    np.random.seed(n)


def rand_simplex(d=6):
    """
    Draw a random point in the simplex.

    :param d: number of coordinate (i.e. we consider the simplex of dimension d - 1).
    :return: a numpy array of length d, whose sum is 1.

    >>> initialize_random_seeds()
    >>> rand_simplex(d=6)  # doctest: +SKIP
    array([0.4236548 , 0.12122838, 0.00393032, 0.05394987, 0.11242599, 0.28481063])
    """
    x = np.sort(np.random.rand(d - 1))
    return np.concatenate((x, [1])) - np.concatenate(([0], x))


def probability(generator, n_samples, test, conditional_on=None):
    """Probability that a random 'something' meets some given test.

    :param generator: a callable that takes no input and that output a (random) something.
    :param test: a function ``something -> bool``.
    :param conditional_on: a function ``something -> bool``. Default: always True.
    :param n_samples: number of samples.
    :return: the probability that a 'something' generated by ``generator`` meets ``test``, conditional on the fact
        that it meets ``conditional_on``, based on a Monte-Carlo estimation of ``n_samples`` trials.

        >>> import random
        >>> initialize_random_seeds()
        >>> def generator():
        ...     return random.random()
        >>> probability(generator, n_samples=1000, test=lambda x: x > .5, conditional_on=lambda x: x > .25)
        0.661
    """
    success = 0
    i_samples = 0
    while i_samples < n_samples:
        something = generator()
        if conditional_on is None or conditional_on(something):
            i_samples += 1
            if test(something):
                success += 1
    return success / n_samples


def _false_for_fraction(f):
    """Decorator to return False when the input is a Fraction (cf. usages below)."""
    def _f(x):
        if isinstance(x, Fraction):
            return False
        return f(x)
    _f.__doc__ = f.__doc__
    return _f


@_false_for_fraction
def isnan(x):
    """Is nan.

        >>> from fractions import Fraction
        >>> isnan(Fraction(1, 10))
        False
        >>> isnan(1)
        False
        >>> isnan(np.nan)
        True
    """
    return np.isnan(x)


@_false_for_fraction
def isposinf(x):
    return np.isposinf(x)


@_false_for_fraction
def isneginf(x):
    return np.isneginf(x)


def sort_ballot(ballot):
    """Put a ballot in alphabetical order.

    :param ballot: a ballot, e.g. 'a', 'ab', 'ba'...
    :return: the same in alphabetical order.

    >>> sort_ballot('a')
    'a'
    >>> sort_ballot('ab')
    'ab'
    >>> sort_ballot('ba')
    'ab'
    """
    return ''.join(sorted(ballot))


def ballot_one(ranking):
    """Ballot for the voter's preferred candidate.

    :param ranking: a ranking.
    :return: the first candidate.

    >>> ballot_one('abc')
    'a'
    """
    return ranking[0]


def ballot_one_two(ranking):
    """Ballot for the voter's two preferred candidates.

    :param ranking: a ranking.
    :return: the ballot.

    >>> ballot_one_two('abc')
    'ab'
    >>> ballot_one_two('bac')
    'ab'
    """
    return min(ranking[:2], ranking[1::-1])


def give_figure(n, singular, plural=None):
    """Combine a number with a unit, whose word can be singular or plural.

    :param n: the number.
    :param singular: the singular word.
    :param plural: the plural word. Default: the singular with an 's'.
    :return: the number and the word.

    >>> give_figure(1, 'apple')
    '1 apple'
    >>> give_figure(2, 'apple')
    '2 apples'
    >>> give_figure(1, 'equilibrium', 'equilibria')
    '1 equilibrium'
    >>> give_figure(2, 'equilibrium', 'equilibria')
    '2 equilibria'
    """
    if n <= 1:
        return str(n) + ' ' + singular
    else:
        if plural is None:
            return str(n) + ' ' + singular + 's'
        else:
            return str(n) + ' ' + plural
