from collections import defaultdict

import seaborn as sns
import streamlit as st
import streamlit.components.v1 as components
import plotly.graph_objects as go
import plotly.express as px
import numpy as np
import pandas as pd
from scipy.stats import norm


st.set_page_config(page_title='Can you be truly random ?', layout = 'wide', page_icon = 'favicon.jpg', initial_sidebar_state = 'auto')

cm = sns.light_palette("green", as_cmap=True)


# Custom CSS to styles
st.markdown("""
    <style>
        button {
            padding-top: 50px !important;
            padding-bottom: 50px !important;
        }
    </style>
""", unsafe_allow_html=True)




def set_state(x):
    if x == 1: st.toast('The journey begins!', icon='😍')
    st.session_state.stage = x

        
def reset_game():
    set_state(0)
    if 'n_buttons' not in st.session_state:
        st.session_state.n_buttons = 2
    st.session_state.history = []
    st.session_state.preds = defaultdict(lambda: defaultdict(int))
    st.session_state.pnl = [0]
    st.session_state.min = [0]
    st.session_state.max = [0]
    st.session_state.max_history = 10

if 'stage' not in st.session_state:
    reset_game()

st.title('Can your brain be random?')
n_plays = 50

if st.session_state.stage == 0:
    st.button('Begin', on_click=set_state, args=[1],use_container_width=True)
    
    st.session_state.n_buttons = st.slider(label="How many buttons to play with?", min_value=2, max_value=5,value=2)
    
    st.markdown(f'You will be presented with {st.session_state.n_buttons} buttons to randomly choose from')
    
    st.markdown(f"At each round, I will try to predict which button you click")
    st.session_state.max_history = st.slider(label="How much memory can I use?", min_value=2, max_value=20,value=10)
    
    st.markdown(f"I will examine {st.session_state.max_history}-long sequences you played in the past")
    st.markdown('This will help me make a guess')
    st.markdown(f"If I get it right, I earn {st.session_state.n_buttons-1} point(s), otherwise you earn 1 point")
    st.subheader(f'This is a fair game, as you have {(1-1/st.session_state.n_buttons):.0%} chances of winning')
    st.markdown(f"Play as long as you want and try to beat me!")
    st.markdown(f"I will start showing you my strategy from {n_plays} rounds")
    
        
    
def make_pred(max_history=st.session_state.max_history,alpha=2):
    history = st.session_state.history
    
    denominator = 0
    scores = {str(i):0 for i in range(st.session_state.n_buttons)}
    recent = np.array(history[-max_history:])
    data = []
    for i in range(1,len(history)):
        past = np.array(history[-i-max_history:-i])
        
        played = history[-i]
        decay = np.exp(-alpha*np.linspace(1,0,len(past)))
        same = past == recent[-len(past):]
        similarity = (same*decay).sum() / decay.sum()
        weight = np.exp(-alpha*i/(len(history)+1)) * similarity # len(history)/(len(history)+i) * similarity
        scores[played] += weight
        denominator += weight
        if similarity > 0.5:
            data.append([played,similarity,weight,i] + [None]* (max_history-len(past)) +list(same))
    if not denominator:
        return get_random_play(),data
    result = {str(i):scores[str(i)]/denominator for i in range(st.session_state.n_buttons)}
    return result, data
    
def get_random_play():
    return {str(i):1/st.session_state.n_buttons for i in range(st.session_state.n_buttons)}
        

def update_pnl(user_win):
    current_score = st.session_state.pnl[-1]
    current_score += -1 if user_win else (st.session_state.n_buttons-1)
    st.session_state.pnl.append(current_score)
    expected_change = (2**0.5) * ((st.session_state.n_buttons-1)/st.session_state.n_buttons)
    st.session_state.min.append(st.session_state.min[-1]-expected_change)
    st.session_state.max.append(st.session_state.max[-1]+expected_change)
    
def user_select(i,choice):
    st.session_state.history.append(str(i))
    if i == choice:
        st.toast("I win!", icon='🤪')
        update_pnl(user_win=False)
    else:
        st.toast('Well done!', icon='😍')
        update_pnl(user_win=True)
    
    # refresh_preds()

def compute_perf():
    data = np.array(st.session_state.pnl)
    if len(data) > 1:
        data = data[1:] - data[:-1]
        perf = data.mean()
        std = data.std()
        win = (data>0).mean()
        sharpe = perf / std
        return win,sharpe
    else:
        return 0,0




if st.session_state.stage == 1:
    pred,data = make_pred()
    
    choice = max(pred,key=pred.get)
    progress = min(1.,float(len(st.session_state.history))/n_plays)
    my_bar = st.progress(progress)
    
    cols = st.columns(st.session_state.n_buttons)
    for i,col in enumerate(cols):
        col.button(str(i),on_click=user_select, args=[str(i),choice],
        use_container_width=True)
        
    col1,col2 = st.columns(2)
    
    win,sharpe = compute_perf()
    
    if progress >= 1.:

        if sharpe > 0.1: 
            emoticon = "🚀"
        elif sharpe < 0:
            emoticon = "📉"
        else:
            emoticon = "🧠"
        text = f'%age win={win:.0%} Sharpe={sharpe:.3f} {emoticon}'
        col1.subheader('My earnings so far...')
        col1.write(text)
        dev = abs(st.session_state.pnl[-1] / st.session_state.max[-1])
        chance = 1 - norm.cdf(dev)
        col1.write(f'There is only {chance:.2%} chances of achieving this by chance')
        fig = px.line(st.session_state.pnl)
        fig.update_layout(showlegend=False)
        fig.update_layout(xaxis={'visible': True, 'showticklabels': False})
        col1.plotly_chart(fig, use_container_width=True)

        cols = ['you_played','similarity','weight','how_long_ago'] +  [f'same_{i}' for i in range(st.session_state.max_history)]
        data = pd.DataFrame(data,columns=cols).set_index('how_long_ago').sort_values('weight',ascending=False)[:10]
        col2.subheader(f'In the past, you played')
        col2.dataframe(data.style.background_gradient(cmap=cm), use_container_width=True)
        won = choice == st.session_state.history[-1]
        col2.write(f'So I chose {choice} (p={pred[choice]:.0%}) and ' + ('won 🚀' if won else 'lost 📉'))
        
    st.button('Start over', on_click=reset_game, args=[],
        use_container_width=True)