from abc import ABC, abstractmethod
from omegaconf import DictConfig
from typing import Any, Dict
from haipr.data import HAIPRData


class BasePredictor(ABC):
    """Abstract base class defining the interface for all predictors."""

    @abstractmethod
    def setup_model(self, data: HAIPRData, cfg: DictConfig) -> None:
        """
        Initialize predictor with configuration and set the alphabet map.
        """
        pass

    @abstractmethod
    def fit_model(
        self,
        dataset: HAIPRData,
        train_indices: Any,
        val_indices: Any,
    ) -> Dict[str, Any]:
        """Train the predictor on given data.
        Returns a dictionary with the following keys:
        - metrics: a dictionary of metrics
        - predictions: a dictionary of predictions
        """
        pass
    
    @abstractmethod
    def prepare_features(self, batch_items: list[dict]) -> dict:
        """Prepare features for the model."""
        pass

    @abstractmethod
    def load_model(self, model: str) -> None:
        """Load the model."""
        pass

    @abstractmethod
    def predict(self, **kwargs) -> Dict[str, Any]:
        """Make predictions"""
        pass
