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="Correlations",
    page_icon="🔗",
    layout="wide"
)

df = st.session_state[DF_KEY]
st.title("🔄 Correlation Analysis")


viz_mode = st.selectbox(
    "Visualization Mode",
    ["Heatmap", "Scatter Matrix"]
)

numeric_df = df.select_dtypes(include=[np.number])

if viz_mode == "Heatmap":
    corr_matrix = numeric_df.corr()
    fig = px.imshow(
        corr_matrix,
        title="Correlation Matrix",
        labels=dict(color="Correlation"),
        color_continuous_scale="RdBu"
    )
    st.plotly_chart(fig, use_container_width=True)

    # Add correlation value display
    st.write("### Correlation Values")
    st.dataframe(
        corr_matrix.round(2),
        use_container_width=True
    )

else:  # Scatter Matrix
    # Allow column selection for scatter matrix
    selected_cols = st.multiselect(
        "Select columns (2-6 recommended)",
        numeric_df.columns,
        default=list(numeric_df.columns[:4])
    )
    if selected_cols:
        fig = px.scatter_matrix(
            numeric_df[selected_cols],
            title="Scatter Plot Matrix",
            height=800
        )
        fig.update_layout(
            title_x=0.5,
            dragmode='select',
            width=None,
            height=None
        )
        st.plotly_chart(fig, use_container_width=True)
    else:
        st.info("Please select at least one column")
