import streamlit as st
import numpy as np
import plotly.express as px
from eda_ts.constants import DF_KEY

st.set_page_config(
    page_title="Distributions",
    page_icon="📊",
    layout="wide"
)

df = st.session_state[DF_KEY]
st.title("📊 Variable Distributions")

numeric_cols = df.select_dtypes(include=[np.number]).columns
selected_cols = st.multiselect("Select columns for distribution",
                                numeric_cols,
                                default=list(numeric_cols))  # Set default to all numeric columns

if selected_cols:
    cols_per_row = st.slider("Plots per row", min_value=1, max_value=3, value=2)

    for i in range(0, len(selected_cols), cols_per_row):
        cols = selected_cols[i:i + cols_per_row]
        figs = st.columns(cols_per_row)

        for j, col in enumerate(cols):
            with figs[j]:
                fig = px.histogram(
                    df,
                    x=col,
                    title=f"Distribution of {col}",
                    height=400
                )
                fig.update_layout(margin=dict(l=20, r=20, t=40, b=20))
                st.plotly_chart(fig, use_container_width=True)

    st.write("### Summary Statistics")
    st.dataframe(df[selected_cols].describe(), use_container_width=True)
else:
    st.info("Please select one or more columns to visualize")