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 the next 3 periods for any numeric metric.
    - Saves a date‐indexed Plotly PNG under /tmp via the safe write monkey‐patch.
    - Returns a text table of the forecast.
    """

    # 0) Read full CSV
    df = pd.read_csv(file_path)

    # 1) Check that both columns actually exist
    if date_col not in df.columns:
        return f"❌ Date column '{date_col}' not found in your data."
    if value_col not in df.columns:
        return f"❌ Metric column '{value_col}' not found in your data."

    # 2) Parse dates
    try:
        df[date_col] = pd.to_datetime(df[date_col])
    except Exception:
        return f"❌ Could not parse '{date_col}' as dates."

    # 3) Coerce metric to numeric & drop invalid rows
    df[value_col] = pd.to_numeric(df[value_col], errors="coerce")
    df = df.dropna(subset=[date_col, value_col])
    if df.empty:
        return f"❌ After coercion, no valid data remains for '{value_col}'."

    # 4) Sort & index by date, collapse duplicates
    df = df.sort_values(date_col).set_index(date_col)
    df = df[[value_col]].groupby(level=0).mean()

    # 5) Infer a frequency and re‐index
    freq = pd.infer_freq(df.index)
    if freq is None:
        freq = "D"  # fallback to daily
    df = df.asfreq(freq)

    # 6) Fit ARIMA (1,1,1)
    try:
        model     = ARIMA(df[value_col], order=(1, 1, 1))
        model_fit = model.fit()
    except Exception as e:
        return f"❌ ARIMA fitting failed: {e}"

    # 7) Produce a proper date‐indexed forecast
    fc_res   = model_fit.get_forecast(steps=3)
    forecast = fc_res.predicted_mean

    # 8) Plot history + 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=date_col,
        yaxis_title=value_col,
        template="plotly_dark",
    )
    fig.write_image("forecast_plot.png")  # lands in /tmp via our monkey‐patch

    # 9) Return the forecast as a text table
    return forecast.to_frame(name="Forecast").to_string()