src/tools/analytics/forecast.py

"""analyze_forecast — project a time series forward with exponential smoothing (§0.9).

First of the ML analytical tools. Buckets rows into periods (reusing analyze_trend's
date parsing + resampling), fits a statsmodels ETS (Holt-Winters, additive trend, no
AutoML — locked decision) OFF the event loop via `run_fit`, and returns the history +
forecast series plus an actual-vs-forecast line chart ("sepaket").

Follows docs/tools/ANALYTICAL_TOOLS_CONVENTIONS.md: pure-ish compute function (async
ONLY to `await run_fit` for the fit), `data` never self-fetched (Pattern A), never
catches its own errors (the invoker's never-throw seam wraps them), deterministic
output (ETS optimisation is deterministic — no seed needed).
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from pydantic import BaseModel, Field

from src.tools.analytics.agg_vocab import normalise_agg
from src.tools.analytics.charts import (
    _LINE_STYLE,
    COLORWAY,
    LAYOUT_PRESET,
    _axis,
    envelope,
)
from src.tools.analytics.descriptive import ColumnNotFoundError
from src.tools.analytics.temporal import (
    FREQ_MAP,
    SUPPORTED_AGGS,
    InvalidFrequencyError,
    UnsupportedAggregationError,
    _clean,
    _parse_date_column,
    _period_label,
)
from src.tools.ml.executor import run_fit
from src.tools.spec_builder import build_description

# ETS with an additive trend and estimated initialisation needs a few points to
# be meaningful; below this we raise rather than emit a garbage projection.
_MIN_POINTS = 4
# A holdout backtest needs points to train on AND points to score against; below this
# the forecast ships UNVALIDATED, which is itself weak evidence (see `quality`).
_MIN_BACKTEST_POINTS = 8
# MASE >= 1 means the model did no better than "next = last" on the holdout — the
# floor the reserved `quality` block reports and the checkpoint's CK7 reads.
_MASE_FLOOR = 1.0


class NotEnoughDataError(ValueError):
    """Too few non-empty periods to fit a forecast (error_code NOT_ENOUGH_DATA)."""


class InvalidHorizonError(ValueError):
    """horizon is not a positive integer (error_code INVALID_HORIZON)."""


class ForecastInput(BaseModel):
    # `data` FIRST — property order follows field order (planner-prompt readability).
    data: str = Field(..., description="Placeholder ${t<id>} of the upstream table (Pattern A).")
    date_column: str = Field(..., description="Column holding dates/timestamps to bucket on.")
    value_column: str = Field(..., description="Numeric column to project into the future.")
    horizon: int = Field(6, description="How many future periods to forecast (>= 1).")
    freq: str = Field("month", description="Period granularity: day, week, month, quarter, year.")
    agg: str = Field("sum", description="How to aggregate the value within each period.")


DESCRIPTION = build_description(
    summary="Project a time series forward N periods with an exponential-smoothing (ETS) model.",
    use_when="the user wants a projection/forecast of a metric into the FUTURE "
    "(next months, next quarter) — 'proyeksi', 'forecast', 'prediksi ke depan'.",
    dont_use_when=[
        "the user wants PAST movement/trend, not a future projection -> analyze_trend",
        "there is no date/time column to project along -> analyze_aggregate / analyze_descriptive",
        "the user wants what DRIVES the metric, not its future values -> analyze_importance",
    ],
    output="a history + forecast series (per-period points) plus an actual-vs-forecast line chart.",
    examples=[
        "proyeksi produksi 6 bulan ke depan",
        "forecast next quarter's revenue",
        "predict signups for the next 12 weeks",
    ],
)


def _ets(values: np.ndarray):
    """An additive-trend ETS fit. Imported lazily so app startup doesn't pay for
    statsmodels unless a forecast actually runs."""
    from statsmodels.tsa.holtwinters import ExponentialSmoothing

    return ExponentialSmoothing(
        values, trend="add", seasonal=None, initialization_method="estimated"
    ).fit()


def _fit_ets(values: np.ndarray, horizon: int):
    """Fit the reported model AND backtest it. Runs inside run_fit's thread.

    Both fits happen in one offloaded call so a forecast costs one thread hop and one
    timeout, not two.

    The backtest is a single holdout: train on everything but the last `h` periods,
    forecast those, and score with **MASE** — mean absolute error divided by the mean
    absolute error of the naive "next = last" forecast on the training span. MASE is
    the scale-free choice (Hyndman & Koehler 2006) and, unlike MAPE, it neither
    explodes near zero nor punishes over-prediction asymmetrically. It also reads
    plainly: **below 1 beats doing nothing, at or above 1 does not.**

    Returns (model, mase, reason). `mase` is None when no backtest was possible, and
    `reason` then says WHY in the user's terms — the two causes are different and a
    wrong explanation is worse than none.
    """
    model = _ets(values)

    n = values.shape[0]
    too_short = "too few periods to hold any out"
    if n < _MIN_BACKTEST_POINTS:
        return model, None, too_short
    # Hold out the forecast horizon, but never more than a quarter of the series and
    # never so much that training drops below the fitting floor.
    holdout = max(1, min(horizon, n // 4, n - _MIN_POINTS))
    train, test = values[:-holdout], values[-holdout:]
    if train.shape[0] < _MIN_POINTS:
        return model, None, too_short

    predicted = np.asarray(_ets(train).forecast(holdout), dtype=float)
    naive_error = float(np.mean(np.abs(np.diff(train))))
    if naive_error == 0:
        # The series never moves, so the naive baseline has zero error and MASE is
        # undefined (not infinite). A forecast of a constant is trivially "right" and
        # tells the user nothing — still weak, but for this reason, not shortness.
        return model, None, "the series never changes, so there is nothing to validate against"
    return model, float(np.mean(np.abs(test - predicted)) / naive_error), None


def _forecast_chart(
    history: list[dict[str, object]],
    forecast_pts: list[dict[str, object]],
    value_column: str,
) -> dict[str, object]:
    """Line chart: actual history + a dashed forecast, connected at the last actual."""
    hist_x = [p["period"] for p in history]
    hist_y = [p["value"] for p in history]

    # Anchor the forecast line to the last real actual so the dashed segment connects.
    anchor = next((p for p in reversed(history) if p["value"] is not None), None)
    fc_x = ([anchor["period"]] if anchor else []) + [p["period"] for p in forecast_pts]
    fc_y = ([anchor["value"]] if anchor else []) + [p["value"] for p in forecast_pts]

    actual = {
        "type": "scatter",
        "mode": "lines+markers",
        "x": hist_x,
        "y": hist_y,
        "name": "actual",
        "line": {**_LINE_STYLE, "color": COLORWAY[0]},
        "marker": {"size": 6},
    }
    forecast = {
        "type": "scatter",
        "mode": "lines+markers",
        "x": fc_x,
        "y": fc_y,
        "name": "forecast",
        "line": {**_LINE_STYLE, "color": COLORWAY[5], "dash": "dash"},
        "marker": {"size": 6},
    }
    title = f"Forecast of {value_column}"
    layout = {
        **LAYOUT_PRESET,
        "title": {"text": title},
        "showlegend": True,
        "xaxis": _axis("period"),
        "yaxis": _axis(value_column),
    }
    return envelope("line", title, [actual, forecast], layout)


async def analyze_forecast(
    df: pd.DataFrame,
    date_column: str,
    value_column: str,
    horizon: int = 6,
    freq: str = "month",
    agg: str = "sum",
) -> dict[str, object]:
    """Forecast `value_column` `horizon` periods ahead over `freq` buckets.

    Async ONLY so the ETS fit can be offloaded via `run_fit`; everything else is
    ordinary synchronous compute. `data` is materialised upstream (Pattern A) and
    arrives as `df`.

    Returns a dict with: freq, agg, horizon, model, history + forecast point lists,
    last_actual / forecast_end / change_abs / change_pct / direction, and `charts`
    (the reserved auto-chart key).

    Raises (all wrapped by the invoker's never-throw seam):
        ColumnNotFoundError, InvalidFrequencyError, UnsupportedAggregationError,
        InvalidHorizonError, NotEnoughDataError, InvalidDateColumnError.
    """
    missing = [c for c in (date_column, value_column) if c not in df.columns]
    if missing:
        raise ColumnNotFoundError(f"columns not found: {missing}")
    if freq not in FREQ_MAP:
        raise InvalidFrequencyError(f"unknown frequency '{freq}'; supported: {list(FREQ_MAP)}")
    # `avg` is the QueryIR's word for this and planner.md teaches it; the tools
    # speak pandas. Accept both, and report whatever the caller actually wrote.
    # (A23 Bug 2 -- see agg_vocab.)
    requested, agg = agg, normalise_agg(agg)
    if agg not in SUPPORTED_AGGS:
        raise UnsupportedAggregationError(
            f"unsupported aggregation '{requested}'; supported: {list(SUPPORTED_AGGS)}"
        )
    if not isinstance(horizon, int) or horizon < 1:
        raise InvalidHorizonError(f"horizon must be a positive integer, got {horizon!r}")

    # Build a datetime-indexed series, then resample into evenly-spaced periods.
    # Same positional build as analyze_trend so date_column == value_column works.
    dates = pd.Series(_parse_date_column(df, date_column).to_numpy(), name="_date")
    values = pd.Series(df[value_column].to_numpy(), name="_value")
    frame = pd.concat([dates, values], axis=1).dropna(subset=["_date"]).set_index("_date")
    resampled = frame["_value"].sort_index().resample(FREQ_MAP[freq]).agg(agg).astype(float)

    observed = resampled.dropna()
    if observed.shape[0] < _MIN_POINTS:
        raise NotEnoughDataError(
            f"need >= {_MIN_POINTS} non-empty {freq} periods to forecast, "
            f"got {observed.shape[0]}"
        )
    # ETS needs a gap-free series; interpolate the empty interior periods.
    filled = resampled.interpolate(limit_direction="both")

    model, mase, mase_reason = await run_fit(
        _fit_ets, filled.to_numpy(dtype=float), horizon, label="analyze_forecast"
    )
    fc_values = [float(v) for v in np.asarray(model.forecast(horizon), dtype=float)]

    future_index = pd.date_range(
        start=filled.index[-1], periods=horizon + 1, freq=FREQ_MAP[freq]
    )[1:]
    history = [
        {"period": _period_label(ts, freq), "value": _clean(val)} for ts, val in resampled.items()
    ]
    forecast_pts = [
        {"period": _period_label(ts, freq), "value": val}
        for ts, val in zip(future_index, fc_values, strict=True)
    ]

    last_actual = float(observed.iloc[-1])
    forecast_end = fc_values[-1] if fc_values else None
    change_abs = (forecast_end - last_actual) if forecast_end is not None else None
    change_pct = (
        (change_abs / last_actual) if (change_abs is not None and last_actual != 0) else None
    )
    if change_abs is None or change_abs == 0:
        direction = "flat"
    else:
        direction = "up" if change_abs > 0 else "down"

    return {
        "freq": freq,
        "agg": agg,
        "horizon": horizon,
        "model": "ETS(add trend)",
        "history": history,
        "forecast": forecast_pts,
        "last_actual": last_actual,
        "forecast_end": forecast_end,
        "change_abs": change_abs,
        "change_pct": change_pct,
        "direction": direction,
        # Reserved `quality` block (conventions §3): the tool reports its own score and
        # its own verdict; CK7 surfaces `weak`. An unvalidated forecast (series too
        # short to hold anything out) counts as weak — shipping a projection nobody
        # checked is exactly the case the plan's R-3 wanted flagged.
        "quality": {
            "metric": "mase",
            "value": None if mase is None else round(mase, 4),
            "floor": _MASE_FLOOR,
            "weak": mase is None or mase >= _MASE_FLOOR,
            "note": (
                f"not validated — {mase_reason}"
                if mase is None
                else "did not beat a naive 'next = last' forecast on the holdout"
                if mase >= _MASE_FLOOR
                else "beats a naive 'next = last' forecast on the holdout"
            ),
        },
        "charts": [_forecast_chart(history, forecast_pts, value_column)],
    }