import streamlit as st
from pysd import read_vensim, read_xmile
import pandas as pd
import matplotlib.pyplot as plt
import tempfile
import os

st.set_page_config(page_title="System Dynamics Simulator", layout="wide")

st.title("📈 System Dynamics Simulator using PySD")
st.markdown("Upload a **.mdl** or **.xmile** file to simulate and visualize system dynamics models.")

uploaded_file = st.file_uploader("Choose a Vensim (.mdl) or XMILE (.xmile) file", type=["mdl", "xmile"])

if uploaded_file is not None:
    suffix = os.path.splitext(uploaded_file.name)[-1].lower()

    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
        tmp_file.write(uploaded_file.getvalue())
        model_path = tmp_file.name

    try:
        if suffix == ".mdl":
            model = read_vensim(model_path)
        elif suffix == ".xmile":
            model = read_xmile(model_path)
        else:
            st.error("Unsupported file format.")
            st.stop()
        
        st.success("✅ Model loaded successfully!")
        
        st.sidebar.header("Simulation Settings (ignored in PySD<3.x)")
        initial_time = st.sidebar.number_input("Initial Time", value=0)
        final_time = st.sidebar.number_input("Final Time", value=100)
        time_step = st.sidebar.number_input("Time Step", value=1.0)
        
        if st.sidebar.button("Run Simulation"):
            result = model.run()
            st.subheader("📊 Simulation Output")
            st.dataframe(result)

            st.subheader("📉 Plot Time-Series")
            selected_vars = st.multiselect("Select variables to plot", result.columns.tolist(), default=result.columns.tolist())
            if selected_vars:
                for var in selected_vars:
                    fig, ax = plt.subplots()
                    ax.plot(result.index, result[var], label=var)
                    ax.set_xlabel("Time")
                    ax.set_ylabel(var)
                    ax.set_title(f"{var} over Time")
                    ax.grid(True)
                    st.pyplot(fig)
            else:
                st.warning("Select at least one variable to plot.")
    except Exception as e:
        st.error(f"❌ Error: {str(e)}")