import streamlit as st

# Must be the first Streamlit command
st.set_page_config(page_title="EPANET Simulation with WNTR + GPT Assistant", layout="wide")

import wntr
import tempfile
import os
import pandas as pd
import matplotlib.pyplot as plt
import openai

# Set OpenAI API Key
openai.api_key = st.secrets.get("OPENAI_API_KEY") or st.text_input("Enter your OpenAI API key", type="password")

st.title("💧 EPANET Simulation with WNTR + GPT Assistant")

# Upload EPANET .inp file
uploaded_file = st.file_uploader("Upload your EPANET .inp file", type=["inp"])
if uploaded_file:
    with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp:
        temp.write(uploaded_file.read())
        inp_path = temp.name

    st.success("File uploaded. Starting simulation...")

    wn = wntr.network.WaterNetworkModel(inp_path)
    sim = wntr.sim.EpanetSimulator(wn)
    results = sim.run_sim()
    pressure = results.node["pressure"]
    st.session_state['pressure_df'] = pressure

    st.subheader("Simulation Complete")
    st.line_chart(pressure)
    st.download_button("Download Pressure Data (CSV)", data=pressure.to_csv().encode('utf-8'),
                       file_name="pressure_results.csv", mime="text/csv")

# GPT-4 Assistant Section
if 'pressure_df' in st.session_state:
    st.subheader("🧠 Ask the GPT-4 Assistant About Simulation Results")

    user_question = st.text_area("Ask a question about the simulation:")
    if st.button("Ask GPT"):
        df = st.session_state['pressure_df']
        prompt = f"""You are a hydraulic engineering assistant. The user has run a water network simulation. 
Here is a summary of node pressures (mean, min, max, std):

{df.describe().to_string()}

Now answer the user's question: "{user_question}"
""" 

        with st.spinner("Thinking..."):
            response = openai.ChatCompletion.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a civil engineering assistant who specializes in hydraulic analysis using EPANET and WNTR."},
                    {"role": "user", "content": prompt}
                ]
            )
            st.success("Answer:")
            st.markdown(response.choices[0].message['content'])