from collections import defaultdict


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


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



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


max_history = 10

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]

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

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

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("At each round, I will try to predict which button you click :-)")
    st.markdown(f"If I get it right, I earn {st.session_state.n_buttons-1} point(s), otherwise you earn 1 point")
    st.markdown(f"Play as long as you want and try to beat me!")
    
        
    

# def get_prev_seqs(history):
#     # st.write(history)
    
#     seqs = []
#     for h in range(max_history+1):
#         if len(history) >= h:
#             previous_seq = ''.join(history[-h:]) if h > 0 else ""
#             # st.write(previous_seq)
#             seqs.append(previous_seq) # from small to largest
#     return seqs
    
# def refresh_preds():
#     played = st.session_state.history[-1]
#     seqs = get_prev_seqs(st.session_state.history[:-1])
#     for seq in seqs:
#         # st.write(f'"{seq}"',played)
#         st.session_state.preds[seq][played] += 1

def make_pred(max_history=max_history,alpha=0.5):
    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:])
    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)))
        similarity = ((past == recent[-len(past):])*decay).sum() / decay.sum()
        weight = len(history)/(len(history)+i) * similarity
        scores[played] += weight
        denominator += weight
    if not denominator:
        return get_random_play()
    result = {str(i):scores[str(i)]/denominator for i in range(st.session_state.n_buttons)}
    return result
    
def get_random_play():
    return {str(i):1/st.session_state.n_buttons for i in range(st.session_state.n_buttons)}

# def make_pred():
#     seqs = get_prev_seqs(st.session_state.history)
#     # st.write('seqs',seqs)
#     scores = defaultdict(float)
#     denominator = 0
#     for i,seq in enumerate(seqs):
#         weight = (i+1)**2
#         preds = st.session_state.preds[seq]
#         total = sum(preds.values())
        
#         if total:
#             for played,value in preds.items():
#                 scores[played] += weight*value/total
#             denominator += weight
#     if denominator:
#         scores = {played:value/denominator for played,value in scores.items()}
#         return scores
#     else:
#         return get_random_play()
        

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 f'%age win={win:.2%} Sharpe={sharpe:.2f}'
    

if st.session_state.stage == 1:
    pred = make_pred()
    choice = max(pred,key=pred.get)
    
    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)

    st.subheader(f'My earnings so far... {compute_perf()} :-)')
    fig = px.line(st.session_state.pnl)
    fig.update_layout(showlegend=False)
    st.plotly_chart(fig, use_container_width=True)
    
    st.button('Start over', on_click=reset_game, args=[],
        use_container_width=True)