#import streamlit as st
#import os
#import numpy as np
#import pandas as pd
#import matplotlib.pyplot as plt
#from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
import streamlit as st
import os
import Groq
import requests

# Dependency imports with error handling
try:
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
    #from groq.api import Client
except ImportError as e:
    st.error(f"Error importing dependencies: {e}")
    st.stop()

# Ensure compatible versions of numpy and pandas
try:
    np_version = np.__version__
    pd_version = pd.__version__

    if tuple(map(int, np_version.split(".")[:2])) < (1, 24):
        raise ImportError(f"Numpy version is {np_version}. Please upgrade to 1.24.0 or newer.")
    if tuple(map(int, pd_version.split(".")[:2])) < (1, 5):
        raise ImportError(f"Pandas version is {pd_version}. Please upgrade to 1.5.0 or newer.")
except Exception as e:
    st.error(f"Version compatibility error: {e}")
    st.stop()

# Initialize the RAG components

GROQ_API_KEY = "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"
client = Groq(api_key = GROQ_API_KEY)

# Initialize the RAG components
try:
    tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
    retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
    model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
except Exception as e:
    st.error(f"Error loading RAG components: {e}")
    st.stop()

# Helper function for energy demand calculation
def calculate_energy(appliances):
    total_energy = 0
    total_power = 0
    for app, details in appliances.items():
        total_energy += details['quantity'] * details['wattage'] * details['hours'] / 1000  # Energy in kWh
        total_power += details['quantity'] * details['wattage']  # Power in Watts
    return total_energy, total_power

# Predefined appliances with their default wattage
predefined_appliances = {
    "Select Appliance": 0,  # Default placeholder
    "LED Light": 10,
    "Air Conditioner": 2000,
    "Refrigerator": 150,
    "Washing Machine": 500,
    "Fan": 75,
    "Laptop": 50,
    "Television": 100,
    "Microwave Oven": 1000,
    "Electric Water Heater": 1500,
    "Dishwasher": 1200,
    "Coffee Maker": 800,
    "Toaster": 1200,
    "Hair Dryer": 1500,
    "Electric Kettle": 1500,
    "Oven": 2000,
    "Vacuum Cleaner": 1000,
    "Iron": 1200,
    "Blender": 300,
    "Clothes Dryer": 1800,
    "Space Heater": 1500,
    "Dehumidifier": 700,
    "Freezer": 400,
    "Electric Grill": 1200,
    "Sewing Machine": 100,
    "Projector": 300,
    "Gaming Console": 200,
    "Electric Fan Heater": 1500
}

# Streamlit App
st.title("SustainaBot: Solar Energy Advisor")
st.sidebar.header("User Input")

# Option for Groq API Key
#groq_api_key = st.sidebar.text_input("Enter your Groq API Key", type="password")

# Inputs
st.sidebar.subheader("Enter Appliance Details")

if "appliances" not in st.session_state:
    st.session_state.appliances = {}

# Function to reset appliance state
def reset_appliances():
    st.session_state.appliances = {}

# Add appliance inputs dynamically
appliance_count = len(st.session_state.appliances)
for i in range(appliance_count):
    with st.sidebar.expander(f"Appliance {i+1} Details"):
        app_name = st.selectbox(f"Appliance {i+1} Name", options=list(predefined_appliances.keys()), key=f"app{i}_name")
        quantity = st.number_input(f"Quantity of {app_name}", min_value=1, step=1, key=f"app{i}_quantity", value=1)
        wattage = predefined_appliances.get(app_name, 0)  # Default wattage
        hours = st.number_input(f"Hours used per day", min_value=1, step=1, key=f"app{i}_hours", value=1)
        power = quantity * wattage
        st.write(f"**Total Power (Watts):** {power} W")
        if app_name and app_name != "Select Appliance":
            st.session_state.appliances[app_name] = {"quantity": quantity, "wattage": wattage, "hours": hours}

# Button to add more appliances
if st.sidebar.button("Add Appliance"):
    st.session_state.appliances[f"Appliance {len(st.session_state.appliances) + 1}"] = {"quantity": 1, "wattage": 0, "hours": 0}

# Button to reset appliances
if st.sidebar.button("Reset Appliances"):
    reset_appliances()

# Input for location and budget
location = st.sidebar.text_input("Enter your Location (City, Country)")
budget = st.sidebar.number_input("Enter your Budget for Solar System (in USD)", min_value=100, step=10)

if st.sidebar.button("Submit"):
    # Calculate energy demand
    daily_energy, total_power = calculate_energy(st.session_state.appliances)
    st.write(f"### Total Daily Energy Demand: {daily_energy:.2f} kWh")
    st.write(f"### Total Power Demand: {total_power:.2f} W")

    # RAG model interaction
    try:
        if not groq_api_key:
            st.warning("No Groq API Key provided. RAG model will be used locally.")
        else:
            st.write("Using Groq API Key for enhanced recommendations.")

        prompt = f"Provide solar recommendations for a location: {location}, daily energy: {daily_energy:.2f} kWh, budget: {budget} USD."
        # Tokenizing and retrieval process
        inputs = tokenizer(prompt, return_tensors="pt")
        question_input = inputs['input_ids']
        doc_scores, docs = retriever(question_input)

        # Generating the response
        generated_output = model.generate(input_ids=question_input, context_input_ids=docs[0]['text'])
        answer = tokenizer.decode(generated_output[0], skip_special_tokens=True)

        st.write("### Solar Recommendations")
        st.write(answer)
    except Exception as e:
        st.error(f"Error during RAG model processing: {e}")

    # Display potential savings and pricing
    savings = daily_energy * 30 * 0.2  # Assuming $0.2/kWh savings
    st.write(f"By switching to solar, you could save approximately **${savings:.2f} per month!**")