import streamlit as st
from eda_ts.constants import DF_KEY
import plotly.express as px
import pandas as pd

st.set_page_config(
    page_title="Seasonality",
    page_icon="📅",
    layout="wide"
)
st.title("📅 Seasonality")

agg_functions = {
    'Mean': 'mean',
    'Median': 'median',
    'Count': 'count',
    'Sum': 'sum',
    'Standard Deviation': 'std',
    'Minimum': 'min',
    'Maximum': 'max'
}
granularity_options = {
    "hour": ("Hour", lambda x: x.hour, None),
    "day": ("Day of Week", lambda x: x.dayofweek, None),
    "month": ("Month", lambda x: x.month, None)
}
label_mapping = {
    'hour': [f"{i:02d}:00" for i in range(24)],
    'day': ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun'],
    'month': ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'],
}



df = st.session_state[DF_KEY]

col1, col2, col3, col4 = st.columns(4)

with col1:
    column_to_visualize = st.selectbox("Select column to visualize", df.columns)

with col2:
    x_granularity = st.selectbox("Select x-axis granularity", list(granularity_options.keys()))

with col3:
    y_granularity = st.selectbox("Select y-axis granularity",
        [key for key in granularity_options.keys() if key != x_granularity]
    )
with col4:
    selected_agg = st.selectbox("Select aggregation function", options=list(agg_functions.keys()), index=0)

df_temp = df.copy()
for gran, (label, func, _) in granularity_options.items():
    df_temp[gran] = df_temp.index.map(func)

# Get unique values for each granularity
available_values = {
    gran: sorted(df_temp[gran].unique())
    for gran in granularity_options.keys()
}

# Create dynamic label mappings based on available data
dynamic_label_mapping = {
    'hour': {i: f"{i:02d}:00" for i in available_values['hour']},
    'day': dict(zip(available_values['day'], ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun'])),
    'month': dict(zip(available_values['month'], ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']))
}

heatmap_data = df_temp.pivot_table(
    index=y_granularity,
    columns=x_granularity,
    values=column_to_visualize,
    aggfunc=agg_functions[selected_agg]
)

# Apply labels only for available values
if x_granularity in dynamic_label_mapping:
    heatmap_data.columns = [dynamic_label_mapping[x_granularity][col] for col in heatmap_data.columns]
if y_granularity in dynamic_label_mapping:
    heatmap_data.index = [dynamic_label_mapping[y_granularity][idx] for idx in heatmap_data.index]

heatmap_df = heatmap_data.reset_index(names=[y_granularity])
heatmap_df = pd.melt(
    heatmap_df,
    id_vars=y_granularity,
    var_name=x_granularity,
    value_name=column_to_visualize
)

fig = px.density_heatmap(
    heatmap_df,
    x=x_granularity,
    y=y_granularity,
    z=column_to_visualize,
    title=f'{selected_agg} of "{column_to_visualize}" by [{x_granularity.capitalize()}] and [{y_granularity.capitalize()}]',
    labels={
        x_granularity: granularity_options[x_granularity][0],
        y_granularity: granularity_options[y_granularity][0],
        column_to_visualize: column_to_visualize
    },
    color_continuous_scale='YlGnBu'
)
fig.update_coloraxes(colorbar_title_text=selected_agg)

fig.update_layout(
    height=600,
    width=1000,
    xaxis_title=granularity_options[x_granularity][0],
    yaxis_title=granularity_options[y_granularity][0]
)
st.plotly_chart(fig, use_container_width=True)