"""Orthogonal PLS Discriminant Analysis (OPLS-DA) classifier for nirs4all.

See pls.py for full documentation and usage examples.
"""
import numpy as np
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.preprocessing import LabelEncoder

from .plsda import PLSDA


class OPLSDA(BaseEstimator, ClassifierMixin):
    """Orthogonal PLS Discriminant Analysis (OPLS-DA) classifier.

    # Explicitly declare estimator type for sklearn compatibility (e.g., StackingClassifier)
    _estimator_type = "classifier"

    OPLS-DA combines OPLS filtering with PLS-DA classification.
    It removes Y-orthogonal variation from X before applying PLS-DA,
    improving class separation and model interpretability.

    Parameters
    ----------
    n_components : int, default=1
        Number of orthogonal components to remove.
    pls_components : int, default=5
        Number of PLS components for the discriminant model.
    scale : bool, default=True
        Whether to scale X before fitting.

    Attributes
    ----------
    classes_ : ndarray of shape (n_classes,)
        Unique class labels.
    n_features_in_ : int
        Number of features seen during fit.
    opls_ : pyopls.OPLS
        Fitted OPLS transformer.
    plsda_ : PLSDA
        Fitted PLS-DA model on filtered data.

    Examples
    --------
    >>> from nirs4all.operators.models.sklearn.pls import OPLSDA
    >>> from sklearn.datasets import make_classification
    >>> X, y = make_classification(n_samples=100, n_features=50, n_classes=2,
    ...                            n_informative=10, random_state=42)
    >>> model = OPLSDA(n_components=1, pls_components=5)
    >>> model.fit(X, y)
    OPLSDA(n_components=1, pls_components=5)
    >>> predictions = model.predict(X)

    See Also
    --------
    PLSDA : Standard PLS-DA without orthogonal filtering.
    OPLS : OPLS for regression tasks.

    References
    ----------
    - Bylesjö, M., et al. (2006). OPLS discriminant analysis: combining
      the strengths of PLS-DA and SIMCA classification. Journal of
      Chemometrics, 20(8-10), 341-351.
    """

    _webapp_meta = {
        "category": "pls",
        "tier": "advanced",
        "tags": ["pls", "opls", "discriminant-analysis", "classification"],
    }

    def __init__(
        self,
        n_components: int = 1,
        pls_components: int = 5,
        scale: bool = True,
    ):
        """Initialize OPLSDA classifier.

        Parameters
        ----------
        n_components : int, default=1
            Number of orthogonal components to remove.
        pls_components : int, default=5
            Number of PLS components for the discriminant model.
        scale : bool, default=True
            Whether to scale X before fitting.
        """
        self.n_components = n_components
        self.pls_components = pls_components
        self.scale = scale

    def fit(self, X, y):
        """Fit the OPLS-DA model.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data.
        y : array-like of shape (n_samples,)
            Target class labels.

        Returns
        -------
        self : OPLSDA
            Fitted estimator.

        Raises
        ------
        ImportError
            If pyopls package is not installed.
        """
        from pyopls import OPLS as PyOPLS

        X = np.asarray(X)
        y = np.asarray(y).ravel()

        self.n_features_in_ = X.shape[1]
        self.classes_ = np.unique(y)

        # Encode y for OPLS fitting (use numeric encoding)
        self._label_encoder = LabelEncoder()
        y_encoded = self._label_encoder.fit_transform(y)

        # Limit components
        max_ortho = min(self.n_components, X.shape[1] - 1, X.shape[0] - 2)
        n_ortho = max(1, max_ortho)

        # Fit OPLS transformer
        self.opls_ = PyOPLS(n_components=n_ortho, scale=self.scale)
        X_filtered = self.opls_.fit_transform(X, y_encoded)

        # Fit PLS-DA on filtered data
        self.plsda_ = PLSDA(n_components=self.pls_components)
        self.plsda_.fit(X_filtered, y)

        return self

    def predict(self, X):
        """Predict class labels for samples in X.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples to predict.

        Returns
        -------
        y_pred : ndarray of shape (n_samples,)
            Predicted class labels.
        """
        X = np.asarray(X)

        # Transform X to remove orthogonal variation
        X_filtered = self.opls_.transform(X)

        # Predict with PLS-DA
        return self.plsda_.predict(X_filtered)

    def predict_proba(self, X):
        """Return pseudo-probabilities (PLS responses).

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples.

        Returns
        -------
        proba : ndarray of shape (n_samples, n_classes)
            Pseudo-probability estimates.
        """
        X = np.asarray(X)

        # Transform X to remove orthogonal variation
        X_filtered = self.opls_.transform(X)

        # Get probabilities from PLS-DA
        return self.plsda_.predict_proba(X_filtered)

    def transform(self, X):
        """Transform X by removing orthogonal variation.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Samples to transform.

        Returns
        -------
        X_filtered : ndarray of shape (n_samples, n_features)
            Transformed samples with orthogonal variation removed.
        """
        X = np.asarray(X)
        return self.opls_.transform(X)

    def get_params(self, deep=True):
        """Get parameters for this estimator.

        Parameters
        ----------
        deep : bool, default=True
            If True, will return the parameters for this estimator and
            contained subobjects that are estimators.

        Returns
        -------
        params : dict
            Parameter names mapped to their values.
        """
        return {
            'n_components': self.n_components,
            'pls_components': self.pls_components,
            'scale': self.scale,
        }

    def set_params(self, **params):
        """Set the parameters of this estimator.

        Parameters
        ----------
        **params : dict
            Estimator parameters.

        Returns
        -------
        self : OPLSDA
            Estimator instance.
        """
        for key, value in params.items():
            setattr(self, key, value)
        return self
