from collections import defaultdict


import streamlit as st
import streamlit.components.v1 as components
import plotly.express as px


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 = 4


def set_state(x):
    if x == 1: st.toast('The journey begins!', icon='😍')
    if x == 2: st.toast('Welcome!', icon='🙃')
    if x == 3: st.toast("Let's play!", icon='🥰')
    if x == 4: st.toast("Meet your friends!", 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]

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

st.title('How random can you truly be?')

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} 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():
    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 {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)
    
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()
    
    

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.plotly_chart(px.line(st.session_state.pnl), use_container_width=True)
    
    st.button('Start over', on_click=reset_game, args=[],
        use_container_width=True)