import altair as alt
import numpy as np
import pandas as pd
import streamlit as st
from scenario_utils import build_scenarios, rank_scenarios

"""
# Welcome to Streamlit!

Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
forums](https://discuss.streamlit.io).

In the meantime, below is an example of what you can do with just a few lines of code:
"""

num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
num_turns = st.slider("Number of turns in spiral", 1, 300, 31)

indices = np.linspace(0, 1, num_points)
theta = 2 * np.pi * num_turns * indices
radius = indices

x = radius * np.cos(theta)
y = radius * np.sin(theta)

df = pd.DataFrame({
    "x": x,
    "y": y,
    "idx": indices,
    "rand": np.random.randn(num_points),
})

st.altair_chart(alt.Chart(df, height=700, width=700)
    .mark_point(filled=True)
    .encode(
        x=alt.X("x", axis=None),
        y=alt.Y("y", axis=None),
        color=alt.Color("idx", legend=None, scale=alt.Scale()),
        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
    ))


st.divider()
st.subheader("Scenario Simulator (Prototype)")

enable_sim = st.checkbox("Enable scenario simulation", value=False)

if enable_sim:
    st.info("Select variables and candidate values to simulate different transplant scenarios.")

    # --- Choose variables (edit names to match your model inputs) ---
    candidate_vars = st.multiselect(
        "Variables to vary (prototype list)",
        options=[
            "DonorType",
            "ConditioningIntensity",
            "Age",
            "HLA_Match"
        ]
    )

    variable_options = {}
    for v in candidate_vars:
        if v in ["Age"]:
            vals = st.slider(f"{v} range", min_value=10, max_value=80, value=(30, 60), step=1)
            variable_options[v] = list(range(vals[0], vals[1] + 1, 5))  # step of 5 years
        else:
            vals = st.text_input(f"{v} candidate values (comma-separated)", "")
            if vals.strip():
                variable_options[v] = [x.strip() for x in vals.split(",") if x.strip()]

    max_scen = st.number_input("Max scenarios", min_value=50, max_value=5000, value=500, step=50)
    objective = st.selectbox(
        "Objective",
        options=["min_gvhd", "min_gvhd_max_survival", "max_survival"],
        index=1
    )

    if st.button("Run simulation"):
        if not variable_options:
            st.error("Select at least one variable and provide candidate values.")
        else:
            baseline_row = input_df.iloc[0]  # <-- adjust to your baseline 1-row dataframe name
            scenario_df = build_scenarios(baseline_row, variable_options, max_scenarios=int(max_scen))

            # TODO: reuse your existing preprocess + inference calls here
            st.write("Generated scenarios:", scenario_df.shape[0])
            st.dataframe(scenario_df.head(20))