import streamlit as st
import requests

# Hugging Face API setup
API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
headers = {"Authorization": "Bearer YOUR_HF_API_KEY"}

def query_llm(prompt):
    payload = {"inputs": prompt}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()[0]['generated_text']

# Static templates for known experiments
experiment_templates = {
    "Vinegar and Baking Soda": {
        "goal": "Observe a chemical reaction and gas formation.",
        "materials": "Vinegar, Baking Soda, Container",
        "default_hypothesis": "Mixing will cause bubbles or fizz due to gas release.",
        "result": "Bubbles and fizz from CO2 gas produced in an acid-base reaction.",
        "explanation": "Acetic acid in vinegar reacts with sodium bicarbonate to form carbon dioxide gas."
    },
    "Lemon Battery": {
        "goal": "Generate electricity using a lemon as a battery.",
        "materials": "Lemon, Copper coin, Zinc nail, Wires, LED",
        "default_hypothesis": "The lemon will generate voltage to light up a small LED.",
        "result": "LED glows slightly due to electron flow.",
        "explanation": "The citric acid acts as electrolyte between the copper and zinc electrodes."
    }
}

# Streamlit UI
st.set_page_config(page_title="Science Lab Assistant", layout="centered")
st.title("🧪 Science Lab Assistant")

# Select known or custom experiment
exp_option = st.selectbox("Choose an experiment", ["Custom"] + list(experiment_templates.keys()))

if exp_option != "Custom":
    data = experiment_templates[exp_option]
    st.subheader("📋 Experiment Summary")
    st.write(f"**Goal:** {data['goal']}")
    st.write(f"**Materials:** {data['materials']}")
    st.write(f"**Suggested Hypothesis:** {data['default_hypothesis']}")

    if st.button("🔍 Show Result & Explanation"):
        st.success(f"**Result:** {data['result']}")
        st.info(f"**Why?:** {data['explanation']}")

else:
    st.subheader("🔬 Describe your experiment")
    experiment_name = st.text_input("Name or short description of your experiment")
    goal = st.text_area("What is the goal of your experiment?")
    materials = st.text_area("List the materials involved")

    if st.button("🧠 Generate Hypothesis and Expected Result"):
        prompt = f"""I am conducting a school science experiment. 
        Experiment: {experiment_name}
        Goal: {goal}
        Materials: {materials}
        Please suggest a hypothesis and likely result with a brief scientific explanation."""
        result = query_llm(prompt)
        st.markdown("### 🤖 Assistant's Suggestion")
        st.write(result)