import gradio as gr
from PIL import Image
from matplotlib import pyplot as plt
import numpy as np
import pickle
import gzip
import seaborn as sns
sns.set_theme()

def flip(n, t, seed=13):
    np.random.seed(seed)
    flips = np.random.randint(0, 2, size=(n, t))
    heads = flips.cumsum(axis=1)
    return heads.astype(np.uint8)

def load(filename):
    # Adapted from https://gist.github.com/thearn/5424244
    file = gzip.GzipFile(filename, 'rb')
    data = file.read()
    object = pickle.loads(data)
    file.close()
    return object

def calc_outcome(p_win, win, loss, leverage):
    n = 1000000
    t = 60
    base=100
    # heads = flip(n, t)
    heads = load('flips.pkl.gz')
    tails = np.arange(t)+1 - heads
    p_loss = 1 - p_win
    log_win = np.log(1+leverage*win)
    log_loss = np.log(1-leverage*loss)
    log_outcome = heads*log_win + tails*log_loss
    outcome = base*(np.exp(log_outcome))
    return outcome

def process(p_win, win, loss, leverage):
    outcome = calc_outcome(p_win/100., win/100., loss/100., leverage/100.)
    ensemble = fig2img(plot_ensemble(outcome)[0])
    time = fig2img(plot_time(outcome)[0])
    richest = fig2img(plot_richest(outcome)[0])
    return ensemble, time, richest

def fig2img(fig):
    """Convert a Matplotlib figure to a PIL Image and return it"""
    import io
    buf = io.BytesIO()
    fig.savefig(buf)
    buf.seek(0)
    img = Image.open(buf)
    plt.close(fig)
    return img

def plot_ensemble(outcome):
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    ax.plot(outcome.mean(axis=0))
    ax.set_xlabel('Time')
    ax.set_ylabel('Average Wealth ($)')
    ax.set_title('Ensemble (Collective) Perspective')
    fig.tight_layout()
    return fig, ax

def plot_richest(outcome):
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    ax.plot(outcome.max(axis=0)/1e6)
    ax.set_xlabel('Time')
    ax.set_ylabel('Wealth ($ million)')
    ax.set_title('Richest Individual')
    fig.tight_layout()
    return fig, ax

def plot_time(outcome):
    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
    ax.plot(np.percentile(outcome, 50, axis=0))
    ax.set_xlabel('Time')
    ax.set_ylabel('Median Wealth ($)')
    ax.set_title('Time (Individual) Perspective')
    fig.tight_layout()
    return fig, ax

css = ""

with gr.Blocks(css=css) as demo:
    with gr.Column():
        with gr.Row():
            gr.Markdown('Choose the probability of a positive outcome.')
        with gr.Row():
            p_win = gr.Slider(0, 100, value=50, label='Win Prob.(%)', interactive=True)
        with gr.Row():
            gr.Markdown('Choose the return, in percentage, for a positive (gain) and negative (loss) outcome. For example, a 50% gain and 40% loss means a bet of $100 will either result in $150 (50% gain) or $60 (40% loss).')
            win = gr.Number(value=50, label='Gain (%)', interactive=True)
            loss = gr.Number(value=40, label='Loss (%)', interactive=True)
        with gr.Row():
            gr.Markdown('Choose the leverage, that is, the percentage of your total assets that you are willing to bet.')
            leverage = gr.Slider(0, 100, value=100, label='Leverage (%)', interactive=True)
    with gr.Row():
        btn = gr.Button('Go!')
    with gr.Column(scale=1, min_width=800):
        with gr.Row():
            ensemble = gr.Image(label="Ensemble Perspective", height=300, width=300, elem_id="plot_ensemble")
            time = gr.Image(label="Time Perspective", height=300, width=300, elem_id="plot_time")
            richest = gr.Image(label="Richest Individual", height=300, width=300, elem_id="plot_richest")
    with gr.Column(scale=1, min_width=800):
        with gr.Row():
            gr.Markdown('The ensemble perspective shows the average wealth of an individual over time (which may be heavily affected by outliers, that is, an extremely rich individual). That is the expected value of the bet assuming ergodicity (in layman terms, "do over").')
            gr.Markdown('The time perspective shows the median wealth, that is, the typical wealth of an individual. It is the likely outcome for a regular individual.')
            gr.Markdown('This plot shows the wealth of the richest individual of them all, the one that got the best lucky streak out of 100,000 simulated individuals.')           

    btn.click(process, [p_win, win, loss, leverage], [ensemble, time, richest])
    
demo.launch()