## This module represents the second chapter of the book
##  "Essential Math for Data Science" - Thomas Nield
##  Chapter 2 - Probability

from scipy.stats import binom, beta
from math import factorial

# 2.0 odds means than an events has twice the probabilities to happen than not
def p_odds_to_probability(o):
	return  (o / (1 + o))

def p_probability_to_odds(p):
	return (p / (1 - p))

class BinomialDistribution:
	@staticmethod
	def example():
		n = 10
		p = 0.9

		for k in range(n + 1):
			probability = binom.pmf(k, n, p)
			print("{0} >> {1}".format(k, probability))

	@staticmethod
	def binomial_coeficient(pool, count):
		return factorial(pool) / (factorial(count) * factorial(pool - count))

	@staticmethod
	def from_scratch(p, n):
		# For each number calc the probability of that exact number of outcomes (no order)
		for k in range(n + 1):
			# 1. Simple combinatory with the binomial coeficient (combinations of k elements out of a pool of n without repetition without order)
			combinatory = BinomialDistribution.binomial_coeficient(n, k)
			# 2. Probability of success, the probability of making it k times
			probability_of_success = p ** k
			# 3. Probability of failure, inverse of the success
			probability_of_failure = (1 - p) ** (n - k)
			k_binomital_distribution_probability = combinatory * probability_of_success * probability_of_failure
			print("[{0}]: {1}".format(k, k_binomital_distribution_probability))

# Study probability of probabilities
# > Given X success rate, get the likelihood of that success rate
# > Returns a continuous function, so the final probability of X or better must be calculated using integrals (the area under the curve)
class BetaDistribution:
	@staticmethod
	def calc(probability, success_count, failure_count):
		# Only calcs the rpbability to the left
		return beta.cdf(probability, success_count, failure_count)

	@staticmethod
	def calc_right(probability, success_count, failure_count):
		return 1.0 - BetaDistribution.calc(probability, success_count, failure_count)

	@staticmethod
	def calc_region(init_probability, end_probability, success_count, failure_count):
		return BetaDistribution.calc(end_probability, success_count, failure_count) - BetaDistribution.calc(init_probability, success_count, failure_count)

p_r = 0.3
p_u = 0.4
class Exercises:
	@staticmethod
	def one():
		# 30% change of rain and 40% change your umbrella order will arrive. P(R AND U)
		return p_r * p_u

	@staticmethod
	def two():
		# Same Ps as previous. P(!R OR U)
		# RU !RU !R!U R!U
		return (1 - p_r)  + p_u - ((1 - p_r) * p_u)

	@staticmethod
	def three():
		p_u_given_rain = 0.2
		return p_u_given_rain * p_r

	@staticmethod
	def four():
		n = 137
		p_bail = 0.4
		p_at_least_50_bail = 0.0
		for x in range(50, n + 1):
			p_at_least_50_bail += binom.pmf(x, n, p_bail)
		return p_at_least_50_bail

	@staticmethod
	def five():
		n = 10
		h = 8
		t = 2
		return  1.0 - BetaDistribution.calc(0.5, 8, 2)


	@staticmethod
	def test():
		print("P(R AND U) = {0}".format(Exercises.one()))
		print("P(!R OR U) = {0}".format(Exercises.two()))
		print("P(U AND R) = {0}".format(Exercises.three()))
		print("P(At least 50 bail) = {0}".format(Exercises.four()))
		print("P(fair) = {0}".format(Exercises.five()))


def run():
	print("=== Probability module ===")
	print(">> Binomial distribution")
	BinomialDistribution.example()
	BinomialDistribution.from_scratch(0.9, 10)
	print(">> Beta distribution")
	print(BetaDistribution.calc(0.9,8,2))
	print(BetaDistribution.calc_region(0.8, 0.9, 8, 2))
	print(">> Exercises")
	Exercises.test()