"""Reference forecasts every other model has to beat."""

from dataclasses import dataclass
from typing import ClassVar

import numpy as np
import pandas as pd

from ..data import log_returns
from ..forecast import Forecast


@dataclass(frozen=True)
class RandomWalk:
    """
    Zero-drift random walk in log price: "it stays about here, give or take".

    Volatility is the trailing standard deviation of daily log returns.
    """

    name: ClassVar[str] = "random_walk"
    vol_window: int = 365

    def forecast(self, history: pd.DataFrame, horizons: np.ndarray) -> Forecast:
        sigma = log_returns(history).iloc[-self.vol_window :].std()
        mean = np.log(history["close"].iloc[-1])
        return Forecast.normal(history.index[-1], horizons, mean, sigma * np.sqrt(horizons))


@dataclass(frozen=True)
class DriftRandomWalk:
    """
    Random walk whose drift is the mean daily log return over the trailing
    `drift_window` days (one halving cycle by default): "it keeps doing what it
    did last cycle".
    """

    name: ClassVar[str] = "drift_rw"
    drift_window: int = 1460
    vol_window: int = 365

    def forecast(self, history: pd.DataFrame, horizons: np.ndarray) -> Forecast:
        returns = log_returns(history)
        mu = returns.iloc[-self.drift_window :].mean()
        sigma = returns.iloc[-self.vol_window :].std()
        mean = np.log(history["close"].iloc[-1]) + mu * horizons
        return Forecast.normal(history.index[-1], horizons, mean, sigma * np.sqrt(horizons))