"""
A distribution is considered to be a dictionary that maps every value to its corresponding probability, also this
module works only for distributions that take finite number of values.
"""
import numpy as np

from collections import defaultdict
from itertools import product
from pprint import pprint


def mean(dist):
    """
    Returns the mean of the given distribution
    """
    return np.sum([i * dist[i] for i in dist])


def cdf(dist, val):
    """
    Computes the CDF for the given distribution and value. i.e P(dist < val)
    """
    return np.sum([dist[i] for i in dist if i <= val])


def nfoldconv(dist, n):
    """
    Returns a distribution that is the sum of `n` independent random variables `dist`. Note that because of python's bad
    floating point precision, numpy is used to calculate the probabilities (full accuracy is still not guaranteed).

    :return: A distribution in the format that is described in the module docs.
    """
    new_dist = defaultdict(int)  # int defaults to 0
    for combination in product(dist, repeat=n):
        new_dist[sum(combination)] += np.prod(map(lambda val: np.float64(dist[val]), combination))

    print sum(new_dist.values())
    return dict(new_dist)


if __name__ == '__main__':
    # b
    fair_dice = {i: 1.0 / 6 for i in range(1, 7)}
    fair_dice_sum = nfoldconv(fair_dice, 6)
    print '\nB:'
    pprint(fair_dice_sum)
    pprint(mean(fair_dice_sum))

    # c
    price_stock_change = {-1: 0.1, 0: 0.25, 1: 0.35, 2: 0.05, 3: 0.25}
    print '\nC:'
    pprint(nfoldconv(price_stock_change, 2))
    pprint(nfoldconv(price_stock_change, 5))

    # d
    stock_after_five_days = nfoldconv(price_stock_change, 5)
    print '\nD:'
    print 1 - cdf(stock_after_five_days, 7)  # CDF computes the probability that we gained less than 7 NIS
    print cdf(stock_after_five_days, -7)