"""Stacking (Stacked Generalization) ensemble implementations for classification and regression."""

import numpy as np
from typing import Dict, List, Any, Optional, Type
from tuiml.evaluation.splitting import KFold

from tuiml.base.algorithms import Classifier, classifier, Regressor, regressor
from tuiml.hub import registry

@classifier(tags=["ensembles", "stacking", "meta"], version="1.0.0")
class StackingClassifier(Classifier):
    """StackingClassifier for **stacked generalization** meta-learning.

    StackingClassifier trains a **meta-classifier** on the outputs of diverse
    base classifiers. Level-one **meta-features** are generated using
    cross-validation to avoid overfitting in the stacking process.

    Overview
    --------
    The algorithm proceeds in two levels:

    1. **Level-0 (Base classifiers):**
       a. Use :math:`k`-fold cross-validation on the training data
       b. For each fold, train all base classifiers on the training portion
       c. Generate meta-features (predictions or probabilities) on the validation portion
    2. **Level-1 (Meta-classifier):**
       a. Train all base classifiers on the full training set (for future prediction)
       b. Train the meta-classifier on the cross-validated meta-features
    3. To predict new data, generate meta-features from level-0 classifiers
       and pass them through the level-1 meta-classifier

    Theory
    ------
    Given :math:`L` base classifiers :math:`h_1, \\ldots, h_L`, the meta-features
    for sample :math:`x_i` are constructed as:

    .. math::
        z_i = \\bigl(h_1(x_i), h_2(x_i), \\ldots, h_L(x_i)\\bigr)

    When ``use_probabilities=True``, each :math:`h_l(x_i)` is a :math:`K`-dimensional
    probability vector, yielding :math:`L \\times K` meta-features.

    The meta-classifier :math:`g` learns the combination function:

    .. math::
        H(x) = g(z) = g\\bigl(h_1(x), h_2(x), \\ldots, h_L(x)\\bigr)

    Cross-validation prevents overfitting by ensuring that meta-features for
    training sample :math:`x_i` are generated by classifiers that did not see
    :math:`x_i` during training.

    Parameters
    ----------
    classifiers : list, default=['NaiveBayesClassifier', 'C45TreeClassifier']
        The base classifiers to use. Can be classifier names (strings),
        classes, or instances.

    meta_classifier : str or object, default='Logistic'
        The meta-classifier used to combine base classifier predictions.

    num_folds : int, default=10
        The number of folds for cross-validation to generate meta-features.
        Must be at least 2.

    use_probabilities : bool, default=True
        Whether to use class probabilities (``True``) or one-hot encoded
        hard labels (``False``) as meta-features.

    random_state : int or None, default=None
        Random seed for reproducibility of the cross-validation splits.

    Attributes
    ----------
    base_estimators_ : list
        The collection of fitted base classifiers (trained on full data).

    meta_estimator_ : Classifier
        The fitted meta-classifier.

    classes_ : np.ndarray
        The unique class labels discovered during ``fit()``.

    Notes
    -----
    **Complexity:**

    - Training: :math:`O(k \\cdot \\sum_{l=1}^{L} C_l + C_{\\text{meta}})` where :math:`k` = num_folds,
      :math:`C_l` = complexity of base classifier :math:`l`, :math:`C_{\\text{meta}}` = meta-classifier complexity
    - Prediction: :math:`O(\\sum_{l=1}^{L} C_l^{\\text{pred}} + C_{\\text{meta}}^{\\text{pred}})` per sample

    **When to use StackingClassifier:**

    - When you have diverse base classifiers with complementary strengths
    - When simple voting does not capture complex relationships between classifiers
    - When you have enough data to support cross-validation without underfitting
    - When you want the meta-learner to learn optimal combination weights automatically

    References
    ----------
    .. [Wolpert1992] Wolpert, D.H. (1992).
           **Stacked Generalization.**
           *Neural Networks*, 5(2), 241-259.
           DOI: `10.1016/S0893-6080(05)80023-1 <https://doi.org/10.1016/S0893-6080(05)80023-1>`_

    .. [Breiman1996c] Breiman, L. (1996).
           **Stacked Regressions.**
           *Machine Learning*, 24(1), 49-64.
           DOI: `10.1007/BF00117832 <https://doi.org/10.1007/BF00117832>`_

    .. [Ting1999] Ting, K.M. and Witten, I.H. (1999).
           **Issues in Stacked Generalization.**
           *Journal of Artificial Intelligence Research*, 10, 271-289.

    See Also
    --------
    :class:`~tuiml.algorithms.ensemble.VotingClassifier` : Simpler combination via voting rules.
    :class:`~tuiml.algorithms.ensemble.AdaBoostClassifier` : Sequential boosting ensemble.
    :class:`~tuiml.algorithms.ensemble.BaggingClassifier` : Parallel bootstrap aggregating.

    Examples
    --------
    Basic usage for stacked generalization:

    >>> from tuiml.algorithms.ensemble import StackingClassifier
    >>> import numpy as np
    >>>
    >>> # Create sample training data
    >>> X_train = np.array([[1, 2], [2, 3], [3, 1], [4, 3], [5, 2]])
    >>> y_train = np.array([0, 0, 1, 1, 1])
    >>>
    >>> # Fit the Stacking classifier
    >>> clf = StackingClassifier(
    ...     classifiers=['NaiveBayesClassifier', 'C45TreeClassifier'],
    ...     meta_classifier='Logistic'
    ... )
    >>> clf.fit(X_train, y_train)
    StackingClassifier(...)
    >>> predictions = clf.predict(X_train)
    """

    def __init__(self, classifiers: List[Any] = None,
                 meta_classifier: Any = 'Logistic',
                 num_folds: int = 10,
                 use_probabilities: bool = True,
                 random_state: Optional[int] = None):
        """Initialize StackingClassifier.

        Parameters
        ----------
        classifiers : list or None, default=None
            Base classifiers. Defaults to ``['NaiveBayesClassifier', 'C45TreeClassifier']``.
        meta_classifier : str or object, default='Logistic'
            The meta-classifier for combining base predictions.
        num_folds : int, default=10
            Number of cross-validation folds for meta-feature generation.
        use_probabilities : bool, default=True
            Whether to use probability outputs as meta-features.
        random_state : int or None, default=None
            Random seed for reproducibility.
        """
        super().__init__()
        self.classifiers = classifiers or ['NaiveBayesClassifier', 'C45TreeClassifier']
        self.meta_classifier = meta_classifier
        self.num_folds = num_folds
        self.use_probabilities = use_probabilities
        self.random_state = random_state
        self.base_estimators_ = None
        self.meta_estimator_ = None
        self.classes_ = None
        self._n_classes = None

    @classmethod
    def get_parameter_schema(cls) -> Dict[str, Dict[str, Any]]:
        return {
            "classifiers": {"type": "array", "default": ["NaiveBayesClassifier", "C45TreeClassifier"],
                           "description": "List of base classifier names"},
            "meta_classifier": {"type": "string", "default": "Logistic",
                               "description": "Meta-classifier name"},
            "num_folds": {"type": "integer", "default": 10, "minimum": 2,
                         "description": "CV folds for meta-features"},
            "use_probabilities": {"type": "boolean", "default": True,
                                 "description": "Use probability outputs"},
            "random_state": {"type": "integer", "default": None,
                            "description": "Random seed"}
        }

    @classmethod
    def get_capabilities(cls) -> List[str]:
        return ["numeric", "nominal", "binary_class", "multiclass"]

    @classmethod
    def get_complexity(cls) -> str:
        return "O(n * base_complexities * num_folds + n * meta_complexity)"

    @classmethod
    def get_references(cls) -> List[str]:
        return ["Wolpert, D.H. (1992). Stacked Generalization. "
                "Neural Networks, 5(2), 241-259."]

    def _get_classifier(self, clf_spec: Any) -> Classifier:
        """Instantiate a classifier from a specification.

        Parameters
        ----------
        clf_spec : str, type, or Classifier
            Classifier name, class, or instance to resolve.

        Returns
        -------
        classifier : Classifier
            A new classifier instance.
        """
        if isinstance(clf_spec, str):
            clf_class = registry.get(clf_spec)
            return clf_class()
        elif isinstance(clf_spec, type):
            return clf_spec()
        elif isinstance(clf_spec, Classifier):
            # Clone by creating new instance
            return type(clf_spec)(**clf_spec.get_params())
        else:
            raise ValueError(f"Invalid classifier specification: {clf_spec}")

    def _get_meta_features(self, X: np.ndarray, estimators: List[Classifier]) -> np.ndarray:
        """Generate meta-features from base classifier predictions.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Input data to generate meta-features for.
        estimators : list of Classifier
            Fitted base classifiers.

        Returns
        -------
        meta_features : np.ndarray of shape (n_samples, n_meta_features)
            Concatenated predictions or probabilities from all base classifiers.
        """
        meta_features = []

        for est in estimators:
            if self.use_probabilities:
                try:
                    proba = est.predict_proba(X)
                    meta_features.append(proba)
                except NotImplementedError:
                    # Fall back to one-hot encoding
                    preds = est.predict(X)
                    one_hot = np.zeros((len(X), self._n_classes))
                    for i, pred in enumerate(preds):
                        idx = np.where(self.classes_ == pred)[0]
                        if len(idx) > 0:
                            one_hot[i, idx[0]] = 1.0
                    meta_features.append(one_hot)
            else:
                preds = est.predict(X)
                one_hot = np.zeros((len(X), self._n_classes))
                for i, pred in enumerate(preds):
                    idx = np.where(self.classes_ == pred)[0]
                    if len(idx) > 0:
                        one_hot[i, idx[0]] = 1.0
                meta_features.append(one_hot)

        return np.hstack(meta_features)

    def fit(self, X: np.ndarray, y: np.ndarray) -> "StackingClassifier":
        """Fit the StackingClassifier classifier.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Training data.
        y : np.ndarray of shape (n_samples,)
            Target labels.

        Returns
        -------
        self : StackingClassifier
            Returns the fitted instance.
        """
        X = np.asarray(X, dtype=float)
        y = np.asarray(y)
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        n_samples = len(y)
        self.classes_ = np.unique(y)
        self._n_classes = len(self.classes_)

        # Generate meta-features using cross-validation
        n_base = len(self.classifiers)
        if self.use_probabilities:
            meta_X = np.zeros((n_samples, n_base * self._n_classes))
        else:
            meta_X = np.zeros((n_samples, n_base * self._n_classes))

        kf = KFold(n_splits=min(self.num_folds, n_samples),
                   shuffle=True, random_state=self.random_state)

        for train_idx, val_idx in kf.split(X):
            X_train, X_val = X[train_idx], X[val_idx]
            y_train = y[train_idx]

            # Train base classifiers on fold
            fold_estimators = []
            for clf_spec in self.classifiers:
                clf = self._get_classifier(clf_spec)
                clf.fit(X_train, y_train)
                fold_estimators.append(clf)

            # Generate meta-features for validation set
            fold_meta = self._get_meta_features(X_val, fold_estimators)
            meta_X[val_idx] = fold_meta

        # Train base classifiers on full training set
        self.base_estimators_ = []
        for clf_spec in self.classifiers:
            clf = self._get_classifier(clf_spec)
            clf.fit(X, y)
            self.base_estimators_.append(clf)

        # Train meta-classifier
        self.meta_estimator_ = self._get_classifier(self.meta_classifier)
        self.meta_estimator_.fit(meta_X, y)

        self._is_fitted = True
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """Predict class labels for samples.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Test data.

        Returns
        -------
        y_pred : np.ndarray of shape (n_samples,)
            Predicted class labels.
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        meta_X = self._get_meta_features(X, self.base_estimators_)
        return self.meta_estimator_.predict(meta_X)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """Predict class probabilities for samples.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Test data.

        Returns
        -------
        proba : np.ndarray of shape (n_samples, n_classes)
            The class probabilities of the input samples.
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        meta_X = self._get_meta_features(X, self.base_estimators_)

        try:
            return self.meta_estimator_.predict_proba(meta_X)
        except NotImplementedError:
            # Fall back to hard predictions
            preds = self.meta_estimator_.predict(meta_X)
            proba = np.zeros((len(X), self._n_classes))
            for i, pred in enumerate(preds):
                idx = np.where(self.classes_ == pred)[0]
                if len(idx) > 0:
                    proba[i, idx[0]] = 1.0
            return proba

    def __repr__(self) -> str:
        if self._is_fitted:
            return (f"StackingClassifier(n_base={len(self.base_estimators_)}, "
                   f"meta={self.meta_classifier})")
        return f"StackingClassifier(classifiers={self.classifiers})"


@regressor(tags=["ensemble", "stacking", "regression"], version="1.0.0")
class StackingRegressor(Regressor):
    """StackingRegressor for **stacked generalization** meta-learning regression.

    StackingRegressor trains a **meta-regressor** on the outputs of diverse
    base regressors. Level-one **meta-features** are generated using
    cross-validation to avoid overfitting in the stacking process.

    Overview
    --------
    The algorithm proceeds in two levels:

    1. **Level-0 (Base regressors):**
       a. Use :math:`k`-fold cross-validation on the training data
       b. For each fold, train all base regressors on the training portion
       c. Generate meta-features (predictions) on the validation portion
    2. **Level-1 (Meta-regressor):**
       a. Train all base regressors on the full training set (for future prediction)
       b. Train the meta-regressor on the cross-validated meta-features
    3. To predict new data, generate meta-features from level-0 regressors
       and pass them through the level-1 meta-regressor

    Theory
    ------
    Given :math:`L` base regressors :math:`h_1, \\ldots, h_L`, the meta-features
    for sample :math:`x_i` are constructed as:

    .. math::
        z_i = \\bigl(h_1(x_i), h_2(x_i), \\ldots, h_L(x_i)\\bigr)

    The meta-regressor :math:`g` learns the combination function:

    .. math::
        H(x) = g(z) = g\\bigl(h_1(x), h_2(x), \\ldots, h_L(x)\\bigr)

    Cross-validation prevents overfitting by ensuring that meta-features for
    training sample :math:`x_i` are generated by regressors that did not see
    :math:`x_i` during training.

    Parameters
    ----------
    regressors : list, default=['AdditiveRegression']
        The base regressors to use. Can be regressor names (strings),
        classes, or instances.

    meta_regressor : str or object, default='AdditiveRegression'
        The meta-regressor used to combine base regressor predictions.

    num_folds : int, default=10
        The number of folds for cross-validation to generate meta-features.
        Must be at least 2.

    random_state : int or None, default=None
        Random seed for reproducibility of the cross-validation splits.

    Attributes
    ----------
    base_estimators_ : list
        The collection of fitted base regressors (trained on full data).

    meta_estimator_ : Regressor
        The fitted meta-regressor.

    Notes
    -----
    **Complexity:**

    - Training: :math:`O(k \\cdot \\sum_{l=1}^{L} C_l + C_{\\text{meta}})` where :math:`k` = num_folds,
      :math:`C_l` = complexity of base regressor :math:`l`, :math:`C_{\\text{meta}}` = meta-regressor complexity
    - Prediction: :math:`O(\\sum_{l=1}^{L} C_l^{\\text{pred}} + C_{\\text{meta}}^{\\text{pred}})` per sample

    **When to use StackingRegressor:**

    - When you have diverse base regressors with complementary strengths
    - When simple averaging does not capture complex relationships between regressors
    - When you have enough data to support cross-validation without underfitting
    - When you want the meta-learner to learn optimal combination weights automatically

    References
    ----------
    .. [Wolpert1992] Wolpert, D.H. (1992).
           **Stacked Generalization.**
           *Neural Networks*, 5(2), 241-259.
           DOI: `10.1016/S0893-6080(05)80023-1 <https://doi.org/10.1016/S0893-6080(05)80023-1>`_

    .. [Breiman1996c] Breiman, L. (1996).
           **Stacked Regressions.**
           *Machine Learning*, 24(1), 49-64.
           DOI: `10.1007/BF00117832 <https://doi.org/10.1007/BF00117832>`_

    See Also
    --------
    :class:`~tuiml.algorithms.ensemble.StackingClassifier` : Stacked generalization for classification.
    :class:`~tuiml.algorithms.ensemble.VotingRegressor` : Simpler combination via averaging.
    :class:`~tuiml.algorithms.ensemble.BaggingRegressor` : Parallel bootstrap aggregating for regression.

    Examples
    --------
    Basic usage for stacked generalization regression:

    >>> from tuiml.algorithms.ensemble import StackingRegressor
    >>> import numpy as np
    >>>
    >>> # Create sample training data
    >>> X_train = np.array([[1, 2], [2, 3], [3, 1], [4, 3], [5, 2]])
    >>> y_train = np.array([1.5, 2.3, 3.1, 4.2, 5.0])
    >>>
    >>> # Fit the Stacking regressor
    >>> reg = StackingRegressor(
    ...     regressors=['AdditiveRegression'],
    ...     meta_regressor='AdditiveRegression'
    ... )
    >>> reg.fit(X_train, y_train)
    StackingRegressor(...)
    >>> predictions = reg.predict(X_train)
    """

    def __init__(self, regressors: List[Any] = None,
                 meta_regressor: Any = 'AdditiveRegression',
                 num_folds: int = 10,
                 random_state: Optional[int] = None):
        """Initialize StackingRegressor.

        Parameters
        ----------
        regressors : list or None, default=None
            Base regressors. Defaults to ``['AdditiveRegression']``.
        meta_regressor : str or object, default='AdditiveRegression'
            The meta-regressor for combining base predictions.
        num_folds : int, default=10
            Number of cross-validation folds for meta-feature generation.
        random_state : int or None, default=None
            Random seed for reproducibility.
        """
        super().__init__()
        self.regressors = regressors or ['AdditiveRegression']
        self.meta_regressor = meta_regressor
        self.num_folds = num_folds
        self.random_state = random_state
        self.base_estimators_ = None
        self.meta_estimator_ = None

    @classmethod
    def get_parameter_schema(cls) -> Dict[str, Dict[str, Any]]:
        """Return JSON Schema for constructor parameters."""
        return {
            "regressors": {"type": "array", "default": ["AdditiveRegression"],
                           "description": "List of base regressor names"},
            "meta_regressor": {"type": "string", "default": "AdditiveRegression",
                               "description": "Meta-regressor name"},
            "num_folds": {"type": "integer", "default": 10, "minimum": 2,
                         "description": "CV folds for meta-features"},
            "random_state": {"type": "integer", "default": None,
                            "description": "Random seed"}
        }

    @classmethod
    def get_capabilities(cls) -> List[str]:
        """Return algorithm capabilities."""
        return ["numeric", "numeric_class"]

    @classmethod
    def get_complexity(cls) -> str:
        """Return time/space complexity."""
        return "O(n * base_complexities * num_folds + n * meta_complexity)"

    @classmethod
    def get_references(cls) -> List[str]:
        """Return academic references."""
        return ["Wolpert, D.H. (1992). Stacked Generalization. "
                "Neural Networks, 5(2), 241-259."]

    def _get_regressor(self, reg_spec: Any) -> Regressor:
        """Instantiate a regressor from a specification.

        Parameters
        ----------
        reg_spec : str, type, or Regressor
            Regressor name, class, or instance to resolve.

        Returns
        -------
        regressor : Regressor
            A new regressor instance.
        """
        if isinstance(reg_spec, str):
            reg_class = registry.get(reg_spec)
            return reg_class()
        elif isinstance(reg_spec, type):
            return reg_spec()
        elif isinstance(reg_spec, Regressor):
            return type(reg_spec)()
        else:
            raise ValueError(f"Invalid regressor specification: {reg_spec}")

    def _get_meta_features(self, X: np.ndarray, estimators: List[Regressor]) -> np.ndarray:
        """Generate meta-features from base regressor predictions.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Input data to generate meta-features for.
        estimators : list of Regressor
            Fitted base regressors.

        Returns
        -------
        meta_features : np.ndarray of shape (n_samples, n_regressors)
            Predictions from all base regressors as columns.
        """
        meta_features = []
        for est in estimators:
            preds = est.predict(X).reshape(-1, 1)
            meta_features.append(preds)
        return np.hstack(meta_features)

    def fit(self, X: np.ndarray, y: np.ndarray) -> "StackingRegressor":
        """Fit the StackingRegressor.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Training data.
        y : np.ndarray of shape (n_samples,)
            Target values.

        Returns
        -------
        self : StackingRegressor
            Returns the fitted instance.
        """
        X = np.asarray(X, dtype=float)
        y = np.asarray(y, dtype=float)
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        n_samples = len(y)
        n_base = len(self.regressors)

        # Generate meta-features using cross-validation
        meta_X = np.zeros((n_samples, n_base))

        kf = KFold(n_splits=min(self.num_folds, n_samples),
                   shuffle=True, random_state=self.random_state)

        for train_idx, val_idx in kf.split(X):
            X_train, X_val = X[train_idx], X[val_idx]
            y_train = y[train_idx]

            fold_estimators = []
            for reg_spec in self.regressors:
                reg = self._get_regressor(reg_spec)
                reg.fit(X_train, y_train)
                fold_estimators.append(reg)

            fold_meta = self._get_meta_features(X_val, fold_estimators)
            meta_X[val_idx] = fold_meta

        # Train base regressors on full training set
        self.base_estimators_ = []
        for reg_spec in self.regressors:
            reg = self._get_regressor(reg_spec)
            reg.fit(X, y)
            self.base_estimators_.append(reg)

        # Train meta-regressor
        self.meta_estimator_ = self._get_regressor(self.meta_regressor)
        self.meta_estimator_.fit(meta_X, y)

        self._is_fitted = True
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        """Predict target values for samples.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Test data.

        Returns
        -------
        y_pred : np.ndarray of shape (n_samples,)
            Predicted target values.
        """
        self._check_is_fitted()
        X = np.asarray(X, dtype=float)
        if X.ndim == 1:
            X = X.reshape(-1, 1)

        meta_X = self._get_meta_features(X, self.base_estimators_)
        return self.meta_estimator_.predict(meta_X)

    def score(self, X: np.ndarray, y: np.ndarray) -> float:
        """Compute R-squared score.

        Parameters
        ----------
        X : np.ndarray of shape (n_samples, n_features)
            Test features.
        y : np.ndarray of shape (n_samples,)
            True target values.

        Returns
        -------
        score : float
            R-squared score.
        """
        self._check_is_fitted()
        y_pred = self.predict(X)
        y = np.asarray(y)

        ss_res = np.sum((y - y_pred) ** 2)
        ss_tot = np.sum((y - np.mean(y)) ** 2)

        if ss_tot == 0:
            return 0.0

        return 1 - (ss_res / ss_tot)

    def __repr__(self) -> str:
        """Return string representation."""
        if self._is_fitted:
            return (f"StackingRegressor(n_base={len(self.base_estimators_)}, "
                   f"meta={self.meta_regressor})")
        return f"StackingRegressor(regressors={self.regressors})"
