import matplotlib.pyplot as plt
import numpy as np

f = 0.25
hs = range(1001)
final_amounts = [(1 + 2 * f)**h * (1 - f)**(1000 - h) for h in hs]

big_enough = [h for h, final_amount in zip(hs, final_amounts) if final_amount >= 10**9]

fig, ax = plt.subplots(figsize=(12, 8))
ax.semilogy(hs, final_amounts)
ax.axhline(1e9, color='red', linestyle='--')
ax.axvspan(min(big_enough), max(big_enough), color='green', alpha=0.2)
ax.set_xlim(0, 1000)
ax.set_ylabel('Final fortune')
ax.set_xlabel('Number of heads')
ax.set_title('{0} heads needed for f = {1}'.format(min(big_enough), f));

def min_h_greater_given_f(f, mult=2):
    ''' Returns the minimimum number of heads needed to end up with at least
        10^9 pounds after 1000 flips if your capital is multiplied by 
        (1 + mult * f) for each head and by (1 - f) for each tail. If even
        1000 heads are not enough, returns 1001.
    '''
    if f == 0:
        # Impossible to end up above 1 if wagering 0.
        h = 1001
    elif f == 1:
        # Need to never get a tail if wagering everything.
        h = 1000
    else:
        exact = (9 - 1000 * np.log10(1 - f)) / (np.log10(1 + mult * f) - np.log10(1 - f))
        h = min(1001, np.ceil(exact))
    return h

fs = np.linspace(0, 1, 1000)
hs = [min_h_greater_given_f(f) for f in fs]

min_h = min(hs)
optimal_fs = [f for f, h in zip(fs, hs) if h == min_h]

fig, (full_ax, zoom_ax) = plt.subplots(1, 2, figsize=(12, 6))

full_ax.plot(fs, hs)
full_ax.set_ylim(0, 1001)
full_ax.set_xlabel('f')
full_ax.set_ylabel('Heads needed')

zoom_ax.plot(fs, hs)
zoom_ax.set_ylim(min_h - 10, min_h + 10)
zoom_ax.set_xlim(min(optimal_fs) - 0.1, max(optimal_fs) + 0.1);
zoom_ax.set_xlabel('f')
zoom_ax.set_ylabel('Heads needed');

from scipy.misc import comb

hs = range(1002)
binomial_coefficients = [comb(1000, h) for h in hs]
pdf = np.array(binomial_coefficients) / 2.**1000
p_geq = np.cumsum(pdf[::-1])[::-1]

fig, (pdf_ax, p_geq_ax) = plt.subplots(1, 2, figsize=(12, 6))

pdf_ax.plot(hs, pdf)
pdf_ax.axis('tight')
pdf_ax.set_ylabel('Probability that H = h')
pdf_ax.set_xlabel('h')

p_geq_ax.plot(hs, p_geq)
p_geq_ax.set_ylabel('Probability that H >= h')
p_geq_ax.set_xlabel('h')
p_geq_ax.axis('tight');

from scipy.stats import binom

p_geq_alternate = [binom(1000, 0.5).sf(h - 1) for h in range(1002)]

print np.allclose(p_geq, p_geq_alternate)

print p_geq[min_h]

def best_p_given_mult_crude(mult, return_f=False):
    ''' For a given value mult of the factor, returns the probability 
        of ending up a billionaire given optimal f.
    '''
    fs = np.linspace(0, 1, 1000)
    min_h = min(min_h_greater_given_f(f, mult) for f in fs)
    best_p = p_geq[min_h]
    return best_p

fig, axs = plt.subplots(1, 2, figsize=(16, 8))
lims_list = [(1, 2), (1.48, 1.52)]

for ax, (x_min, x_max) in zip(axs, lims_list):
    mults = np.linspace(x_min, x_max, 1000)
    ps = [best_p_given_mult_crude(mult) for mult in mults]

    ax.plot(mults, ps)

    ax.axis('tight')
    ax.set_ylabel('P(billionaire) given optimal f')
    ax.set_xlabel('Factor on heads')

def perform_bisection(best_p_finder):
    low, high = 1., 2.
    num_digits = 15
    while (high - low) > 10**-num_digits:
        mid = (low + high) / 2
        mid_p = best_p_finder(mid)
        if mid_p < 0.5:
            low = mid
        else:
            high = mid
        
    print 'A factor of {0:0.{precision}f} is high enough, but {1:0.{precision}f} is not.'.format(high, low, precision=num_digits + 1)
    return high

min_factor = perform_bisection(best_p_given_mult_crude)

smaller_mult = 1.504673877089423

fig, axs = plt.subplots(1, 2, figsize=(12, 6))
num_points_list = [1000, 100000]

for ax, num_points in zip(axs, num_points_list):
    fs = np.linspace(0, 1, num_points)
    hs = [min_h_greater_given_f(f, mult=smaller_mult) for f in fs]

    min_h = min(hs)
    optimal_fs = [f for f, h in zip(fs, hs) if h == min_h]

    ax.plot(fs, hs)
    ax.axvspan(min(optimal_fs), max(optimal_fs), color='green', alpha=0.2)
    ax.set_ylim(min_h - 10, min_h + 10)
    ax.set_xlim(0.05, 0.3);
    ax.set_xlabel('f')
    ax.set_title('{0:,} points'.format(num_points))
    ax.set_ylabel('Heads needed');

from scipy.optimize import minimize_scalar

def best_p_given_mult(mult, return_f=False):
    ''' For a given value mult of the factor, returns the probability 
        of ending up a billionaire given optimal f.
    '''
    get_fractional_h = lambda f: (9 - 1000 * np.log10(1 - f)) / (np.log10(1 + mult * f) - np.log10(1 - f))
    minimization_result = minimize_scalar(get_fractional_h,
                                          bounds=(0.001, 0.999),
                                          method='Bounded',
                                          tol=1e-16,
                                         )
    fractional_h = minimization_result.fun
    min_h = min(1001, np.ceil(fractional_h))
    best_p = p_geq[min_h]
    if return_f:
        return best_p, minimization_result.x
    else:
        return best_p

even_smaller_mult = perform_bisection(best_p_given_mult)

smaller_mult - even_smaller_mult

p, optimal_f = best_p_given_mult(even_smaller_mult, return_f=True)
min_h_greater_given_f(optimal_f, even_smaller_mult)