import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import openai
import os

# Set your OpenAI key here or use Hugging Face Secrets Manager
openai.api_key = os.getenv("OPENAI_API_KEY")

st.set_page_config(page_title="Smart Factory RAG Assistant", layout="wide")

st.title("🏭 Industry 5.0 | Smart Factory RAG Assistant")

# File Upload
uploaded_file = st.file_uploader("📤 Upload your factory CSV data", type=["csv"])

if uploaded_file:
    df = pd.read_csv(uploaded_file)
    st.success("✅ File uploaded and loaded!")

    # Basic Preview
    st.subheader("📄 Data Preview")
    st.dataframe(df.head())

    # Descriptive Stats
    st.subheader("📊 Descriptive Statistics")
    st.dataframe(df.describe().T)

    # Correlation Analysis
    st.subheader("🔗 Parameter Correlation Heatmap")
    fig, ax = plt.subplots(figsize=(10, 6))
    corr = df.corr(numeric_only=True)
    sns.heatmap(corr, annot=True, cmap="coolwarm", fmt=".2f", ax=ax)
    st.pyplot(fig)

    # Anomaly Detection
    st.subheader("⚠️ Anomaly Detection using Isolation Forest")
    num_df = df.select_dtypes(include='number').dropna()
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(num_df)

    iso = IsolationForest(contamination=0.05)
    df['Anomaly'] = iso.fit_predict(X_scaled)
    anomalies = df[df['Anomaly'] == -1]
    st.write(f"Detected {len(anomalies)} anomalies")
    st.dataframe(anomalies.head(10))

    # Prepare context for GPT
    st.subheader("🧠 Role-Based Decision Assistant")
    role = st.selectbox("Select your role", ["Engineer", "Operator"])
    question = st.text_input("Ask a question based on the data analysis")

    if question:
        with st.spinner("Thinking..."):
            summary = df.describe().to_string()
            corr_text = corr.to_string()
            anomaly_count = len(anomalies)

            context = f"""
You are a highly skilled {role} working in a smart manufacturing facility.

Here is a summary of the uploaded data:

STATISTICAL SUMMARY:
{summary}

PARAMETER CORRELATIONS:
{corr_text}

ANOMALY DETECTION:
{anomaly_count} anomalies detected using Isolation Forest method.

Based on this context, answer the following question in a clear, technically accurate manner and suggest best decisions from the point of view of a {role}.
"""

            final_prompt = f"""{context}
QUESTION: {question}
ANSWER:"""

            try:
                response = openai.ChatCompletion.create(
                    model="gpt-4",
                    messages=[
                        {"role": "system", "content": f"You are an expert {role} in a smart factory."},
                        {"role": "user", "content": final_prompt}
                    ],
                    temperature=0.5,
                    max_tokens=500
                )
                answer = response['choices'][0]['message']['content']
                st.success("✅ Recommendation:")
                st.markdown(f"**{answer}**")
            except Exception as e:
                st.error(f"⚠️ Error calling GPT API: {str(e)}")

else:
    st.info("📂 Please upload a factory CSV data file to begin analysis.")