import gradio as gr
import pandas as pd
import numpy as np
import pickle

# Load the trained model from the pickle file
with open('best_arima_models.pkl', 'rb') as f:
    model = pickle.load(f)

def predict_demand(mapped_code, num_months):
    try:
        print(f"Received mapped code: {mapped_code}")
        print(f"Number of months for prediction: {num_months}")
        
        # Retrieve the specific model for the mapped code
        if mapped_code not in model:
            return None, f"No model found for Mapped Code: {mapped_code}"
        
        model_for_code = model[mapped_code]
        
        # Generate a date range for the prediction period
        dates = pd.date_range(start=pd.Timestamp.today(), periods=num_months, freq='M')
        
        # Make predictions
        future_steps = len(dates)
        forecast = model_for_code.forecast(steps=future_steps)
        print(f"Forecast: {forecast}")

        # Prepare a DataFrame for display
        df = pd.DataFrame({
            'Date': dates.strftime('%Y-%m'),
            'Predicted Demand': forecast
        })

        return df, None
    
    except Exception as e:
        print(f"Error occurred: {e}")
        return None, f"An error occurred: {str(e)}"

# Gradio Interface Definition
gr.Interface(
    fn=predict_demand,
    inputs=[
        gr.Textbox(label="Mapped Code", placeholder="Enter mapped code here"),
        gr.Slider(minimum=1, maximum=12, step=1, label="Number of Months")
    ],
    outputs=[
        gr.Dataframe(label="Predicted Demand"),
        gr.Textbox(label="Error Message")
    ],
    title="Demand Forecasting",
    description="Enter the mapped code and the number of months to predict future demand."
).launch()