import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sentence_transformers import SentenceTransformer
from transformers import pipeline
from sklearn.ensemble import IsolationForest
from io import BytesIO

# App Config
st.set_page_config(
    page_title="FactoryGPT 5.0 – Maintenance Dashboard",
    page_icon="🧠",
    layout="wide"
)

# Load NLP + Anomaly Models
EMBED_MODEL = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
GEN_MODEL = pipeline('text2text-generation', model='google/flan-t5-base')

# Custom Style
st.markdown("""
    <style>
    html, body, [class*="css"] {
        font-family: 'Segoe UI', sans-serif;
        background-color: #0f1117;
        color: #f0f0f0;
    }
    .card {
        background-color: #1a1c23;
        padding: 1rem;
        border-radius: 10px;
        margin-bottom: 1rem;
        box-shadow: 0 0 8px rgba(88,166,255,0.2);
    }
    .metric-box {
        background-color: #1f2937;
        padding: 0.75rem;
        border-radius: 8px;
        margin-top: 0.5rem;
    }
    </style>
""", unsafe_allow_html=True)

# Header
st.markdown("""
    <div style='text-align: center;'>
        <h1 style='color: #58a6ff;'>🏭 FactoryGPT 5.0 – Technical Maintenance Dashboard</h1>
        <p style='color: #bbb;'>Anomaly Monitoring • Parameter Trends • Role-Based Intelligence</p>
        <hr style='border-top: 2px solid #888;'>
    </div>
""", unsafe_allow_html=True)

# Upload CSV
uploaded_file = st.sidebar.file_uploader("📂 Upload sensor log (CSV)", type=["csv"])

if uploaded_file:
    df = pd.read_csv(uploaded_file)
    numeric_cols = df.select_dtypes(include=np.number).columns.tolist()

    # Run Anomaly Detection
    iso = IsolationForest(contamination=0.02)
    labels = iso.fit_predict(df[numeric_cols])
    df['status'] = ['❌ Fault Detected' if x == -1 else '✅ Healthy' for x in labels]
    df['maintenance_flag'] = ['🔧 Inspect Required' if x == -1 else '🟢 Stable' for x in labels]

    # NLP Embeddings for RAG
    if 'chunks' not in st.session_state or 'embeddings' not in st.session_state:
        chunks = [
            f"[Entry {i}] " + ", ".join([f"{col}: {row[col]:.2f}" for col in numeric_cols])
            for i, row in df.iterrows()
        ]
        embeddings = EMBED_MODEL.encode(chunks)
        st.session_state.chunks = chunks
        st.session_state.embeddings = embeddings

    # Dashboard Layout
    col1, col2 = st.columns(2)

    with col1:
        st.markdown("### 🧠 Machine Health Summary")
        status_counts = df['status'].value_counts()
        fig1, ax1 = plt.subplots()
        ax1.bar(status_counts.index, status_counts.values, color=["red", "green"])
        ax1.set_title("Health Status Count")
        ax1.set_ylabel("Instances")
        ax1.set_facecolor("#0f1117")
        ax1.tick_params(colors='white')
        ax1.title.set_color('white')
        ax1.yaxis.label.set_color('white')
        for spine in ax1.spines.values():
            spine.set_edgecolor('white')
        st.pyplot(fig1)

    with col2:
        st.markdown("### 🧪 Parameters Monitored")
        st.markdown(f"""
            <div class="card">
                <div style='line-height: 1.8;'>{' • '.join(numeric_cols)}</div>
            </div>
        """, unsafe_allow_html=True)

    st.markdown("### 📉 Sensor Trend (with Fault Overlay)")
    col3, col4 = st.columns(2)
    time_col = col3.selectbox("🕒 Time Column", ["None"] + list(df.columns))
    trend_param = col4.selectbox("📌 Select Parameter", numeric_cols)

    if time_col != "None":
        df = df.sort_values(by=time_col)

    x_vals = df[time_col] if time_col != "None" else df.index
    y_vals = df[trend_param]

    fig2, ax2 = plt.subplots()
    ax2.plot(x_vals, y_vals, label=trend_param, color="skyblue")
    ax2.scatter(
        x_vals[df['status'] == '❌ Fault Detected'],
        y_vals[df['status'] == '❌ Fault Detected'],
        color='red', label='Fault Detected', zorder=5
    )
    ax2.set_title(f"{trend_param} Trend")
    ax2.set_xlabel("Time" if time_col != "None" else "Index")
    ax2.set_ylabel(trend_param)
    ax2.legend()
    ax2.set_facecolor("#0f1117")
    ax2.tick_params(colors='white')
    ax2.title.set_color('white')
    ax2.xaxis.label.set_color('white')
    ax2.yaxis.label.set_color('white')
    for spine in ax2.spines.values():
        spine.set_edgecolor('white')
    st.pyplot(fig2)

    # Export Fault Records
    st.markdown("### 📤 Export Anomalies")
    anomalies_df = df[df['status'] == '❌ Fault Detected']
    buffer = BytesIO()
    anomalies_df.to_csv(buffer, index=False)
    st.download_button(
        label="⬇️ Download Fault Log",
        data=buffer.getvalue(),
        file_name="fault_anomalies.csv",
        mime="text/csv"
    )

    # Fault Correlation
    st.markdown("### 🏷️ Fault Distribution by Machine/Component")
    metadata_cols = [col for col in df.columns if 'machine' in col.lower() or 'component' in col.lower()]
    for meta in metadata_cols:
        st.markdown(f"**{meta} – Fault Frequency**")
        meta_counts = df[df['status'] == '❌ Fault Detected'][meta].value_counts()
        st.bar_chart(meta_counts)

    # Role Assistant
    st.markdown("### 💬 Technical Assistant by Role")
    roles = {
        "Operator": {
            "description": "Focus on equipment behavior. Spot abnormal patterns and guide simple actions.",
            "style": "Explain simply. Emphasize safety and when to alert maintenance."
        },
        "Maintenance": {
            "description": "Diagnose machine issues. Recommend parts to inspect or replace.",
            "style": "Use technical language. Mention symptoms and sensor causes."
        },
        "Engineer": {
            "description": "Analyze system behavior. Identify root causes or instability.",
            "style": "Use RCA format. Discuss fault thresholds, control issues, and next steps."
        }
    }

    role = st.selectbox("👤 Select your role", roles.keys())
    user_input = st.text_input("🧠 Ask a question")

    if user_input:
        query_vec = EMBED_MODEL.encode([user_input])[0]
        sims = np.dot(st.session_state.embeddings, query_vec)
        top_idxs = np.argsort(sims)[-5:][::-1]
        context = "\n".join([st.session_state.chunks[i] for i in top_idxs])

        prompt = (
            f"ROLE: {role}\n"
            f"RESPONSIBILITIES: {roles[role]['description']}\n"
            f"COMMUNICATION STYLE: {roles[role]['style']}\n\n"
            f"DATA CONTEXT:\n{context}\n\n"
            f"QUESTION:\n{user_input}\n\n"
            f"ANSWER:\n"
        )
        result = GEN_MODEL(prompt, max_length=400)[0]['generated_text']
        st.markdown(f"**🤖 Response:**\n\n{result}")

else:
    st.info("👈 Upload a sensor CSV file to generate your technical dashboard.")