"""
Models built from parts.

Most ideas about Bitcoin prices say something about either the expected return
(drift) or the size of the uncertainty (volatility). A Composite pairs one of
each, so an A/B test can swap exactly one part and hold the other fixed.
"""

from dataclasses import dataclass
from typing import Protocol

import numpy as np
import pandas as pd

from ..forecast import Forecast
from .shape import Normal


class Drift(Protocol):
    def expected_log_return(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        """Expected cumulative log return from the origin to each horizon."""
        ...


class Volatility(Protocol):
    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        """Standard deviation of the cumulative log return at each horizon."""
        ...


class Shape(Protocol):
    def standard_quantiles(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        """Quantiles at LEVELS of a mean-0, sd-1 distribution, shape (H, N_LEVELS)."""
        ...


@dataclass(frozen=True)
class Composite:
    """Log price: mean from `drift`, spread from `volatility`, normal unless `shape` says."""

    name: str
    drift: Drift
    volatility: Volatility
    shape: Shape = Normal()

    def forecast(self, history: pd.DataFrame, horizons: np.ndarray) -> Forecast:
        mean = np.log(history["close"].iloc[-1]) + self.drift.expected_log_return(history, horizons)
        sd = self.volatility.sd(history, horizons)
        z = self.shape.standard_quantiles(history, horizons)
        return Forecast(history.index[-1], horizons, mean[:, None] + sd[:, None] * z)