app.py

import streamlit as st
import pandas as pd
import math
from fpdf import FPDF
import os
import matplotlib.pyplot as plt
from groq import Groq

# ========== 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 ==========
groq_api_key = os.getenv("GROQ_API_KEY")
groq_client = Groq(api_key=groq_api_key) if groq_api_key else 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, R, S, E, corrosion):
    return (P * R) / (S * E - 0.6 * P) + corrosion

def head_thickness(P, R, S, E, corrosion, head_type):
    if head_type == "Ellipsoidal":
        return (0.5 * P * R) / (S * E - 0.1 * P) + corrosion
    elif head_type == "Torispherical":
        return (0.885 * P * R) / (S * E - 0.1 * P) + corrosion
    elif head_type == "Hemispherical":
        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")
        R = st.number_input("Internal Radius (mm)", value=1000.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")

        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, R, S, E, corrosion)
        st.metric("Required Shell Thickness (mm)", f"{t_shell_calc:.2f}")

    # --- HEAD TAB ---
    with tabs[1]:
        t_head_calc = head_thickness(P, R, S, E, corrosion, head_type)
        st.metric(f"Required {head_type} Head Thickness (mm)", 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:", "✅ 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) else "❌ No")

    # --- SUMMARY TAB ---
    with tabs[5]:
        summary_data = {
            "Shell Thickness": t_shell_calc,
            "Head Thickness": t_head_calc,
            "Nozzle Safe": safe_nozzle,
            "PWHT Required": pwht_required(thickness),
            "Impact Test Required": impact_test_required(thickness),
        }
        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")
        pdf.chapter_body(str(summary_data))
        pdf_file = "results.pdf"
        pdf.output(pdf_file)
        with open(pdf_file, "rb") as f:
            st.download_button("📄 Download PDF Report", f, "results.pdf")

    # --- AI EXPLANATION TAB ---
    with tabs[6]:
        st.markdown("### 🤖 Ask AI for Explanation")
        if groq_client:
            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}"
                    chat_completion = groq_client.chat.completions.create(
                        messages=[{"role": "user", "content": prompt}],
                        model="llama-3.1-8b-instant",
                    )
                    explanation = chat_completion.choices[0].message.content
                    st.success("✅ AI Explanation Generated Below")
                    st.write(explanation)

            if st.session_state.ai_done:
                st.info("✨ AI explanation already generated. Rerun to refresh.")
        else:
            st.info("ℹ️ Add your GROQ_API_KEY in Hugging Face secrets to enable AI explanations.")

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}.")