import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


st.set_page_config(
    page_title='Widgets',
    page_icon=":1234:"
)

mobility_url = 'https://raw.githubusercontent.com/UIUC-iSchool-DataViz/is445_data/main/mobility.csv'
df = pd.read_csv(mobility_url)



st.title("Widgets in Streamlit apps")

st.markdown("""Using markdown for a reminder, we can use
            widgets in Streamlit (similar to `ipywidgets`).
 """)

st.subheader('Feedback Widget')

st.markdown("""We could try the 
            [feedback widget](https://docs.streamlit.io/develop/api-reference/widgets/st.feedback).
 """)

st.markdown("Using the following code:")
st.code("""
sentiment_mapping = ["one", "two", "three", "four", "five"]
selected = st.feedback("stars")
if selected is not None:
    st.markdown(f"You selected {sentiment_mapping[selected]} star(s).")
         """)

# sentiment_mapping = ["one", "two", "three", "four", "five"]
# selected = st.feedback("stars")
# if selected is not None:
#     st.markdown(f"You selected {sentiment_mapping[selected]} star(s).")

st.write("How are you feeling right now?")
sentiment_mapping = ["one", "two", "three", "four", "five"]
selected = st.feedback("stars")
if selected is not None: # only run if a star is selected
    if selected < 1:
        st.markdown("Sorry to hear you are so sad :(")
    elif selected < 3:
        st.markdown("A solid medium is great!")
    else:
        st.markdown("Fantastic to hear you are having a great day!")

st.subheader("Connecting Widgets to Plots")

st.markdown("We'll start with a static plot:")

# bins along student to teacher ratio
bins = np.linspace(df['Student_teacher_ratio'].min(),
                   df['Student_teacher_ratio'].max(), 10)
#bins # note -- this will be "pandas-like" in view

table = df.pivot_table(index='State', 
                       columns=pd.cut(df['Student_teacher_ratio'],bins),
                       aggfunc='size')

fig, ax = plt.subplots()
extent = [bins.min(),bins.max(), 0, len(table.index)] # xmin, xmax, ymin, ymax
ax.imshow(table.values, cmap='hot', interpolation='nearest',extent=extent)
ax.set_yticks(range(len(table.index)))
ax.set_yticklabels(table.index)
st.pyplot(fig)

# trick for imshow command -- save as buffer so that its not so
#  big and we can format the size the way we want
# (might not need to do this for all plots)

from io import BytesIO
fig, ax = plt.subplots(figsize=(4,8))
extent = [bins.min(),bins.max(), 0, len(table.index)] # xmin, xmax, ymin, ymax
ax.imshow(table.values, cmap='hot', interpolation='nearest',extent=extent)
ax.set_yticks(range(len(table.index)))
ax.set_yticklabels(table.index)
#st.pyplot(fig)

buf = BytesIO()
fig.tight_layout()
fig.savefig(buf, format='png')
st.image(buf, width=500)

st.markdown("""Let's make use of a 
            [multiselect widget](https://docs.streamlit.io/develop/api-reference/widgets/st.multiselect) """)

fig_col, controls_col = st.columns([2,1], 
                                   vertical_alignment='center')

states_selected = controls_col.multiselect("Which states do you want to view", table.index.values)

if len(states_selected) > 0:
    df_subset = df[df['State'].isin(states_selected)]
    #st.write(df_subset) # used to debug our selection
    table_subset = df_subset.pivot_table(index='State', 
                       columns=pd.cut(df_subset['Student_teacher_ratio'],bins),
                       aggfunc='size')
    fig, ax = plt.subplots(figsize=(4,8))
    extent = [bins.min(),bins.max(), 0, len(table_subset.index)] # xmin, xmax, ymin, ymax
    ax.imshow(table_subset.values, cmap='hot', interpolation='nearest',extent=extent)
    ax.set_yticks(range(len(table_subset.index)))
    ax.set_yticklabels(table_subset.index)
    #st.pyplot(fig)

    buf = BytesIO()
    fig.tight_layout()
    fig.savefig(buf, format='png')
    #st.image(buf, width=500)
    fig_col.image(buf, width=400)
else:
    #st.write(df) # used to debug selection
    #pass
    fig, ax = plt.subplots(figsize=(4,8))
    extent = [bins.min(),bins.max(), 0, len(table.index)] # xmin, xmax, ymin, ymax
    ax.imshow(table.values, cmap='hot', interpolation='nearest',extent=extent)
    ax.set_yticks(range(len(table.index)))
    ax.set_yticklabels(table.index)
    #st.pyplot(fig)

    buf = BytesIO()
    fig.tight_layout()
    fig.savefig(buf, format='png')
    #st.image(buf, width=500)
    fig_col.image(buf, width=400)