"""Model wrapper for standardizing prediction interfaces."""

import numpy as np
import pandas as pd


class ModelWrapper:
    """A wrapper to standardize model prediction for XAI libraries."""

    def __init__(self, model):
        self.model = model

    def _ensure_2d(self, data):
        """Helper to ensure data is a 2D numpy array."""
        if isinstance(data, pd.DataFrame):
            data = data.values
        if isinstance(data, np.ndarray) and data.ndim == 1:
            data = data.reshape(1, -1)
        return data

    def predict_proba(self, data):
        """
        Provides a probability prediction function for classifiers.
        Returns a 2D array of shape (n_samples, n_classes).
        """
        data = self._ensure_2d(data)
        if hasattr(self.model, "predict_proba"):
            return self.model.predict_proba(data)
        else:
            # Fallback for models without predict_proba but used in classification context
            scores = self.predict_score(data)
            prob_pos = 1 / (1 + np.exp(-scores))
            return np.vstack([1 - prob_pos, prob_pos]).T

    def predict_score(self, data):
        """
        Provides a raw score for regression or anomaly models.
        Returns a 1D array of shape (n_samples,).
        """
        data = self._ensure_2d(data)
        if hasattr(self.model, "decision_function"):
            return self.model.decision_function(data)
        else:  # For standard regression models
            return self.model.predict(data)
