import os
import pandas as pd
import streamlit as st
import matplotlib.pyplot as plt
import seaborn as sns
from groq import Groq
from typing import List

# ---- Setup Groq Client ---- #
client = Groq(api_key=("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"))

# ---- Streamlit App Config ---- #
st.set_page_config(page_title="🚦 Traffic Optimization RAG App")

# ---- Load and Hardcode CSV Data ---- #
@st.cache_data
def load_data():
    return pd.read_csv("Traffic.csv")  # Replace with your file path if needed

df = load_data()

# ---- Utility: Summarize traffic data to context ---- #
def summarize_traffic_data(df: pd.DataFrame) -> List[str]:
    summaries = []
    for index, row in df.iterrows():
        summary = (
            f"On {row['Date']} ({row['Day of the week']}) at {row['Time']}, "
            f"there were {row['CarCount']} cars, {row['BikeCount']} bikes, "
            f"{row['BusCount']} buses, and {row['TruckCount']} trucks. "
            f"Total traffic: {row['Total']}. Situation: {row['Traffic Situation']}."
        )
        summaries.append(summary)
    return summaries

# ---- Prompt Constructor ---- #
def generate_traffic_prompt(user_query: str, context: List[str]) -> str:
    context_text = "\n".join(context)
    prompt = f"""
Context:
{context_text}

User Query:
{user_query}

Based on the context above, generate a traffic optimization strategy.
"""
    return prompt

# ---- Groq Query ---- #
def get_optimization_recommendation(prompt: str) -> str:
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="gpt-3.5",
        stream=False
    )
    return response.choices[0].message.content

# ---- Streamlit App ---- #
st.title("🚦 Real-Time Traffic Optimization using RAG + Groq")

# ---- Data Visualization ---- #
st.subheader("📊 Traffic Data Visualization")
st.dataframe(df.head(10))

# Handle 'Date' and 'Time' columns
df['Date'] = df['Date'].astype(str)
df['Time'] = df['Time'].astype(str)

# Handle missing values if any
df['Date'] = df['Date'].fillna('')
df['Time'] = df['Time'].fillna('')

# Combine 'Date' and 'Time' to create a timestamp
df['Timestamp'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], errors='coerce')  # `errors='coerce'` will turn invalid dates into NaT

# Check if 'Timestamp' is created successfully
st.write(f"Timestamp column created successfully: {df['Timestamp'].head()}")

# ---- Traffic Volume Over Time ---- #
st.write("### Traffic Volume Over Time")
fig, ax = plt.subplots(figsize=(10, 4))
df_sorted = df.sort_values("Timestamp")
ax.plot(df_sorted['Timestamp'], df_sorted['Total'], marker='o')
ax.set_xlabel("Time")
ax.set_ylabel("Total Traffic Volume")
ax.set_title("Traffic Volume Over Time")
st.pyplot(fig)

# ---- Vehicle Count Distribution ---- #
st.write("### Vehicle Count Distribution")
fig2, ax2 = plt.subplots(figsize=(10, 4))
df[['CarCount', 'BikeCount', 'BusCount', 'TruckCount']].plot(kind='box', ax=ax2)
ax2.set_title("Distribution of Vehicle Counts")
st.pyplot(fig2)

# ---- User Query and RAG Output ---- #
user_query = st.text_area("Enter your traffic-related query")
if user_query:
    with st.spinner("Processing traffic data and generating strategy..."):
        traffic_context = summarize_traffic_data(df)
        prompt = generate_traffic_prompt(user_query, traffic_context[:10])  # Limit to first 10 rows
        result = get_optimization_recommendation(prompt)
        st.success("Strategy Generated:")
        st.write(result)

st.markdown("---")
st.caption("🔁 This app analyzes traffic data using RAG + Groq and visualizes traffic patterns.")