from matplotlib.pyplot import gcf, legend, plot, savefig, show, title, xlabel, ylabel
import numpy
from numpy import arange, array, cumsum, zeros_like
from numpy.random import random

BLOCK_REWARD = 1

def difficulties(diff_increase, epochs, length):
    exps = numpy.floor(numpy.linspace(0, epochs, epochs * length, endpoint=False))
    return diff_increase ** exps

def payouts(diffs, hashrate, pool_hashrate):
    pool_rewards = random(len(diffs)) < pool_hashrate / diffs
    return hashrate / pool_hashrate * pool_rewards

def accumulate(times, hashrate, pool_hashrate, diff_increase, epochs, epoch_length):
    diffs = difficulties(diff_increase, epochs, epoch_length)
    pays = zeros_like(diffs)
    for _ in xrange(times):
        pays += payouts(diffs, hashrate, pool_hashrate)
    return pays / times

def plot_cum_payouts(times, hashrate, pool_hashrate, diff_increase, epochs, epoch_length, linetype, label, **plot_kwargs):
    pays = accumulate(times, hashrate, pool_hashrate, diff_increase, epochs, epoch_length)
    cum_pays = cumsum(pays)
    ticks = arange(len(pays))
    plot(ticks, cum_pays, linetype, label=label, **plot_kwargs)

def plot_times_cum_payouts(times, hashrate, pool_hashrate, diff_increase, epochs, epoch_length, linetype, label):
    for i in xrange(times):
        label_i = label if i == 0 else None
        plot_cum_payouts(1, hashrate, pool_hashrate, diff_increase, epochs, epoch_length, linetype, label_i)

def plot_cloud():
    times = 100
    inverse_hashrate = 1000
    red_pool_percent = 50
    blue_pool_percent = 20
    difficulty_inc_percent = 15
    epochs = 10
    epoch_length = 100
    plot_times_cum_payouts(times, 1.0/inverse_hashrate, red_pool_percent/100.0, (1 + difficulty_inc_percent/100.0), epochs, epoch_length, 'r-', "%d%% pool" % red_pool_percent)
    plot_times_cum_payouts(times, 1.0/inverse_hashrate, blue_pool_percent/100.0, (1 + difficulty_inc_percent/100.0), epochs, epoch_length, 'b-', "%d%% pool" % blue_pool_percent)
    legend(loc='upper left')
    title("{times} earnings simulations: {difficulty_inc_percent}% difficulty increase every {epoch_length} blocks, {block_reward} btc block reward, miner hash rate = 1/{inverse_hashrate} initial network hash rate".format(
        times=times,
        difficulty_inc_percent=difficulty_inc_percent,
        epoch_length=epoch_length,
        block_reward=BLOCK_REWARD,
        inverse_hashrate=inverse_hashrate))
    xlabel('block')
    ylabel('cumulative earnings')
    gcf().set_size_inches(16,12)
    savefig('simulated_pool_variance.png', dpi=50)
    show()

def plot_average():
    times = 10000
    inverse_hashrate = 1000
    red_pool_percent = 50
    blue_pool_percent = 20
    difficulty_inc_percent = 15
    epochs = 10
    epoch_length = 100
    plot_cum_payouts(times, 1.0/inverse_hashrate, red_pool_percent/100.0, (1 + difficulty_inc_percent/100.0), epochs, epoch_length, 'r-', "%d%% pool" % red_pool_percent, linewidth=10)
    plot_cum_payouts(times, 1.0/inverse_hashrate, blue_pool_percent/100.0, (1 + difficulty_inc_percent/100.0), epochs, epoch_length, 'b--', "%d%% pool" % blue_pool_percent, linewidth=10)
    legend(loc='upper left')
    title("average earnings from {times} simulations: {difficulty_inc_percent}% difficulty increase every {epoch_length} blocks, {block_reward} btc block reward, miner hash rate = 1/{inverse_hashrate} initial network hash rate".format(
        times=times,
        difficulty_inc_percent=difficulty_inc_percent,
        epoch_length=epoch_length,
        block_reward=BLOCK_REWARD,
        inverse_hashrate=inverse_hashrate))
    xlabel('block')
    ylabel('cumulative earnings')
    gcf().set_size_inches(16,12)
    savefig('simulated_pool_average.png', dpi=50)
    show()

plot_cloud()
plot_average()