# app.py - FactoryRAG: Q&A over Sensor Logs (Streamlit + Hugging Face)

import streamlit as st
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline

# Set page config
st.set_page_config(page_title="FactoryRAG - Sensor Logs", layout="wide")
st.title("🏭 FactoryRAG: Human-Centric AI for Sensor Log Analysis")

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

# Load CSV sensor log file
@st.cache_data
def load_logs():
    df = pd.read_csv("01-04T184148_000_mode1.csv")
    return df

# Convert log rows to natural language chunks
def convert_to_chunks(df):
    chunks = []
    for idx, row in df.iterrows():
        sentence = f"Log entry {idx}: " + ", ".join([f"{col}: {row[col]:.2f}" for col in df.columns])
        chunks.append(sentence)
    return chunks

# Load and embed logs
df = load_logs()
st.write("📊 Sensor Data Snapshot:", df.head())

if 'chunks' not in st.session_state:
    st.session_state.chunks = convert_to_chunks(df)
    st.session_state.embeddings = EMBED_MODEL.encode(st.session_state.chunks)

# Ask a question
query = st.text_input("🔍 Ask something about the sensor logs:")

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

    prompt = f"Answer based on the following logs:\n{context}\n\nQuestion: {query}"
    response = GEN_MODEL(prompt, max_length=256)[0]['generated_text']

    st.subheader("🤖 FactoryGPT Answer")
    st.write(response)

    st.markdown("### 🧑‍🏭 Human Feedback")
    st.radio("Is this answer acceptable?", ["Approve", "Correct", "Escalate"], horizontal=True)

    with st.expander("📄 Retrieved Log Context"):
        st.code(context)

elif len(df) == 0:
    st.warning("⚠️ No data loaded. Please check your CSV file.")