import streamlit as st
import plotly.graph_objects as go
import plotly.express as px
import numpy as np

# Set the title of the app
st.title("Homer Simpson Meta-Learning with Hierarchical Reinforcement Learning Intrinsic Reward Lecture")

# Display the image with a caption
st.image("homer.webp", caption="Homer Simpson Meta-Learning HRL Lecture", use_column_width=True)

# Display and play the audio files
st.write("Audio Playback Meta-Learning with HRL Intrinsic Reward Lecture:")
st.audio("h0.wav", format="audio/wav")
st.audio("h1.wav", format="audio/wav")
st.write("Oh, sweet Homer's doughnuts! If that second .wav file ain't playin', just download the darn thing! Mmm... downloads...")
st.audio("h2.wav", format="audio/wav")
st.image("intrinsic_reward_formulation.png", caption='Intrinsic Reward Formulation')

# Define the grid and visitations
grid = np.zeros((6, 6))
visitations = {
    (0, 0): 7035, (1, 0): 3579, (2, 0): 1359, (2, 1): 1707, (3, 1): 520, (4, 1): 227,
    (4, 2): 243, (5, 1): 217, (5, 2): 181, (5, 0): 241, (4, 0): 267, (5, 3): 179,
    (4, 3): 1034, (3, 3): 2163, (2, 3): 2080, (0, 1): 3313, (1, 1): 3015, (0, 2): 1846,
    (0, 3): 1104, (0, 4): 351, (1, 4): 518, (1, 3): 1497, (1, 2): 2236, (2, 2): 2239,
    (2, 4): 842, (1, 5): 238, (2, 5): 217, (0, 5): 341, (3, 5): 382, (4, 5): 1872,
    (4, 4): 2038, (3, 4): 1684, (3, 0): 383, (3, 2): 1102, (5, 4): 198
}

# Fill the grid with visitations
for (x, y), count in visitations.items():
    grid[x, y] = count

# Calculate the total number of visitations
total_visitations = sum(visitations.values())

# Calculate the percentages
percentages = {state: (count / total_visitations) * 100 for state, count in visitations.items()}

# Print the percentages in the specified order
order = [
    (0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5),
    (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5),
    (2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5),
    (3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5),
    (4, 0), (4, 1), (4, 2), (4, 3), (4, 4), (4, 5),
    (5, 0), (5, 1), (5, 2), (5, 3), (5, 4), (5, 5)
]

st.title("State Visitations Visualization")

st.write("### State Visitations Percentages:")
for state in order:
    st.write(f"State {state}: {percentages.get(state, 0):.2f}%")

# Create a pie chart
labels = [f"State {state}" for state in visitations.keys()]
values = list(visitations.values())

fig_pie = go.Figure(data=[go.Pie(labels=labels, values=values)])
fig_pie.update_layout(title_text="State Visitations Pie Chart")
st.plotly_chart(fig_pie)

# Create a heatmap
fig_heatmap = px.imshow(grid, labels=dict(x="Column", y="Row", color="Visitations"),
                        x=list(range(6)), y=list(range(6)), title="State Visitations Heatmap")
fig_heatmap.update_xaxes(side="top")
st.plotly_chart(fig_heatmap)