import os
import pandas as pd
import plotly.express as px
from groq import Groq
import gradio as gr



api = os.environ.get('GroqApi')
client = Groq(api_key=api)

# Function to process the traffic dataset
def process_traffic_data(csv_file):
    data = pd.read_csv(csv_file)
    peak_hours = data.groupby("Time")["Total"].sum().idxmax()
    worst_situation = data.groupby("Traffic Situation")["Total"].sum().idxmax()
    vehicle_contribution = (
        data[["CarCount", "BikeCount", "BusCount", "TruckCount"]]
        .sum()
        .apply(lambda x: round((x / data["Total"].sum()) * 100, 2))
    )
    summary = (
        f"Peak traffic hour: {peak_hours}. Worst traffic situation: {worst_situation}.\n"
        f"Vehicle Contribution to Traffic:\n"
        f"- Cars: {vehicle_contribution['CarCount']}%\n"
        f"- Bikes: {vehicle_contribution['BikeCount']}%\n"
        f"- Buses: {vehicle_contribution['BusCount']}%\n"
        f"- Trucks: {vehicle_contribution['TruckCount']}%\n"
    )
    return data, summary

# Function to generate suggestions using Groq API
def generate_suggestions(summary):
    input_content = (
        f"Here is the detailed traffic data summary:\n{summary}"
        f"Suggest realistic traffic optimization strategies with supporting stats."
    )
    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": input_content}],
        model="llama3-8b-8192",
    )
    return chat_completion.choices[0].message.content

# Function to create traffic visualization
def create_visualization(data):
    fig = px.bar(
        data,
        x="Day of the week",
        y="Total",
        color="Traffic Situation",
        title="Total Traffic by Day and Traffic Situation",
        labels={"Total": "Total Vehicles", "Day of the week": "Day"},
    )
    return fig.to_html()

# Gradio interface function
def analyze_traffic(file):
    data, summary = process_traffic_data(file.name)
    suggestions = generate_suggestions(summary)
    visualization_html = create_visualization(data)
    return summary, suggestions, visualization_html

# Gradio interface definition
with gr.Blocks() as demo:
    gr.Markdown("# Traffic Optimization Application")
    gr.Markdown(
        "Upload a CSV file with traffic data to generate traffic flow optimization suggestions and visualize traffic patterns."
    )

    with gr.Row():
        file_input = gr.File(label="Upload Traffic Data CSV")
        analyze_button = gr.Button("Analyze Traffic Data")

    with gr.Row():
        summary_output = gr.Textbox(label="Traffic Data Summary", lines=8)
        suggestions_output = gr.Textbox(label="Optimization Suggestions", lines=8)

    visualization_output = gr.HTML(label="Traffic Visualization")

    analyze_button.click(
        analyze_traffic, inputs=file_input, outputs=[summary_output, suggestions_output, visualization_output]
    )

# Launch Gradio app
if __name__ == "__main__":
    demo.launch()