gradio_app.py

"""
Gradio app to load a CSV file and forecast upcoming monthly budgets using a trained model.
Usage:
  python gradio_app.py
Then open the local URL printed by Gradio, upload a CSV, and click forecast.
"""

import gradio as gr
import pandas as pd
from pathlib import Path
from typing import Tuple, Dict, Any

from budget_forecasting_real_data import (
    clean_and_preprocess,
    predict_future_budgets,
    load_model,
)

MODEL_PATH = Path("best_model_linear_regression.joblib")


def load_model_bundle() -> Tuple[Dict[str, Any], str]:
    """Load the persisted model bundle."""
    if not MODEL_PATH.exists():
        return {}, f"Missing model bundle at {MODEL_PATH}"
    
    try:
        bundle = load_model(MODEL_PATH)
        return bundle, ""
    except Exception as exc:
        return {}, f"Failed to load model: {exc}"


MODEL_BUNDLE, MODEL_ERROR = load_model_bundle()


def forecast_with_csv(csv_file, n_months: int) -> pd.DataFrame:
    """
    Process uploaded CSV and forecast future months.
    
    Args:
        csv_file: Uploaded file (Gradio returns file path as string)
        n_months: Number of months to forecast
    
    Returns:
        DataFrame with forecasts or error message
    """
    if MODEL_ERROR:
        return pd.DataFrame({"error": [MODEL_ERROR]})
    
    if csv_file is None:
        return pd.DataFrame({"error": ["Please upload a CSV file"]})
    
    try:
        # csv_file is a path string when uploaded via Gradio
        df_raw = pd.read_csv(csv_file)
        df_processed = clean_and_preprocess(df_raw)
        
        # Extract model and scaler from bundle
        best_result = {
            "model": MODEL_BUNDLE["model"],
            "scaler": MODEL_BUNDLE["scaler"],
        }
        feature_cols = MODEL_BUNDLE.get("feature_columns", [])
        
        # Forecast
        n = max(1, min(int(n_months), 24))  # clamp to 1..24
        future_df = predict_future_budgets(
            df_processed, best_result, feature_cols, n_future_months=n
        )
        return future_df
    except Exception as exc:
        return pd.DataFrame({"error": [str(exc)]})


def build_interface():
    with gr.Blocks(title="Budget Forecasting Console") as demo:
        gr.Markdown(
            """
            # Budget Forecasting Console
            Upload a CSV with monthly budget data and forecast upcoming months using a trained Linear Regression model.
            
            **CSV Format Required:**
            - Must contain columns: `month` (YYYY-MM format) and `monthly_budget_pkr` (numeric)
            - Example: 2026-01, 5975.77
            
            **How it works:**
            1. Upload your CSV file
            2. Select forecast horizon (1-24 months)
            3. Click "Run forecast" to see predictions
            """
        )

        with gr.Row():
            csv_upload = gr.File(
                label="Upload CSV",
                file_types=[".csv"],
                type="filepath"
            )

        with gr.Row():
            n_slider = gr.Slider(
                minimum=1,
                maximum=24,
                value=1,
                step=1,
                label="Months to forecast",
                info="Forecast horizon (months ahead)",
            )
            run_btn = gr.Button("Run forecast", variant="primary")

        output_df = gr.Dataframe(
            headers=["month", "predicted_monthly_budget_pkr"],
            datatype=["str", "number"],
            label="Forecasts",
            interactive=False,
        )

        run_btn.click(
            forecast_with_csv,
            inputs=[csv_upload, n_slider],
            outputs=output_df
        )

        if MODEL_ERROR:
            gr.Markdown(f"⚠️ **Model Load Error:** {MODEL_ERROR}")

    return demo


def main():
    demo = build_interface()
    demo.launch()


if __name__ == "__main__":
    main()