import gradio as gr
import pandas as pd
from autots import AutoTS
import plotly.graph_objects as go

# Function to handle forecasting
def forecast_queues(data):
    # Convert the input string to a pandas DataFrame
    data = pd.read_csv(data)
    
    # Ensure the 'date' column is in the correct datetime format
    data['date'] = pd.to_datetime(data['date'], errors='coerce')
    
    # Check for missing values and drop them if necessary
    data = data.dropna(subset=['date', 'queue_length'])
    
    # Setup AutoTS for time-series forecasting
    model = AutoTS(forecast_length=12, frequency='D', ensemble=True, model_list="all")
    model = model.fit(data, date_col='date', value_col='queue_length')  # Adjust column names as per your CSV
    
    # Generate forecast for the next 12 time periods
    forecast = model.predict()
    
    # Get the forecasted values
    forecast_df = forecast.forecast
    forecast_values = forecast_df['queue_length'].values.tolist()
    
    # Create a Plotly figure to visualize the forecast
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=forecast_df.index, y=forecast_values, mode='lines+markers', name='Forecast'))
    fig.update_layout(title="Queue Length Forecast", xaxis_title="Date", yaxis_title="Queue Length")
    
    return fig

# Gradio Interface function
def gradio_interface(file):
    return forecast_queues(file.name)

# Define Gradio interface with file input and plot output
iface = gr.Interface(fn=gradio_interface, 
                     inputs=gr.File(label="Upload your historical queue data (CSV with date and queue_length columns)"), 
                     outputs=gr.Plot(), 
                     live=True, 
                     title="Queue Length Forecasting with AutoTS",
                     description="Upload a CSV file containing historical queue data with a date and queue_length columns, and get a forecast for the next 12 periods.")

# Launch the app
if __name__ == "__main__":
    iface.launch()