import os
import streamlit as st
import faiss
import numpy as np
import pandas as pd
import plotly.express as px
from sentence_transformers import SentenceTransformer
from groq import Groq

# Load pre-trained embedding model
model = SentenceTransformer("all-MiniLM-L6-v2")

# Load FAISS index
index_path = "faiss_index.idx"
if os.path.exists(index_path):
    index = faiss.read_index(index_path)
else:
    index = None

# Load dataset (replace with actual datasets)
dataset_path = "electrolysis_data.csv"
if os.path.exists(dataset_path):
    df = pd.read_csv(dataset_path)
else:
    df = pd.DataFrame()  # Empty DataFrame as fallback

# ✅ Set Groq API Key (Replace with your actual key)
os.environ["GROQ_API_KEY"] = "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"
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Function for Retrieval-Augmented Generation (RAG)
def rag_pipeline(query):
    if index is not None:
        query_embedding = model.encode([query]).astype("float32")
        D, I = index.search(query_embedding, k=3)
        retrieved_docs = [df.iloc[i]["description"] for i in I[0] if i < len(df)]
        context = " ".join(retrieved_docs)
    else:
        context = "No relevant documents found."

    # Generate response using Llama 3 (via Groq API)
    response = client.chat.completions.create(
        messages=[
            {"role": "system", "content": f"Context: {context}"},
            {"role": "user", "content": query}
        ],
        model="llama3-70b-8192"
    )

    return response.choices[0].message.content

# Streamlit UI
st.set_page_config(page_title="HydroGen-AI", layout="wide")
st.title("🔬 HydroGen-AI: AI-Powered Hydrogen Techno-Economic Analyzer")

st.markdown("🚀 **Techno-Economic & Simulation-Based Analysis of Hydrogen Production via Electrolysis**")

# Left Panel: User Input
st.sidebar.header("🔧 Input Parameters")
water_source = st.sidebar.selectbox("Water Source", ["Atmospheric", "Seawater", "Groundwater", "Municipal"])
water_cost = st.sidebar.number_input("Water Cost ($/m³)", min_value=0.0, value=0.5)
purification_cost = st.sidebar.number_input("Water Purification Cost ($/m³)", min_value=0.0, value=0.1)
water_quantity = st.sidebar.number_input("Water Quantity (m³)", min_value=0.1, value=10.0)
method = st.sidebar.selectbox("Electrolysis Method", ["Alkaline", "PEM", "SOEC"])
current_density = st.sidebar.number_input("Current Density (A/cm²)", min_value=0.0, value=0.5)
voltage = st.sidebar.number_input("Cell Voltage (V)", min_value=1.0, value=1.8)
temperature = st.sidebar.number_input("Cell Temperature (°C)", min_value=25.0, value=60.0)
membrane = st.sidebar.selectbox("Membrane Type", ["Nafion", "AEM", "Zirfon"])
electrode = st.sidebar.selectbox("Electrode Material", ["Nickel", "Platinum", "Graphite"])

# Analysis Button
if st.sidebar.button("Analyze"):
    query = (f"Analyze hydrogen production with {method} electrolysis at {current_density} A/cm², {voltage} V, "
             f"{temperature}°C, using {membrane} membrane and {electrode} electrodes. Water: {water_quantity} m³ of {water_source} "
             f"(Cost: ${water_cost}, Purification: ${purification_cost}). Provide cost breakdown, efficiency, and production rate.")

    result = rag_pipeline(query)

    # Display Results
    st.subheader("📊 Analysis Result")
    st.write(result)
    st.success("✅ Analysis completed successfully!")

    # Visualization
    fig = px.bar(df, x="Parameter", y="Value", title="Techno-Economic Breakdown")
    st.plotly_chart(fig)