import numpy as np

n_teams = 10
n_persons_per_team = 3
n_beans = 10
n_rounds = 10
n_die_change = 7
target_rain = 10
penalty = 4

def initialize():
    beans = n_beans * np.ones((n_teams, n_persons_per_team))
    forecast_teams = np.ones((n_teams)) # receive regional forecast
    drr_teams = np.ones((n_teams))  # have disaster risk reduction
    return beans, forecast_teams, drr_teams

def get_forecast_bids(beans):
    """ Defines how each person or team will bid for regional forecast

    Example
    -------

    return np.random.randint(0, np.max(beans) * .4, size=beans.shape)
    """
    return np.ones(beans.shape)

def get_drr_bids(beans):
    """ Defines how each person or team will bid for disaster risk reduction

    Example
    -------

    bids = np.zeros(beans.shape)
    for i in range(beans.shape[0]):
        for j in range(beans.shape[1]):
            bids[i, j] = np.random.randint(0, beans[i, j] * .2)
    """
    bids = np.ones(beans.shape)
    return bids


beans, forecast_teams, drr_teams = initialize()

# perform forecast bids
forecast_bids = get_forecast_bids(beans)
forecast_team_bids = np.sum(forecast_bids, axis=1)
sort_index = np.argsort(forecast_team_bids)
forecast_teams[sort_index[:n_teams/2]] = 0

# Winning teams pay their beans
beans = beans - (forecast_teams[:, None] * forecast_bids)
print beans.T
print forecast_teams

# perform drr bids
drr_bids = get_drr_bids(beans)
drr_team_bids = np.sum(drr_bids, axis=1)
sort_index = np.argsort(drr_team_bids)
drr_teams[sort_index[:-2]] = 0

# Winning teams pay their beans
beans = beans - (drr_teams[:, None] * drr_bids)
print beans.T
print forecast_teams
print drr_teams

def get_insurance_payments(regional_predictions, drr_teams, beans, round_idx):
    # determine the likelihood of a flood
    likelihood = (7 - (target_rain - regional_predictions)) / 6
    # if likelihood > .2 pay, or if you didn't have a prediction pay one bean
    payments = ((likelihood > .2) + (regional_predictions < 1))[:, None] * np.ones(beans.shape)
    return (payments * (beans > 0)).astype(int)

def generate_rainfall(n_sides):
    regional_rainfall = np.random.randint(1, n_sides, size=(n_teams))
    local_rainfall = np.random.randint(1, 7, size=(n_teams, n_persons_per_team))
    total_rainfall = local_rainfall + regional_rainfall[:, None]
    flooded = (total_rainfall >= target_rain).astype(np.int)
    return regional_rainfall, flooded

def adjust_beans(beans, payments, flooded, round_idx, drr_teams):
    if round_idx < 3:
        drr_penalty = 4
    else:
        drr_penalty = 2
    penalized = np.maximum(flooded - payments, 0)
    beans_to_remove = drr_penalty * penalized * (drr_teams[:, None] == 1) + penalty * penalized * (drr_teams[:, None] == 0)
    already_in_crisis = beans * (beans < 0)
    beans[already_in_crisis < 0] = 0
    beans = beans - payments - beans_to_remove
    beans_joining_crisis = beans < 0 
    beans = beans * (beans > 0)
    in_crisis = already_in_crisis - beans_joining_crisis
    return beans + in_crisis

print beans.T
for turn in range(n_rounds):
    n_sides = 6
    if turn == 6: # 7th round
        n_sides = 8
    regional_rainfall, flooded = generate_rainfall(n_sides=n_sides)
    payments = get_insurance_payments(regional_rainfall * forecast_teams, drr_teams, beans, turn + 1)
    beans = adjust_beans(beans.copy(), payments, flooded, turn + 1, drr_teams)
    if turn % 2 == 0:
        print turn + 1, flooded.T - payments.T
print beans.T

!/software/challenges/nbconvert/nbconvert.py -f reveal /software/challenges/payforpredictions/payforpredictions.ipynb