src/models/base.py

"""Abstract base class for all prediction models."""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional

import numpy as np
import pandas as pd


@dataclass
class PredictionResult:
    """Standard prediction output from all models."""
    direction: np.ndarray        # -1 (down), 0 (flat), 1 (up)
    direction_proba: np.ndarray  # probability for each class [n_samples, 3]
    magnitude: np.ndarray        # expected return %
    volatility: np.ndarray       # expected volatility
    confidence: np.ndarray       # model confidence score [0, 1]


class BaseModel(ABC):
    """Abstract base class all models must implement."""

    def __init__(self, name: str, stock_type: str, horizon: int):
        self.name = name
        self.stock_type = stock_type
        self.horizon = horizon
        self.is_fitted = False

    @abstractmethod
    def fit(
        self,
        X_train: pd.DataFrame,
        y_train: pd.DataFrame,
        X_val: Optional[pd.DataFrame] = None,
        y_val: Optional[pd.DataFrame] = None,
    ) -> "BaseModel":
        """Train the model. y_train has columns: direction, magnitude, volatility."""
        ...

    @abstractmethod
    def predict(self, X: pd.DataFrame) -> PredictionResult:
        """Generate predictions."""
        ...

    def save(self, path: str) -> None:
        """Save model artifacts to disk."""
        raise NotImplementedError

    @classmethod
    def load(cls, path: str) -> "BaseModel":
        """Load model artifacts from disk."""
        raise NotImplementedError

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name={self.name!r}, type={self.stock_type!r}, horizon={self.horizon})"