import streamlit as st
import os
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from groq import Groq
from io import BytesIO

GROQ_API_KEY = "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"  # Replace with your Groq API key
client = Groq(api_key=GROQ_API_KEY)

# Function to upload the dataset
def upload_file():
    uploaded_file = st.file_uploader("Upload your Excel dataset", type=["xlsx"])
    if uploaded_file is not None:
        return pd.read_excel(uploaded_file)
    return None

# Function to construct the graph and visualize
def construct_graph(dataset):
    G = nx.DiGraph()
    
    # Add nodes and edges based on the dataset
    for _, row in dataset.iterrows():
        # Replace with your column names if necessary
        G.add_edge(row['Vendor Name'], row['Meter Number'], weight=row['Consumption (HCF)'])
    
    # Visualize the graph (water distribution network)
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G)  # Layout for visualization
    nx.draw(G, pos, with_labels=True, node_color='skyblue', node_size=1500, edge_color='gray')
    
    # Render the graph as an image and return it
    buf = BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)
    return buf

# Function to query Groq for optimization suggestions based on serial number
def query_groq(serial_number):
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": f"Optimize water distribution network for Serial Number {serial_number} to minimize consumption fluctuations."
            }
        ],
        model="llama3-8b-8192",
    )
    return chat_completion.choices[0].message.content

# Streamlit UI
st.set_page_config(page_title="Water Distribution Network Optimization", page_icon="💧")

# Title and description
st.title("💧 Intelligent Water Distribution Network Optimization")
st.write("Upload your water distribution dataset and optimize the network based on specific serial numbers.")

# Step 1: Upload file
dataset = upload_file()

if dataset is not None:
    # Show uploaded dataset details
    st.subheader("Dataset Preview:")
    st.write(dataset.head())

    # Check if there's a serial number column, add if missing
    if 'Serial Number' not in dataset.columns:
        dataset['Serial Number'] = range(1, len(dataset) + 1)  # Adding Serial Number column
    
    # Step 2: Select Serial Number for analysis
    serial_number = st.number_input("Enter Serial Number for analysis", min_value=1, max_value=int(dataset['Serial Number'].max()), step=1)
    
    if serial_number in dataset['Serial Number'].values:
        # Filter the dataset for the selected serial number
        filtered_data = dataset[dataset['Serial Number'] == serial_number]
        st.write(f"Filtered data for Serial Number {serial_number}:", filtered_data)
        
        # Step 3: Display the graph
        st.subheader(f"Graph for Serial Number {serial_number}")
        graph_image = construct_graph(filtered_data)
        st.image(graph_image, caption=f"Water Distribution Network for Serial Number {serial_number}", use_column_width=True)
        
        # Step 4: Query Groq for optimization
        if st.button(f"Optimize for Serial Number {serial_number}"):
            optimization_response = query_groq(serial_number)
            st.subheader("Optimization Suggestions:")
            st.write(optimization_response)

else:
    st.info("Please upload a dataset to get started.")

# Footer
st.markdown("""
    ---
    Developed by [Your Name]. 
    For more information, visit our [website](https://example.com).
""")