import streamlit as st
from normstrategy import MVStrategy, ScalingStrategy, KNNStrategy

if "data" in st.session_state:
    data = st.session_state.original_data
    st.session_state.original_data = data.copy()

    for column, series in data.items():
        col1, col2 = st.columns(2)
        missing_count = series.isna().sum()
        choices = MVStrategy.list_available(data, column, series)
        option = col1.selectbox(
            f"Missing values of {column} ({missing_count})",
            choices,
            index=1,
            key=f"mv-{column}",
        )
        if isinstance(option, KNNStrategy):
            option.training_features = st.multiselect("Training columns", option.training_features, default=option.available_features, key=f"cols-{column}")
            option.n_neighbors = st.number_input("Number of neighbors", min_value=1, max_value=option.count_max(data, column), value=option.n_neighbors, key=f"neighbors-{column}")
        # Always re-get the series to avoid reusing an invalidated series pointer
        data = option.apply(data, column, data[column])

        choices = ScalingStrategy.list_available(data, series)
        option = col2.selectbox(
            "Scaling",
            choices,
            key=f"scaling-{column}",
        )
        data = option.apply(data, column, data[column])

    st.write(data)
    st.session_state.data = data
else:
    st.error("file not loaded")