"""Assess a betting strategy.

Copyright 2018, Georgia Institute of Technology (Georgia Tech)
Atlanta, Georgia 30332
All Rights Reserved

Template code for CS 4646/7646

Georgia Tech asserts copyright ownership of this template and all derivative
works, including solutions to the projects assigned in this course. Students
and other users of this template code are advised not to share it with others
or to make it available on publicly viewable websites including repositories
such as github and gitlab.  This copyright statement should not be removed
or edited.

We do grant permission to share solutions privately with non-students such
as potential employers. However, sharing with other current or future
students of CS 7646 is prohibited and subject to being investigated as a
GT honor code violation.

-----do not edit anything above this line---

Student Name: Tucker Balch (replace with your name)
GT User ID: tb34 (replace with your User ID)
GT ID: 900897987 (replace with your GT ID)
"""

import numpy as np
import matplotlib.pyplot as plt

def author():
    return 'tb34' # replace tb34 with your Georgia Tech username.

def gtid():
    return 900897987 # replace with your GT ID number

def get_spin_result(win_prob):
    result = False
    if np.random.random() <= win_prob:
        result = True
    return result

def run_experiment(max_winnings=80, max_bets=1000, initial_credit=0):
    win_prob = 18 / 38 # 18 black numbers out of 38 total numbers
    current_bet, episode_winnings, bet_amount = 0, 0, 1
    winnings = np.zeros(max_bets)

    # Keep making bets until we reach desired winnings or betting limit
    while episode_winnings < max_winnings and current_bet < max_bets:
        if get_spin_result(win_prob):
            episode_winnings += bet_amount
            bet_amount = 1
        else:
            episode_winnings -= bet_amount
            bet_amount *= 2
        winnings[current_bet] = episode_winnings
        current_bet += 1

        # Handle experiment 2 where we have a initial maximum bankroll
        if initial_credit > 0:
            current_credit = initial_credit + episode_winnings
            if current_credit <= 0:
                break
            if bet_amount > current_credit:
                bet_amount = current_credit

    # Fill remaining fields with last value
    while current_bet < max_bets:
        winnings[current_bet] = episode_winnings
        current_bet += 1

    return winnings

def configure_plot():
    plt.figure()
    axes = plt.gca()
    axes.set_xlim([0, 300])
    axes.set_ylim([-256, 100])
    plt.xlabel("#bets []")
    plt.ylabel("win [$]")

def experiment_1_figure_1(number_runs=10):
    configure_plot()
    for _ in range(number_runs):
        winnings = run_experiment()
        plt.plot(winnings)
    plt.savefig('figure_1.png')

def experiment_1_figure_2(number_runs=1000):
    configure_plot()
    runs = np.array([run_experiment() for _ in range(number_runs)])
    winnings_mean = runs.mean(axis=0)
    winnings_std = runs.std(axis=0)

    plt.plot(winnings_mean, linewidth=0.7)
    plt_std_setting = {'ls': '-', 'color': 'blue', 'linewidth': 0.3}
    plt.plot(winnings_mean + winnings_std, **plt_std_setting)
    plt.plot(winnings_mean - winnings_std, **plt_std_setting)
    plt.savefig('figure_2.png')

    experiment_1_figure_3(runs)

def experiment_1_figure_3(runs, figurename='figure_3.png'):
    configure_plot()
    winnings_median = np.median(a=runs, axis=0)
    winnings_std = runs.std(axis=0)
    plt.plot(winnings_median, linewidth=0.7)
    plt_std_setting = {'ls': '-', 'color': 'blue', 'linewidth': 0.3}
    plt.plot(winnings_median + winnings_std, **plt_std_setting)
    plt.plot(winnings_median - winnings_std, **plt_std_setting)
    plt.savefig(figurename)

def experiment_2_figure_4(number_runs=1000):
    configure_plot()
    runs = np.array([run_experiment(initial_credit=256)
                    for _ in range(number_runs)])
    winnings_mean = runs.mean(axis=0)
    winnings_std = runs.std(axis=0)
    plt.plot(winnings_mean, linewidth=0.7)
    plt_std_setting = {'ls': '-', 'color': 'blue', 'linewidth': 0.3}
    plt.plot(winnings_mean + winnings_std, **plt_std_setting)
    plt.plot(winnings_mean - winnings_std, **plt_std_setting)
    plt.savefig('figure_4.png')
    experiment_1_figure_3(runs, figurename='figure_5.png')

def test_code():
    np.random.seed(gtid()) # do this only once
    experiment_1_figure_1()
    experiment_1_figure_2()
    experiment_2_figure_4()

if __name__ == "__main__":
    test_code()