"""SARIMA forecast for BBCA closing prices.

The order comes from the notebook's auto_arima search:
SARIMA(0,1,0)(2,1,0)[5], where 5 is one trading week.
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass

import numpy as np
import pandas as pd
from statsmodels.tsa.statespace.sarimax import SARIMAX

ORDER = (0, 1, 0)
SEASONAL_ORDER = (2, 1, 0, 5)
FIT_WINDOW = 250  # about one year of trading days
MAX_HORIZON = 10


@dataclass(frozen=True)
class Forecast:
    dates: pd.DatetimeIndex
    mean: np.ndarray
    lower: np.ndarray
    upper: np.ndarray
    base_date: pd.Timestamp  # last observation the forecast starts from
    base_price: float
    base_is_live: bool
    fit_points: int


def fit(values: np.ndarray):
    """Fit SARIMA on a plain array (no date index, so no frequency warnings)."""
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        model = SARIMAX(
            np.asarray(values, dtype=float),
            order=ORDER,
            seasonal_order=SEASONAL_ORDER,
            enforce_stationarity=False,
            enforce_invertibility=False,
        )
        return model.fit(disp=False)


def next_trading_days(after: pd.Timestamp, n: int) -> pd.DatetimeIndex:
    """Weekdays after `after`. IDX holidays are not removed."""
    return pd.bdate_range(after + pd.Timedelta(days=1), periods=n)


def forecast(
    closes: pd.Series,
    horizon: int,
    live_price: float | None = None,
    live_date: pd.Timestamp | None = None,
) -> Forecast:
    horizon = int(min(max(horizon, 1), MAX_HORIZON))
    series = closes.iloc[-FIT_WINDOW:]
    values = series.to_numpy(dtype=float)
    base_date = series.index[-1]
    base_is_live = False
    if live_price is not None and live_date is not None:
        # Treat the running price as today's close. The forecast then moves
        # with the market during trading hours.
        values = np.append(values, live_price)
        base_date = live_date
        base_is_live = True

    result = fit(values)
    frame = result.get_forecast(steps=horizon).summary_frame(alpha=0.05)
    return Forecast(
        dates=next_trading_days(base_date, horizon),
        mean=frame["mean"].to_numpy(),
        lower=frame["mean_ci_lower"].to_numpy(),
        upper=frame["mean_ci_upper"].to_numpy(),
        base_date=base_date,
        base_price=float(values[-1]),
        base_is_live=base_is_live,
        fit_points=len(values),
    )