app.py

import streamlit as st
import pandas as pd
from fpdf import FPDF
import io
import openai

# ========== PAGE CONFIG ==========
st.set_page_config(page_title="ASME Calculator", layout="wide")

# App Title + Description
st.markdown("<h1 style='text-align: center;'>🛠️ ASME CALCULATOR</h1>", unsafe_allow_html=True)
st.markdown(
    "<p style='text-align: center;'>This tool calculates <b>required thicknesses, nozzle reinforcement, "
    "PWHT, and impact test requirements</b><br>using ASME Section VIII Division 1 formulas.</p>",
    unsafe_allow_html=True
)

# ========== API CLIENT ==========
if "OPENAI_API_KEY" in st.secrets:
    openai.api_key = st.secrets["OPENAI_API_KEY"]
else:
    openai.api_key = None

# ========== PDF GENERATOR ==========
class PDF(FPDF):
    def __init__(self):
        super().__init__()
        self.add_page()
        self.set_font("Helvetica", "", 12)

    def header(self):
        self.set_font("Helvetica", "B", 14)
        self.cell(0, 10, "ASME VIII Div.1 Vessel Design Report", 0, 1, "C")

    def chapter_title(self, title):
        self.set_font("Helvetica", "B", 12)
        self.cell(0, 10, title, 0, 1, "L")

    def chapter_body(self, body):
        self.set_font("Helvetica", "", 11)
        self.multi_cell(0, 8, body)

# ========== CALCULATION FUNCTIONS ==========
def shell_thickness(P, D, S, E, corrosion):
    R = D / 2  # convert diameter to radius
    return (P * R) / (S * E - 0.6 * P) + corrosion

def head_thickness(P, D, S, E, corrosion, head_type):
    R = D / 2
    if head_type == "Ellipsoidal":  # UG-32
        return (0.5 * P * R) / (S * E - 0.1 * P) + corrosion
    elif head_type == "Torispherical":  # Approx formula
        return (0.885 * P * R) / (S * E - 0.1 * P) + corrosion
    elif head_type == "Hemispherical":  # UG-32
        return (P * R) / (2 * S * E - 0.2 * P) + corrosion
    return None

def nozzle_reinforcement(P, d, t_shell, t_nozzle, S, E):
    return (P * d) / (2 * S * E) <= (t_shell + t_nozzle)

def pwht_required(thickness, material="CS"):
    return material == "CS" and thickness > 38

def impact_test_required(thickness, MDMT=-20, material="CS"):
    return material == "CS" and (MDMT < -29 and thickness > 12)

# ========== SESSION STATE ==========
if "run_done" not in st.session_state:
    st.session_state.run_done = False
if "ai_done" not in st.session_state:
    st.session_state.ai_done = False

# ========== SIDEBAR INPUTS ==========
with st.sidebar.expander("📥 Manual Design Inputs", expanded=True):
    input_mode = st.radio("Input Mode:", ["Manual Entry", "Upload CSV"])
    run_calculation = False

    if input_mode == "Manual Entry":
        P = st.number_input("Design Pressure (MPa)", value=2.0, format="%.2f")
        D = st.number_input("Internal Diameter (mm)", value=2000.0, format="%.1f")
        S = st.number_input("Allowable Stress (MPa)", value=120.0, format="%.1f")
        corrosion = st.number_input("Corrosion Allowance (mm)", value=1.5, format="%.2f")

        joint_method = st.radio("Joint Efficiency Selection", ["Preset (UW-12)", "Manual Entry"])
        if joint_method == "Preset (UW-12)":
            E = st.selectbox("Select E (Joint Efficiency)", [1.0, 0.85, 0.7, 0.65, 0.6, 0.45])
        else:
            E = st.number_input("Manual Joint Efficiency (0-1)", value=0.85, min_value=0.1, max_value=1.0)

        head_type = st.selectbox("Head Type", ["Ellipsoidal", "Torispherical", "Hemispherical"])
        d_nozzle = st.number_input("Nozzle Diameter (mm)", value=200.0, format="%.1f")
        t_shell = st.number_input("Shell Thickness Provided (mm)", value=12.0, format="%.1f")
        t_nozzle = st.number_input("Nozzle Thickness Provided (mm)", value=10.0, format="%.1f")
        thickness = st.number_input("Governing Thickness (mm)", value=40.0, format="%.1f")
        MDMT = st.number_input("MDMT (°C)", value=-20.0, format="%.1f")

        if st.button("🚀 Run Calculation", use_container_width=True):
            st.session_state.run_done = True
            run_calculation = True

        if st.session_state.run_done:
            st.success("✅ Calculations completed! See results in the tabs.")

    else:
        uploaded_file = st.file_uploader("Upload CSV File", type=["csv"])
        if uploaded_file:
            df = pd.read_csv(uploaded_file)
            st.dataframe(df.head())
            if st.button("🚀 Run Calculation", use_container_width=True):
                st.session_state.run_done = True
                run_calculation = True

            if st.session_state.run_done:
                st.success("✅ Calculations completed! See results in the tabs.")

# ========== TABS ==========
tabs = st.tabs(["Shell", "Head", "Nozzle", "PWHT", "Impact Test", "Summary", "AI Explanation"])

if st.session_state.run_done:
    # --- SHELL TAB ---
    with tabs[0]:
        t_shell_calc = shell_thickness(P, D, S, E, corrosion)
        st.metric("Required Shell Thickness (mm) [UG-27]", f"{t_shell_calc:.2f}")

    # --- HEAD TAB ---
    with tabs[1]:
        t_head_calc = head_thickness(P, D, S, E, corrosion, head_type)
        st.metric(f"Required {head_type} Head Thickness (mm) [UG-32]", f"{t_head_calc:.2f}")

    # --- NOZZLE TAB ---
    with tabs[2]:
        safe_nozzle = nozzle_reinforcement(P, d_nozzle, t_shell, t_nozzle, S, E)
        st.write("Nozzle Reinforcement Check [UG-37]:", "✅ Safe" if safe_nozzle else "❌ Not Safe")

    # --- PWHT TAB ---
    with tabs[3]:
        st.write("PWHT Required:", "✅ Yes" if pwht_required(thickness) else "❌ No")

    # --- IMPACT TEST TAB ---
    with tabs[4]:
        st.write("Impact Test Required:", "✅ Yes" if impact_test_required(thickness, MDMT) else "❌ No")

    # --- SUMMARY TAB ---
    with tabs[5]:
        summary_data = {
            "Shell Thickness (UG-27)": round(t_shell_calc, 2),
            "Head Thickness (UG-32)": round(t_head_calc, 2),
            "Nozzle Safe (UG-37)": safe_nozzle,
            "PWHT Required": pwht_required(thickness),
            "Impact Test Required": impact_test_required(thickness, MDMT),
        }
        df_summary = pd.DataFrame([summary_data])
        st.dataframe(df_summary)

        # CSV export
        csv = df_summary.to_csv(index=False).encode("utf-8")
        st.download_button("📥 Download Results (CSV)", csv, "results.csv")

        # PDF export
        pdf = PDF()
        pdf.chapter_title("Calculation Summary")
        for k, v in summary_data.items():
            pdf.chapter_body(f"{k}: {v}")
        pdf_bytes = pdf.output(dest="S").encode("latin1")
        st.download_button("📄 Download PDF Report", pdf_bytes, "results.pdf", "application/pdf")

    # --- AI EXPLANATION TAB ---
    with tabs[6]:
        st.markdown("### 🤖 Ask AI for Explanation")
        if not openai.api_key:
            st.warning("⚠️ Add your OpenAI API key in Streamlit secrets to enable AI explanations.")
        else:
            if st.button("✨ Ask AI", use_container_width=True):
                st.session_state.ai_done = True
                with st.spinner("AI is preparing explanation..."):
                    prompt = f"Explain these ASME vessel design results in simple terms: {summary_data}"

                    try:
                        chat_completion = openai.ChatCompletion.create(
                            model="gpt-4o-mini",
                            messages=[
                                {"role": "system", "content": "You are an ASME design expert who explains results clearly."},
                                {"role": "user", "content": prompt}
                            ],
                            temperature=0.3
                        )
                        explanation = chat_completion["choices"][0]["message"]["content"]
                        st.success("✅ AI Explanation Generated Below")
                        st.write(explanation)

                    except Exception as e:
                        st.error(f"⚠️ Error generating AI explanation: {e}")

            if st.session_state.ai_done:
                st.info("✨ AI explanation already generated. Rerun to refresh.")

else:
    # Placeholders
    for i, msg in enumerate([
        "Shell results", "Head results", "Nozzle results",
        "PWHT decision", "Impact Test decision",
        "Summary", "AI explanation"
    ]):
        with tabs[i]:
            st.info(f"Run calculation to see {msg}.")

# ========== FOOTER ==========
st.markdown("---")
st.caption("Disclaimer: This tool is for demo/educational purposes. Final design must be verified by a qualified engineer per ASME VIII Div.1.")