app.py

import os
import json
import re
import torch
import gradio as gr
import pandas as pd

# --- LangChain Imports ---
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.document_loaders import CSVLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings 
from langchain_classic.chains import RetrievalQA        
from langchain_core.prompts import PromptTemplate  

# --- 1. Setup API Key for Hugging Face Spaces ---
# HF Spaces uses standard environment variables instead of Colab secrets
GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
if not GOOGLE_API_KEY:
    raise ValueError("🔴 GOOGLE_API_KEY not found. Please add it to your Hugging Face Space Secrets.")

# --- 2. Build the RAG Chain & Feedback System ---
FMEA_DATA_FILE = '10000fmea_data.csv'
FEEDBACK_FILE = 'fmea_feedback.csv'
QA_CHAIN = None
feedback_vector_store = None
embeddings = None

DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"✅ Using device: {DEVICE}")

# --- FEEDBACK LOOP PART 1: Saving, Normalizing, and Loading Feedback ---
def normalize_action(text: str) -> str:
    return re.sub(r'\s+', ' ', str(text).strip().lower())

def load_feedback_stats():
    if not os.path.exists(FEEDBACK_FILE):
        return {}
    try:
        feedback_df = pd.read_csv(FEEDBACK_FILE)
        if feedback_df.empty:
            return {}
        stats = feedback_df.groupby('action')['rating'].agg(['mean', 'count']).to_dict('index')
        return stats
    except pd.errors.EmptyDataError:
        return {}

def save_feedback(action, rating, display_df):
    if not action:
        return "Please select a recommendation from the table first.", display_df
    norm_action = normalize_action(action)
    new_feedback = pd.DataFrame([{'action': norm_action, 'rating': int(rating)}])
    if not os.path.exists(FEEDBACK_FILE):
        new_feedback.to_csv(FEEDBACK_FILE, index=False)
    else:
        new_feedback.to_csv(FEEDBACK_FILE, mode='a', header=False, index=False)
    build_feedback_db()
    
    msg = f"✅ Rating of {rating}/10 saved for: {action}"
    
    # Update the displayed table dynamically
    if display_df is not None and not display_df.empty:
        try:
            feedback_stats = load_feedback_stats()
            default_stat = {'mean': 0, 'count': 0}
            stats_list = [feedback_stats.get(normalize_action(act), default_stat) for act in display_df['Recommended Action']]
            display_df['Avg. Feedback'] = [f"{stat['mean']:.2f}/10 ({int(stat['count'])})" for stat in stats_list]
        except Exception as e:
            print(f"Error updating display_df: {e}")

    return msg, display_df

def build_feedback_db():
    global feedback_vector_store
    if not os.path.exists(FEEDBACK_FILE):
        return
    try:
        feedback_df = pd.read_csv(FEEDBACK_FILE)
        if feedback_df.empty:
            return
    except pd.errors.EmptyDataError:
        return

    avg_ratings = feedback_df.groupby('action')['rating'].mean()
    highly_rated_actions = avg_ratings[avg_ratings > 7].index.tolist()

    if highly_rated_actions and embeddings:
        print(f"Found {len(highly_rated_actions)} highly-rated actions. Building feedback vector store...")
        feedback_vector_store = FAISS.from_texts(highly_rated_actions, embeddings)
        print("✅ Feedback vector store is ready.")

# --- build_rag_chain ---
def build_rag_chain():
    global QA_CHAIN, embeddings
    try:
        print("Initializing local HuggingFace embedding model...")
        embeddings = HuggingFaceEmbeddings(
            model_name='all-MiniLM-L6-v2',
            model_kwargs={'device': DEVICE} 
        )
        print("✅ Local embedding model loaded.")

        build_feedback_db()

        print(f"Loading FMEA data from {FMEA_DATA_FILE}...")
        loader = CSVLoader(file_path=FMEA_DATA_FILE, source_column="Failure_Mode")
        documents = loader.load()
        print(f"✅ Successfully loaded {len(documents)} records.")

        print("Creating embeddings and building main FAISS vector store...")
        main_vector_store = FAISS.from_documents(documents, embeddings)
        print("✅ Main vector store created successfully.")

        llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2)

        prompt_template = """
        You are an expert FMEA analyst. Your task is to generate the TOP 3 recommended actions for the given failure.
        The user has provided their current S, O, and D scores.
        For EACH recommendation, you must also estimate the revised S, O, and D scores (1-10) that would result *after* that action is successfully implemented.

        -   **new_S (Severity):** This score should *usually* stay the same as the original Severity.
        -   **new_O (Occurrence):** This score should be *lower* than the original Occurrence.
        -   **new_D (Detection):** This score should be *lower* than the original Detection (as the action makes the failure easier to detect).

        CONTEXT (Historical data and user feedback):
        {context}

        QUESTION (The new failure and its current scores):
        {question}

        INSTRUCTIONS:
        Format your entire response as a single, valid JSON object with a key "recommendations" which is a list of 3 objects.
        Each object must have these keys: "rank", "action", "department", "ai_score", "new_S", "new_O", "new_D".

        - "rank": The rank of the recommendation (1, 2, 3).
        - "action": The recommended action text.
        - "department": The most likely responsible department.
        - "ai_score": Confidence score (1-100) for this recommendation.
        - "new_S": Your estimated new Severity score (1-10).
        - "new_O": Your estimated new Occurrence score (1-10).
        - "new_D": Your estimated new Detection score (1-10).
        """
        PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"])

        # Included the token-saving "k": 2 limit
        retriever = main_vector_store.as_retriever(search_kwargs={"k": 2})
        QA_CHAIN = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=retriever,
            chain_type_kwargs={"prompt": PROMPT}
        )
        print("✅ RAG model is ready.")
        return True
    except Exception as e:
        print(f"🔴 An error occurred during RAG setup: {e}")
        return False

# --- 3. Gradio Interface Logic ---
def fmea_rag_interface(mode, effect, cause, severity, occurrence, detection):
    if QA_CHAIN is None:
        return "RAG Model is not initialized.", pd.DataFrame(), ""

    rpn = severity * occurrence * detection
    rpn_text = f"Current RPN (S×O×D): {int(rpn)}"

    query = (
        f"For a failure with Failure Mode='{mode}', Effect='{effect}', and Cause='{cause}', "
        f"what are the top 3 most appropriate recommended actions? "
        f"The current scores are: Severity={severity}, Occurrence={occurrence}, Detection={detection}."
    )

    docs = QA_CHAIN.retriever.invoke(query)
    context_from_history = "\n---\n".join([doc.page_content for doc in docs])

    context_from_feedback = ""
    if feedback_vector_store:
        feedback_docs = feedback_vector_store.similarity_search(query, k=3)
        if feedback_docs:
            feedback_actions = "\n".join([doc.page_content for doc in feedback_docs])
            context_from_feedback = f"\n\n--- Highly-Rated Actions from User Feedback ---\n{feedback_actions}"

    combined_context = f"--- Historical FMEA Entries ---\n{context_from_history}{context_from_feedback}"

    try:
        result = QA_CHAIN.invoke({"query": query, "context": combined_context})
        json_text = result["result"].strip().replace("```json", "").replace("```", "")
        data = json.loads(json_text)
        output_df = pd.DataFrame(data['recommendations'])

        feedback_stats = load_feedback_stats()
        default_stat = {'mean': 0, 'count': 0}
        stats_list = [feedback_stats.get(normalize_action(action), default_stat) for action in output_df['action']]
        output_df['avg_feedback'] = [stat['mean'] for stat in stats_list]
        output_df['feedback_count'] = [stat['count'] for stat in stats_list]

        output_df['new_S'] = output_df['new_S'].astype(int)
        output_df['new_O'] = output_df['new_O'].astype(int)
        output_df['new_D'] = output_df['new_D'].astype(int)
        output_df['new_RPN'] = output_df['new_S'] * output_df['new_O'] * output_df['new_D']

        rpn_change_list = [f"{int(rpn)} ➔ {int(new_rpn)}" for new_rpn in output_df['new_RPN']]

        display_df = pd.DataFrame({
            "Rank": output_df['rank'],
            "Recommended Action": output_df['action'],
            "Department": output_df['department'],
            "AI Confidence": [f"{score}%" for score in output_df['ai_score']],
            "Avg. Feedback": [f"{avg:.2f}/10 ({int(count)})" for avg, count in zip(output_df['avg_feedback'], output_df['feedback_count'])],
            "Revised RPN": rpn_change_list
        })

    except Exception as e:
        print(f"Error parsing LLM output: {e}")
        return rpn_text, pd.DataFrame({"Error": [f"Could not parse AI response: {e}"]}), None

    return rpn_text, display_df, output_df

def get_level_info(val):
    levels = {
        10: "Hazardous", 9: "Serious", 8: "Extreme", 7: "Major", 
        6: "Significant", 5: "Moderate", 4: "Minor", 3: "Slight", 
        2: "Very Slight", 1: "No Effect"
    }
    return gr.update(info=f"Level: {levels.get(val, '')}")

# --- 6. Main Application Execution ---
if build_rag_chain():
    print("\n🚀 Launching Gradio Interface...")
    with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.blue)) as demo:
        gr.Markdown("<h1>Pangun ReliAI-FMEA</h1>")

        with gr.Group():
            gr.Markdown("## FMEA Inputs ")
            with gr.Row():
                with gr.Column(scale=2):
                    f_mode = gr.Textbox(label="Failure Mode", placeholder="e.g., Engine Overheating")
                    f_effect = gr.Textbox(label="Effect", placeholder="e.g., Reduced vehicle performance")
                    f_cause = gr.Textbox(label="Cause", placeholder="e.g., Coolant leak")
                with gr.Column(scale=1):
                    f_sev = gr.Slider(1, 10, value=5, step=1, label="Severity", info="Level: Moderate")
                    f_occ = gr.Slider(1, 10, value=5, step=1, label="Occurrence", info="Level: Moderate")
                    f_det = gr.Slider(1, 10, value=5, step=1, label="Detection", info="Level: Moderate")

        f_sev.change(fn=get_level_info, inputs=f_sev, outputs=f_sev)
        f_occ.change(fn=get_level_info, inputs=f_occ, outputs=f_occ)
        f_det.change(fn=get_level_info, inputs=f_det, outputs=f_det)

        submit_btn = gr.Button("Get AI Recommendations", variant="primary")

        with gr.Group():
            gr.Markdown("## 💡 Top 3 AI-Generated Recommendations")
            rpn_output = gr.Textbox(label="Current RPN", interactive=False)
            recommendations_output = gr.DataFrame(
                headers=["Rank", "Recommended Action", "Department", "AI Confidence", "Avg. Feedback", "Revised RPN"],
                datatype=["number", "str", "str", "str", "str", "str"]
            )
            df_state = gr.State()

        with gr.Group():
            gr.Markdown("## ⭐ Provide Feedback")
            gr.Markdown("Click a row in the table above to select it, then submit your rating.")
            selected_action_text = gr.Textbox(label="Selected for Feedback", interactive=False)
            rating_slider = gr.Slider(minimum=1, maximum=10, step=1, value=8, label="Your Rating (1-10)")
            submit_feedback_btn = gr.Button("Submit Rating")
            feedback_status = gr.Textbox(label="Feedback Status", interactive=False)

        # FIX 1: New safer update_selection function
        def update_selection(table_df, evt: gr.SelectData):
            # Safety check if the table is empty
            if table_df is None or len(table_df) == 0: 
                return ""
            
            # evt.index gives us [row_index, column_index] of the click
            row_idx = evt.index[0]
            
            # "Recommended Action" is the 2nd column in your UI table (index 1)
            selected_action = table_df.iloc[row_idx, 1]
            return selected_action

        submit_btn.click(
            fn=fmea_rag_interface,
            inputs=[f_mode, f_effect, f_cause, f_sev, f_occ, f_det],
            outputs=[rpn_output, recommendations_output, df_state]
        )

        # FIX 2: Trigger relies on the visible table (recommendations_output) instead of df_state
        recommendations_output.select(
            fn=update_selection,
            inputs=[recommendations_output], # <-- Pass the visible table directly!
            outputs=[selected_action_text]
        )

        submit_feedback_btn.click(
            fn=save_feedback,
            inputs=[selected_action_text, rating_slider, recommendations_output],
            outputs=[feedback_status, recommendations_output]
        )

    # Simplified launch command for Hugging Face
    demo.launch()