import numpy as np
import random
from fractions import Fraction


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

    Parameters
    ----------
    n : int
        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.

    Parameters
    ----------
    d : int
        Number of coordinates. In other words, we consider the simplex of dimension `d - 1`.

    Returns
    -------
    numpy.ndarray
        A numpy array of length `d`, whose sum is 1.

    Examples
    --------
        >>> 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.

    Parameters
    ----------
    generator : callable
        A callable that takes no input and that output a (random) `something`.
    test : callable
        A function ``something -> bool``.
    conditional_on : callable
        A function ``something -> bool``. Default: always True.
    n_samples : int
        Number of samples.

    Returns
    -------
    float
        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.

    Examples
    --------
        >>> 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.

    Parameters
    ----------
    x : Number

    Returns
    -------
    bool
        True if `x` is nan.

    Notes
    -----
    This extends the usual numpy function ``isnan`` to fractions.

    Examples
    --------
        >>> 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):
    """Is positive infinity.

    Parameters
    ----------
    x : Number

    Returns
    -------
    bool
        True if `x` is positive infinity.

    Notes
    -----
    This extends the usual numpy function ``isposinf`` to fractions.

    Examples
    --------
        >>> from fractions import Fraction
        >>> isposinf(Fraction(1, 10))
        False
        >>> isposinf(1)
        False
        >>> isposinf(np.inf)
        True
    """
    return np.isposinf(x)


@_false_for_fraction
def isneginf(x):
    """Is negative infinity.

    Parameters
    ----------
    x : Number

    Returns
    -------
    bool
        True if `x` is negative infinity.

    Notes
    -----
    This extends the usual numpy function ``isneginf`` to fractions.

    Examples
    --------
        >>> from fractions import Fraction
        >>> isposinf(Fraction(1, 10))
        False
        >>> isneginf(1)
        False
        >>> isneginf(- np.inf)
        True
    """
    return np.isneginf(x)


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

    Parameters
    ----------
    ballot : str
        A ballot, e.g. ``'a'``, ``'ab'``, ``'ba'``, etc.

    Returns
    -------
    str
        The same ballot in alphabetical order.

    Examples
    --------
        >>> 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.

    Parameters
    ----------
    ranking : str
        A ranking.

    Returns
    -------
    str
        The first candidate.

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


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

    Parameters
    ----------
    ranking : str
        A ranking.

    Returns
    -------
    str
        The ballot with the two first candidates.

    Examples
    --------
        >>> 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.

    Parameters
    ----------
    n : int
        The number.
    singular : str
        The singular word.
    plural : str
        The plural word. Default: the singular with a final 's'.

    Returns
    -------
    str
        The number and the word.

    Examples
    --------
        >>> 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
