forecast.py

# WriterAgent - AI Writing Assistant for LibreOffice
# Copyright (c) 2026 KeithCu (modifications and relicensing)
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Trusted venv forecast compute — runs in user venv worker."""

from __future__ import annotations

import logging
from typing import Any

from plugin.scripting.calc_functions_common import (
    ANALYSIS_DATA_TABLE_ROWS as DATA_TABLE_ROWS,
    FORECAST_HELPER_NAMES as HELPER_NAMES,
)
from plugin.scripting.venv.coerce import (
    CoerceResult,
    ok_result as _ok_result,
    error_result as _error_result,
    parse_trusted_spec as _parse_trusted_spec,
    missing_package_error as _missing_package_error,
    resolve_df as _resolve_df,
    table_from_df as _table_from_df,
)

log = logging.getLogger(__name__)

_MIN_FORECAST_POINTS = 8
_MIN_DECOMPOSE_CYCLES = 2


def _require_statsmodels(helper: str) -> Any | None:
    try:
        import statsmodels  # noqa: F401

        return statsmodels
    except ImportError:
        return None


def _prepare_time_series(
    data: Any,
    *,
    date_col: str,
    value_col: str,
    headers: bool,
    header_row: int,
    sheet_hint: str | None,
    helper: str,
) -> tuple[CoerceResult | None, Any | None, dict[str, Any] | None]:
    """Return (coerced, series, error_dict)."""
    import pandas as pd

    coerced = _resolve_df(data, headers=headers, header_row=header_row, sheet_hint=sheet_hint)
    df = coerced.df
    if date_col not in df.columns:
        return None, None, _error_result("UNKNOWN_COLUMN", f"Column {date_col!r} not found", helper=helper)
    if value_col not in df.columns:
        return None, None, _error_result("UNKNOWN_COLUMN", f"Column {value_col!r} not found", helper=helper)

    work = df[[date_col, value_col]].copy()
    work[date_col] = pd.to_datetime(work[date_col], errors="coerce")
    work[value_col] = pd.to_numeric(work[value_col], errors="coerce")
    work = work.dropna(subset=[date_col, value_col]).sort_values(date_col)
    if work.empty:
        return None, None, _error_result("INSUFFICIENT_DATA", "No valid date/value rows after coercion", helper=helper)

    # STL and Holt-Winters require a unique index. Two sheet rows on the same
    # day used to be left as-is and then either error or distort the season.
    n_before = len(work)
    work = work.groupby(date_col, as_index=False, sort=True)[value_col].mean()
    collapsed = n_before - len(work)
    if collapsed:
        coerced.metadata["duplicate_dates_collapsed"] = int(collapsed)

    series = work.set_index(date_col)[value_col]
    if not series.index.is_monotonic_increasing:
        series = series.sort_index()
    if len(series) < _MIN_FORECAST_POINTS:
        return None, None, _error_result(
            "INSUFFICIENT_DATA",
            f"Need at least {_MIN_FORECAST_POINTS} observations; got {len(series)}",
            helper=helper,
        )
    return coerced, series, None


# One seasonal cycle for a regular calendar. Annual and sub-hour series have
# no cycle we can safely guess, so they stay trend-only unless the caller
# passes a period.
_SEASON_BY_FREQ_BASE: dict[str, int] = {
    "D": 7,
    "C": 7,
    "B": 5,
    "W": 52,
    "M": 12,
    "ME": 12,
    "MS": 12,
    "BM": 12,
    "BMS": 12,
    "BME": 12,
    "Q": 4,
    "QE": 4,
    "QS": 4,
    "BQ": 4,
    "BQS": 4,
    "H": 24,
    "BH": 24,
}


def _freq_base(freq: str) -> str:
    return str(freq).split("-", 1)[0].upper()


def _weekday_daily(series: Any, step: Any) -> bool:
    """True when the typical gap is one day and the index has no weekend."""
    import pandas as pd

    if not isinstance(step, pd.Timedelta):
        return False
    days = float(step.total_seconds() / 86400.0)
    if not 0.5 <= days <= 1.5:
        return False
    weekdays = {int(day) for day in series.index.dayofweek}
    return bool(weekdays) and weekdays.isdisjoint({5, 6})


def _period_from_gaps(series: Any) -> int | None:
    """Guess a cycle when ``infer_freq`` fails (weekend gaps, short spans)."""
    import pandas as pd

    deltas = series.index.to_series().diff().dropna()
    if deltas.empty:
        return None
    step = deltas.median()
    if not isinstance(step, pd.Timedelta):
        return None
    days = float(step.total_seconds() / 86400.0)
    if _weekday_daily(series, step):
        return 5
    if 0.5 <= days <= 1.5:
        return 7
    if 6 <= days <= 8:
        return 52
    if 27 <= days <= 32:
        return 12
    if 85 <= days <= 95:
        return 4
    return None


def _candidate_seasonal_period(series: Any) -> int | None:
    import pandas as pd

    freq = pd.infer_freq(series.index)
    if freq is not None:
        mapped = _SEASON_BY_FREQ_BASE.get(_freq_base(freq))
        if mapped is not None:
            return mapped
        # Known but non-seasonal (yearly, minutes): do not fall through to a
        # gap guess that might invent a cycle.
        return None
    return _period_from_gaps(series)


def _infer_seasonal_periods(series: Any, seasonal_periods: int | None) -> int | None:
    """Season length from the dates, or the caller's explicit period.

    The previous rule used the row count alone (``n >= 24`` → 12, ``n >= 14``
    → 7). A 30-row daily sheet was fit as a 12-day season, and future dates
    were still daily because those used ``infer_freq`` separately. An explicit
    period still wins. An inferred period is used only when the series covers
    at least two full cycles; otherwise the caller stays trend-only.
    """
    if seasonal_periods is not None and int(seasonal_periods) > 1:
        return int(seasonal_periods)
    candidate = _candidate_seasonal_period(series)
    if candidate is None:
        return None
    if len(series) < candidate * _MIN_DECOMPOSE_CYCLES:
        return None
    return candidate


def _seasonality_skip_reason(series: Any) -> str:
    candidate = _candidate_seasonal_period(series)
    if candidate is None:
        return "could not infer a seasonal period from the dates; used trend-only"
    need = candidate * _MIN_DECOMPOSE_CYCLES
    return (
        f"not enough cycles for a seasonal model "
        f"(need {need} observations for period {candidate}); used trend-only"
    )


def _duplicate_date_flags(coerced: CoerceResult | None) -> list[str]:
    if coerced is None:
        return []
    collapsed = coerced.metadata.get("duplicate_dates_collapsed")
    if not collapsed:
        return []
    return [f"Aggregated {int(collapsed)} duplicate date rows by mean"]


def _future_index(series: Any, periods: int) -> tuple[Any, str]:
    """Dates after the last observation, and the frequency label used.

    Weekday-only sheets often have no inferred frequency (the Friday–Monday
    gap breaks ``infer_freq``). The median gap is still one day, which used
    to emit Saturday and Sunday forecast dates.
    """
    import pandas as pd

    last_date = series.index[-1]
    freq = pd.infer_freq(series.index)
    if freq is not None:
        step = pd.tseries.frequencies.to_offset(freq)
        label = str(freq)
    else:
        deltas = series.index.to_series().diff().dropna()
        step = deltas.median() if not deltas.empty else pd.Timedelta(days=30)
        if _weekday_daily(series, step):
            step = pd.tseries.frequencies.to_offset("B")
            label = "B"
        else:
            label = "median_step"
    future = pd.date_range(start=last_date + step, periods=periods, freq=step)
    return future, label


def _forecast_moving_average(series: Any, *, periods: int) -> tuple[Any, str, dict[str, Any]]:
    import pandas as pd

    window = max(2, min(12, len(series) // 3))
    last_ma = float(series.rolling(window).mean().iloc[-1])
    future_dates, freq_label = _future_index(series, periods)
    forecast_df = pd.DataFrame({"date": future_dates, "forecast": [last_ma] * periods})
    metrics = {
        "model": "moving_average",
        "periods": periods,
        "n_obs": int(len(series)),
        "window": window,
        "freq": freq_label,
    }
    return forecast_df, "moving_average", metrics


# HoltWintersResults has forecast/predict/simulate and no get_prediction.
# get_prediction lives on ETSResults / statespace results (statsmodels
# tsa.holtwinters.results.HoltWintersResults). Calling it always raised
# AttributeError, which was swallowed as "confidence intervals unavailable".
_HW_SIM_REPETITIONS = 400
_HW_INTERVALS_UNAVAILABLE = (
    "Holt-Winters prediction intervals are omitted. "
    "HoltWintersResults has no analytic get_prediction; "
    "simulate() did not return a 95% band. The forecast column is the point forecast."
)


def _forecast_point(forecast_vals: Any, idx: int) -> float:
    if hasattr(forecast_vals, "iloc"):
        return float(forecast_vals.iloc[idx])
    return float(forecast_vals[idx])


def _holt_winters_intervals(fit: Any, periods: int) -> tuple[Any, Any] | None:
    """95% band from ``HoltWintersResults.simulate`` state-space paths.

    Returns ``(lower, upper)`` aligned by position with the horizon, or
    ``None`` when simulation cannot build a band. Callers must say so in
    ``interval_note`` instead of inventing lower/upper.
    """
    import numpy as np
    import pandas as pd

    try:
        sims = fit.simulate(
            periods,
            repetitions=_HW_SIM_REPETITIONS,
            error="add",
            rng=np.random.default_rng(0),
        )
    except TypeError:
        # statsmodels before 0.14 named this argument random_state.
        try:
            sims = fit.simulate(
                periods,
                repetitions=_HW_SIM_REPETITIONS,
                error="add",
                random_state=0,
            )
        except Exception:
            log.exception("Holt-Winters interval simulation failed")
            return None
    except Exception:
        log.exception("Holt-Winters interval simulation failed")
        return None

    if not isinstance(sims, pd.DataFrame) or len(sims) != periods or sims.shape[1] < 2:
        return None
    lower = sims.quantile(0.025, axis=1)
    upper = sims.quantile(0.975, axis=1)
    if len(lower) != periods or len(upper) != periods:
        return None
    if not np.isfinite(lower.to_numpy()).all() or not np.isfinite(upper.to_numpy()).all():
        return None
    return lower, upper


def _forecast_holt_winters(series: Any, *, periods: int, seasonal_periods: int) -> tuple[Any, str, dict[str, Any], list[str]]:
    import pandas as pd
    from statsmodels.tsa.holtwinters import ExponentialSmoothing

    model = ExponentialSmoothing(
        series,
        trend="add",
        seasonal="add",
        seasonal_periods=seasonal_periods,
    )
    fit = model.fit(optimized=True)
    forecast_vals = fit.forecast(periods)
    future_dates, freq_label = _future_index(series, periods)

    rows: list[dict[str, Any]] = []
    flags: list[str] = []
    intervals = _holt_winters_intervals(fit, periods)
    if intervals is None:
        flags.append(_HW_INTERVALS_UNAVAILABLE)
        for idx, dt in enumerate(future_dates):
            rows.append({"date": dt, "forecast": _forecast_point(forecast_vals, idx)})
    else:
        lower, upper = intervals
        for idx, dt in enumerate(future_dates):
            rows.append({
                "date": dt,
                "forecast": _forecast_point(forecast_vals, idx),
                "lower": float(lower.iloc[idx]),
                "upper": float(upper.iloc[idx]),
            })

    forecast_df = pd.DataFrame(rows)
    metrics: dict[str, Any] = {
        "model": "holt_winters",
        "periods": periods,
        "n_obs": int(len(series)),
        "seasonal_periods": seasonal_periods,
        "freq": freq_label,
    }
    if hasattr(fit, "aic"):
        metrics["aic"] = float(fit.aic)
    if hasattr(fit, "sse"):
        metrics["sse"] = float(fit.sse)
    if intervals is None:
        metrics["interval_note"] = {
            "available": False,
            "method": "HoltWintersResults.simulate",
            "message": _HW_INTERVALS_UNAVAILABLE,
        }
    return forecast_df, "holt_winters", metrics, flags


def _forecast_arima(series: Any, *, periods: int) -> tuple[Any, str, dict[str, Any], list[str]]:
    import pandas as pd
    from statsmodels.tsa.arima.model import ARIMA

    model = ARIMA(series, order=(1, 1, 1))
    fit = model.fit()
    pred = fit.get_forecast(steps=periods)
    forecast_vals = pred.predicted_mean
    conf = pred.conf_int()
    future_dates, freq_label = _future_index(series, periods)

    rows = []
    for idx, dt in enumerate(future_dates):
        row = {"date": dt, "forecast": float(forecast_vals.iloc[idx])}
        if conf is not None and len(conf) > idx:
            row["lower"] = float(conf.iloc[idx, 0])
            row["upper"] = float(conf.iloc[idx, 1])
        rows.append(row)
    forecast_df = pd.DataFrame(rows)
    metrics: dict[str, Any] = {
        "model": "arima",
        "periods": periods,
        "n_obs": int(len(series)),
        "order": "(1,1,1)",
        "freq": freq_label,
    }
    if hasattr(fit, "aic"):
        metrics["aic"] = float(fit.aic)
    return forecast_df, "arima", metrics, []


def forecast_time_series(
    data: Any,
    *,
    periods: int = 12,
    model: str = "auto",
    date_col: str = "Date",
    value_col: str = "Value",
    seasonal_periods: int | None = None,
    headers: bool = True,
    header_row: int = 0,
    sheet_hint: str | None = None,
) -> dict[str, Any]:
    """Forward predictions on a date-indexed series."""
    helper = "forecast_time_series"
    coerced, series, err = _prepare_time_series(
        data,
        date_col=date_col,
        value_col=value_col,
        headers=headers,
        header_row=header_row,
        sheet_hint=sheet_hint,
        helper=helper,
    )
    if err is not None:
        return err
    if series is None:
        return _error_result("INSUFFICIENT_DATA", "No valid time series", helper=helper)

    # A huge horizon used to be passed straight to statsmodels. Same idea as
    # monte_carlo capping sims at 1_000_000.
    periods = max(1, min(int(periods), 10_000))
    model_name = str(model or "auto").strip().lower()
    flags: list[str] = _duplicate_date_flags(coerced)

    if model_name == "moving_average":
        forecast_df, used_model, metrics = _forecast_moving_average(series, periods=periods)
    else:
        if _require_statsmodels(helper) is None and model_name != "moving_average":
            if model_name == "auto":
                forecast_df, used_model, metrics = _forecast_moving_average(series, periods=periods)
                flags.append("statsmodels missing; used moving_average fallback")
            else:
                return _missing_package_error(helper, "statsmodels")
        else:
            explicit_season = seasonal_periods is not None and int(seasonal_periods) > 1
            season = _infer_seasonal_periods(series, seasonal_periods)
            if season is None and model_name == "auto" and not explicit_season:
                flags.append(_seasonality_skip_reason(series))
            try:
                if model_name in ("auto", "holt_winters") and season is not None and len(series) >= season * _MIN_DECOMPOSE_CYCLES:
                    forecast_df, used_model, metrics, hw_flags = _forecast_holt_winters(series, periods=periods, seasonal_periods=season)
                    flags.extend(hw_flags)
                elif model_name in ("auto", "arima"):
                    forecast_df, used_model, metrics, arima_flags = _forecast_arima(series, periods=periods)
                    flags.extend(arima_flags)
                elif model_name == "holt_winters":
                    if season is None:
                        return _error_result(
                            "INSUFFICIENT_DATA",
                            "holt_winters requires seasonal_periods or enough data to infer seasonality",
                            helper=helper,
                        )
                    forecast_df, used_model, metrics, hw_flags = _forecast_holt_winters(series, periods=periods, seasonal_periods=season)
                    flags.extend(hw_flags)
                else:
                    return _error_result("FORECAST_FAILED", f"Unknown model {model_name!r}", helper=helper)
            except Exception as exc:
                if model_name == "auto":
                    try:
                        forecast_df, used_model, metrics, arima_flags = _forecast_arima(series, periods=periods)
                        flags.extend(arima_flags)
                        flags.append(f"primary model failed ({exc}); used arima")
                    except Exception:
                        forecast_df, used_model, metrics = _forecast_moving_average(series, periods=periods)
                        flags.append(f"statsmodels forecast failed ({exc}); used moving_average")
                else:
                    return _error_result("FORECAST_FAILED", str(exc), helper=helper)

    metrics.setdefault("model", used_model)
    table = _table_from_df(forecast_df, name="forecast")
    return _ok_result(helper, metrics=metrics, tables=[table], flags=flags, metadata=coerced.metadata if coerced else {})


def decompose_time_series(
    data: Any,
    *,
    date_col: str = "Date",
    value_col: str = "Value",
    model: str = "additive",
    period: int | None = None,
    headers: bool = True,
    header_row: int = 0,
    sheet_hint: str | None = None,
) -> dict[str, Any]:
    """Trend / seasonal / residual decomposition via statsmodels."""
    import pandas as pd

    helper = "decompose_time_series"
    if _require_statsmodels(helper) is None:
        return _missing_package_error(helper, "statsmodels")

    coerced, series, err = _prepare_time_series(
        data,
        date_col=date_col,
        value_col=value_col,
        headers=headers,
        header_row=header_row,
        sheet_hint=sheet_hint,
        helper=helper,
    )
    if err is not None:
        return err
    if series is None:
        return _error_result("INSUFFICIENT_DATA", "No valid time series", helper=helper)

    decomp_model = str(model or "additive").strip().lower()
    if decomp_model not in ("additive", "multiplicative"):
        return _error_result("FORECAST_FAILED", f"model must be additive or multiplicative, got {decomp_model!r}", helper=helper)

    season = period if period is not None else _infer_seasonal_periods(series, None)
    if season is None or season < 2:
        return _error_result(
            "INSUFFICIENT_DATA",
            "decompose_time_series needs an explicit period, or a date frequency with at least two full seasonal cycles",
            helper=helper,
        )
    if len(series) < season * _MIN_DECOMPOSE_CYCLES:
        return _error_result(
            "INSUFFICIENT_DATA",
            f"Need at least {season * _MIN_DECOMPOSE_CYCLES} observations for period={season}; got {len(series)}",
            helper=helper,
        )

    try:
        from statsmodels.tsa.seasonal import seasonal_decompose

        result = seasonal_decompose(series, model=decomp_model, period=season)
    except Exception as exc:
        return _error_result("FORECAST_FAILED", str(exc), helper=helper)

    decomp_df = pd.DataFrame(
        {
            "date": series.index,
            "observed": series.values,
            "trend": result.trend,
            "seasonal": result.seasonal,
            "resid": result.resid,
        }
    )
    table = _table_from_df(decomp_df, name="decomposition", max_rows=DATA_TABLE_ROWS)
    metrics = {
        "model": decomp_model,
        "period": season,
        "n_obs": int(len(series)),
    }
    return _ok_result(
        helper,
        metrics=metrics,
        tables=[table],
        flags=_duplicate_date_flags(coerced),
        metadata=coerced.metadata if coerced else {},
    )


def _robust_z_scores(values: Any) -> Any:
    """MAD-based robust z-scores; fall back to std when MAD is zero."""
    import numpy as np

    arr = np.asarray(values, dtype=float)
    median = float(np.median(arr))
    mad = float(np.median(np.abs(arr - median)))
    if mad > 0:
        return 0.6745 * (arr - median) / mad
    std = float(np.std(arr))
    if std > 0:
        return (arr - float(np.mean(arr))) / std
    return np.zeros_like(arr)


def anomaly_detection_time_series(
    data: Any,
    *,
    date_col: str = "Date",
    value_col: str = "Value",
    period: int | None = None,
    method: str = "stl_residual",
    threshold: float = 3.0,
    include_all: bool = False,
    headers: bool = True,
    header_row: int = 0,
    sheet_hint: str | None = None,
) -> dict[str, Any]:
    """Flag temporal outliers via STL residuals and robust z-scores."""
    import pandas as pd

    helper = "anomaly_detection_time_series"
    if _require_statsmodels(helper) is None:
        return _missing_package_error(helper, "statsmodels")

    method_name = str(method or "stl_residual").strip().lower()
    if method_name != "stl_residual":
        return _error_result("FORECAST_FAILED", f"Unknown method {method_name!r}; only stl_residual is supported", helper=helper)

    coerced, series, err = _prepare_time_series(
        data,
        date_col=date_col,
        value_col=value_col,
        headers=headers,
        header_row=header_row,
        sheet_hint=sheet_hint,
        helper=helper,
    )
    if err is not None:
        return err
    if series is None:
        return _error_result("INSUFFICIENT_DATA", "No valid time series", helper=helper)

    season = period if period is not None else _infer_seasonal_periods(series, None)
    if season is None or season < 2:
        return _error_result(
            "INSUFFICIENT_DATA",
            "anomaly_detection_time_series needs an explicit period, or a date frequency with at least two full seasonal cycles",
            helper=helper,
        )
    if len(series) < season * _MIN_DECOMPOSE_CYCLES:
        return _error_result(
            "INSUFFICIENT_DATA",
            f"Need at least {season * _MIN_DECOMPOSE_CYCLES} observations for period={season}; got {len(series)}",
            helper=helper,
        )

    try:
        from statsmodels.tsa.seasonal import STL

        stl_result = STL(series, period=season, robust=True).fit()
    except Exception as exc:
        return _error_result("FORECAST_FAILED", str(exc), helper=helper)

    expected = stl_result.trend + stl_result.seasonal
    resid = stl_result.resid
    scores = _robust_z_scores(resid.values)
    threshold_val = float(threshold)

    scores_df = pd.DataFrame(
        {
            "date": series.index,
            "observed": series.values,
            "expected": expected.values,
            "residual": resid.values,
            "score": scores,
        }
    )
    anomalies_df = scores_df[scores_df["score"].abs() > threshold_val].copy()
    tables = [_table_from_df(anomalies_df, name="anomalies")]
    if include_all:
        tables.append(_table_from_df(scores_df, name="all_scores"))

    metrics = {
        "n_anomalies": int(len(anomalies_df)),
        "period": season,
        "threshold": threshold_val,
        "method": method_name,
        "n_obs": int(len(series)),
    }
    return _ok_result(
        helper,
        metrics=metrics,
        tables=tables,
        flags=_duplicate_date_flags(coerced),
        metadata=coerced.metadata if coerced else {},
    )


def _dispatch_helper(name: str, data: Any, params: dict[str, Any], *, headers: bool, header_row: int, context: dict[str, Any]) -> dict[str, Any]:
    sheet_hint = context.get("sheet_name") if isinstance(context.get("sheet_name"), str) else None
    common: dict[str, Any] = {"headers": headers, "header_row": header_row, "sheet_hint": sheet_hint}

    if name == "forecast_time_series":
        return forecast_time_series(
            data,
            periods=int(params.get("periods", 12)),
            model=str(params.get("model", "auto")),
            date_col=str(params.get("date_col", "Date")),
            value_col=str(params.get("value_col", "Value")),
            seasonal_periods=params.get("seasonal_periods"),
            **common,
        )
    if name == "decompose_time_series":
        return decompose_time_series(
            data,
            date_col=str(params.get("date_col", "Date")),
            value_col=str(params.get("value_col", "Value")),
            model=str(params.get("model", "additive")),
            period=params.get("period"),
            **common,
        )
    if name == "anomaly_detection_time_series":
        return anomaly_detection_time_series(
            data,
            date_col=str(params.get("date_col", "Date")),
            value_col=str(params.get("value_col", "Value")),
            period=params.get("period"),
            method=str(params.get("method", "stl_residual")),
            threshold=float(params.get("threshold", 3.0)),
            include_all=bool(params.get("include_all", False)),
            **common,
        )
    return _error_result("UNKNOWN_HELPER", f"Forecast helper {name!r} not found", helper=name)


def run_forecast(
    spec: dict[str, Any] | str,
    data: Any,
    context: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Spec-driven dispatcher — single trusted entry for host RPC."""
    parsed = _parse_trusted_spec(spec, helper_names=HELPER_NAMES, context=context)
    if isinstance(parsed, dict):
        return parsed
    helper, params, headers, header_row, ctx, _spec = parsed

    try:
        result = _dispatch_helper(helper, data, params, headers=headers, header_row=header_row, context=ctx)
    except Exception as exc:
        log.exception("Forecast helper %s failed", helper)
        return _error_result("FORECAST_FAILED", str(exc), helper=helper)

    if isinstance(result, dict) and result.get("status") == "ok" and ctx:
        result["context"] = {k: v for k, v in ctx.items() if k in ("sheet_name", "range_a1", "task_hint")}

    return result