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

st.set_page_config(page_title="System Dynamics Simulator with Insights", layout="wide")
st.title("📈 System Dynamics Simulator using PySD + HF Transformers")
st.markdown("Upload a **.mdl** or **.xmile** file, simulate it, and gain insights using Hugging Face LLMs.")

HF_API_TOKEN = st.secrets["HF_API_TOKEN"] if "HF_API_TOKEN" in st.secrets else st.text_input("Enter Hugging Face API Token", type="password")
HF_MODEL_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"

def ask_llm(prompt):
    if not HF_API_TOKEN:
        return "🔒 API token required."
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    payload = {"inputs": prompt}
    response = requests.post(HF_MODEL_URL, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()[0]["generated_text"]
    else:
        return f"⚠️ Error: {response.text}"

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!")
        if st.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)

        st.subheader("🧠 Ask a Question about the Model")
        question = st.text_area("What would you like to know about the system dynamics model?")
        if st.button("Get Insight from LLM"):
            model_info = str(model.components)[:1000]  # trim context
            prompt = f"""The following is a system dynamics model fragment:

{model_info}

Question: {question}
"""
            response = ask_llm(prompt)
            st.markdown(f"**Answer:**\n\n{response}")

    except Exception as e:
        st.error(f"❌ Error: {str(e)}")