tools/forecaster.py

# tools/forecaster.py

import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
import plotly.graph_objects as go

def forecast_metric_tool(file_path: str, date_col: str, value_col: str):
    """
    Forecast next 3 periods for any numeric metric, and return a textual table.
    Saves a date-indexed PNG under /tmp via our safe monkey-patch.
    """
    # 1) Load & parse
    df = pd.read_csv(file_path)
    try:
        df[date_col] = pd.to_datetime(df[date_col])
    except Exception:
        return f"❌ Could not parse '{date_col}' as dates."
    
    # 2) Coerce metric to numeric & drop invalid
    df[value_col] = pd.to_numeric(df[value_col], errors="coerce")
    df = df.dropna(subset=[date_col, value_col])
    
    if df.empty:
        return f"❌ No valid data for '{value_col}'."
    
    # 3) Sort and set index, infer frequency
    df = df.sort_values(date_col)
    df.set_index(date_col, inplace=True)
    freq = pd.infer_freq(df.index)
    if freq is None:
        # fallback to daily if pandas can't infer
        freq = "D"
    df = df.asfreq(freq)
    
    # 4) Fit ARIMA
    try:
        model     = ARIMA(df[value_col], order=(1, 1, 1))
        model_fit = model.fit()
    except Exception as e:
        return f"❌ ARIMA fitting failed: {e}"
    
    # 5) Get a proper date-indexed forecast
    fc_res   = model_fit.get_forecast(steps=3)
    forecast = fc_res.predicted_mean  # a pd.Series with a DatetimeIndex
    
    # 6) Plot historical + forecast
    fig = go.Figure()
    fig.add_scatter(
        x=df.index, y=df[value_col],
        mode="lines", name=value_col
    )
    fig.add_scatter(
        x=forecast.index, y=forecast,
        mode="lines+markers", name="Forecast"
    )
    fig.update_layout(
        title=f"{value_col} Forecast",
        xaxis_title=str(date_col),
        yaxis_title=str(value_col),
        template="plotly_dark"
    )
    fig.write_image("forecast_plot.png")
    
    # 7) Return the forecast table as text
    tbl = forecast.to_frame(name="Forecast")
    return tbl.to_string()