from __future__ import annotations
from typing import Tuple, NoReturn
from ...base import BaseEstimator
from ...metrics.loss_functions import misclassification_error
import numpy as np
from itertools import product


class DecisionStump(BaseEstimator):
    """
    A decision stump classifier for {-1,1} labels according to the CART algorithm

    Attributes
    ----------
    self.threshold_ : float
        The threshold by which the data is split

    self.j_ : int
        The index of the feature by which to split the data

    self.sign_: int
        The label to predict for samples where the value of the j'th feature is equal or above the threshold
    """
    def __init__(self) -> DecisionStump:
        """
        Instantiate a Decision stump classifier
        """
        super().__init__()
        self.threshold_, self.j_, self.sign_ = None, None, None

    def _fit(self, X: np.ndarray, y: np.ndarray) -> NoReturn:
        """
        fits a decision stump to the given data

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_features)
            Input data to fit an estimator for

        y : ndarray of shape (n_samples, )
            Responses of input data to fit to
        """
        err_list = []
        for sign in [-1, 1]:
            for i, feature in enumerate(X.T):
                thr, thr_err = self._find_threshold(feature, y, sign)
                err_list.append([i, thr, thr_err, sign])
        self.j_, self.threshold_, _, self.sign_ = min(err_list, key=lambda t: t[2])

    def _predict(self, X: np.ndarray) -> np.ndarray:
        """
        Predict responses for given samples using fitted estimator

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_features)
            Input data to predict responses for

        Returns
        -------
        responses : ndarray of shape (n_samples, )
            Predicted responses of given samples

        Notes
        -----
        Feature values strictly below threshold are predicted as `-sign` whereas values which equal
        to or above the threshold are predicted as `sign`
        """
        feature = X[:, self.j_]
        return np.where(feature >= self.threshold_, self.sign_, -self.sign_)

    @staticmethod
    def _weighted_misclassification_error(y_true: np.ndarray, y_pred: np.ndarray):
        sum_mistakes = np.sum(np.abs(y_true * (np.sign(y_true) != y_pred)))
        return sum_mistakes / len(y_true)

    def _find_threshold(self, values: np.ndarray, labels: np.ndarray, sign: int) -> Tuple[float, float]:
        """
        Given a feature vector and labels, find a threshold by which to perform a split
        The threshold is found according to the value minimizing the misclassification
        error along this feature

        Parameters
        ----------
        values: ndarray of shape (n_samples,)
            A feature vector to find a splitting threshold for

        labels: ndarray of shape (n_samples,)
            The labels to compare against

        sign: int
            Predicted label assigned to values equal to or above threshold

        Returns
        -------
        thr: float
            Threshold by which to perform split

        thr_err: float between 0 and 1
            Misclassificaiton error of returned threshold

        Notes
        -----
        For every tested threshold, values strictly below threshold are predicted as `-sign` whereas values
        which equal to or above the threshold are predicted as `sign`
        """
        sorted_values = np.sort(values)
        sorted_values_and_limits = np.concatenate((sorted_values, [sorted_values[-1] + 1]))
        err_list = []
        for thr_value in sorted_values_and_limits:
            pred_labels = np.where(values >= thr_value, sign, -sign)
            err = self._weighted_misclassification_error(labels, pred_labels)
            err_list.append((thr_value, err))
        return min(err_list, key=lambda t: t[1])

    def _loss(self, X: np.ndarray, y: np.ndarray) -> float:
        """
        Evaluate performance under misclassification loss function

        Parameters
        ----------
        X : ndarray of shape (n_samples, n_features)
            Test samples

        y : ndarray of shape (n_samples, )
            True labels of test samples

        Returns
        -------
        loss : float
            Performance under missclassification loss function
        """
        y_pred = self._predict(X)
        return misclassification_error(y, y_pred)
