import numpy as np

import pyqtgraph as pg
from pyqtgraph.Qt import QtCore


initial_dollar = 100

factor_if_heads = 1.8
factor_if_tails = 0.5

population_size = 10000
num_tosses_visualized = 500

def want_to_play(money, toss) -> bool:
    return True
    if toss > 4:
        return False

rng = np.random.default_rng(seed=4)

def evolve(want_to_play, money, toss):
    # if money < 2:
    #     return 1
    # if not want_to_play(money, toss):
    #     return money

    if rng.choice([0,1]):
        return money * factor_if_heads
    else:
        return money * factor_if_tails

population_histories = [[initial_dollar] for _ in range(population_size)]
for toss in range(num_tosses_visualized):
    for hist in population_histories:
        hist.append(evolve(want_to_play, hist[-1], toss))


app = pg.mkQApp("Plotting Example")
#mw = QtWidgets.QMainWindow()
#mw.resize(800,800)

win = pg.GraphicsLayoutWidget(show=True, title="Basic plotting examples")
win.resize(1000,600)
win.setWindowTitle('pyqtgraph example: Plotting')

# Enable antialiasing for prettier plots
pg.setConfigOptions(antialias=True)

p2 = win.addPlot(title="Multiple curves")


population_histories = np.array(population_histories)
# population_histories = np.log(population_histories + 1)
p2.setLogMode(False, True)
p2.plot(population_histories.mean(0), pen=(255, 0, 0))
p2.plot(np.median(population_histories, 0), pen=(0, 255, 0))
for percentile in range(0, 101):
    print('percentile', percentile)
    p2.plot(np.percentile(population_histories, percentile, 0), pen=(100, 10, 0))

# for hist in population_histories:
#     p2.plot(hist, pen=(33,33,33))

# p2.plot(np.random.normal(size=100), pen=(255,0,0), name="Red curve")
# p2.plot(np.random.normal(size=110)+5, pen=(0,255,0), name="Green curve")
# p2.plot(np.random.normal(size=120)+10, pen=(0,0,255), name="Blue curve")


if __name__ == '__main__':
    pg.exec()