# Python standard library imports
from itertools import accumulate

# External imports
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import OneHotEncoder
from sklearn.utils.validation import check_array, check_is_fitted

def _halfway_points(values):
    """
    Generate the halfway points between the values in a sorted list.

    Parameters
    ----------
    values : array-like, shape [n_values]
        The values to generate the halfway points between.

    Returns
    -------
    list
        The halfway points between the values in `values`.

    Examples
    --------
    >>> _halfway_points([1, 2, 3])
    [1.5, 2.5]
    >>> _halfway_points([1, 2, 3, 25])
    [1.5, 2.5, 14.0]
    """
    return [(values[i] + values[i + 1]) / 2 for i in range(len(values) - 1)]

class NumericBinarizer(BaseEstimator, TransformerMixin):
    f"""
    Encode numerical features as a one-hot numeric array.

    The input to this encoder should be an array-like of integers or floats.
    The features are encoded using greater-than-or-equal conditions on the 
    midpoint between each unique value in each column. This creates a binary 
    column for each of these conditions, and the resulting array has a one-hot
    encoding of the original features.

    Attributes
    ----------
    n_features_in_ : int
        The number of features passed to :term:`fit`.

    n_features_out_ : int
        The number of features generated by :term:`transform`.

    feature_names_in_ : list of str
        The names of the features passed to :term:`fit` if `X` is a DataFrame.
        Otherwise `features_in_` is a list of strings, ["x0", ..., 
        "x(n_features_in_-1)"].

    column_values_ : list of arrays
        The unique values for each feature passed to :term:`fit`. This is 
        stored in a list of length n_features_in_.

    Examples
    --------
    Given a dataset with two features, we let the encoder find the unique
    midpoint values per feature and transform the data to a binary one-hot
    encoding.

    >>> enc = NumericBinarizer()
    >>> enc.fit_transform([[0, 0], [1, 1], [0, 2], [1, 3]])
    array([[1., 1., 1., 1.],
       [0., 0., 1., 1.],
       [1., 0., 0., 1.],
       [0., 0., 0., 0.]])
    array([[0., 0.],
       [1., 1.],
       [0., 2.],
       [1., 3.]])


    An example with the well known Iris dataset, which uses set_output to 
    return a dataframe.

    >>> from sklearn.datasets import load_iris
    >>> enc = NumericBinarizer()
    >>> X = load_iris(as_frame=True).data
    >>> enc.fit_transform(X)
    array([[0., 0., 0., ..., 1., 1., 1.],
        [0., 0., 0., ..., 1., 1., 1.],
        [0., 0., 0., ..., 1., 1., 1.],
        ...,
        [0., 0., 0., ..., 1., 1., 1.],
        [0., 0., 0., ..., 0., 1., 1.],
        [0., 0., 0., ..., 1., 1., 1.]])
    >>> enc.set_output(transform='pandas').fit_transform(X)
        sepal length (cm) <= 4.35  sepal length (cm) <= 4.45  sepal length (cm) <= 4.55  sepal length (cm) <= 4.65  sepal length (cm) <= 4.75  sepal length (cm) <= 4.85  ...  petal width (cm) <= 1.95  petal width (cm) <= 2.05  petal width (cm) <= 2.1500000000000004  petal width (cm) <= 2.25  petal width (cm) <= 2.3499999999999996  petal width (cm) <= 2.45
    0                          0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    1                          0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    2                          0.0                        0.0                        0.0                        0.0                        1.0                        1.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    3                          0.0                        0.0                        0.0                        1.0                        1.0                        1.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    4                          0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    ..                         ...                        ...                        ...                        ...                        ...                        ...  ...                       ...                       ...                                     ...                       ...                                     ...                       ...
    145                        0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       0.0                       0.0                                     0.0                       0.0                                     1.0                       1.0
    146                        0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    147                        0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       0.0                       1.0                                     1.0                       1.0                                     1.0                       1.0
    148                        0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       0.0                       0.0                                     0.0                       0.0                                     1.0                       1.0
    149                        0.0                        0.0                        0.0                        0.0                        0.0                        0.0  ...                       1.0                       1.0                                     1.0                       1.0                                     1.0                       1.0

    [150 rows x 119 columns]
    >>> enc.inverse_transform(_)
        sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)
    0                  5.1               3.5                1.4               0.2
    1                  4.9               3.0                1.4               0.2
    2                  4.7               3.2                1.3               0.2
    3                  4.6               3.1                1.5               0.2
    4                  5.0               3.6                1.4               0.2
    ..                 ...               ...                ...               ...
    145                6.7               3.0                5.2               2.3
    146                6.3               2.5                5.0               1.9
    147                6.5               3.0                5.2               2.0
    148                6.2               3.4                5.4               2.3
    149                5.9               3.0                5.1               1.8

    [150 rows x 4 columns]
    """

    def __init__(self):
        pass

    def fit(self, X, y=None, columns=None):
        """
        Fit NumericBinarizer to X.

        If the user would like to produce a transformed dataset with legible 
        feature names, then `X~ should be a pandas DataFrame with named 
        columns, or provide the `columns` parameter with a list of strings.

        Parameters
        ----------
        X : array-like, shape [n_samples, n_features]
            The data to determine the unique values of each feature.

        y : default=None
            Ignored. This parameter exists only for compatibility with
            :class:`~sklearn.pipeline.Pipeline`.

        columns : array-like, shape [n_features], default=None
            The names of the features in X (optional)
            
        Returns
        -------
        self 
            Fitted encoder.
        """
        # Store the column names. columns check has to be done before check_array because it converts DataFrame to
        # ndarray.
        self.feature_names_in_ = columns
        if hasattr(X, 'columns'):
            self.feature_names_in_ = X.columns

        X = check_array(X, accept_sparse=False)

        # In the case that the user did not provide column names, we create them.
        if self.feature_names_in_ is None:
            self.feature_names_in_ = [f'x{i}' for i in range(X.shape[1])]

        self.n_features_in_ = X.shape[1]

        # Compute the values present in each column
        self.column_values_ = [np.unique(column) for column in X.T]

        # Number of transformed features per original feature
        self.n_features_out_ = sum(
            [len(values) - 1 for values in self.column_values_])

        # Return the transformer
        return self

    def get_feature_names_out(self, *args, **params):
        """
        Generates the names of the transformed features. This is necessary to 
        provide the `set_output` api that enables SciKit-Learn Transformers to
        return DataFrames (with named columns) instead of ndarrays.
        """
        check_is_fitted(self, ['n_features_in_', 'n_features_out_', 'column_values_', 'feature_names_in_'])
        return np.concatenate(
            [[f'{self.feature_names_in_[i]} <= {val}' for val in _halfway_points(self.column_values_[i])] for i in
             range(self.n_features_in_)])

    def transform(self, X):
        """
        Transform X using one-hot encoding between midpoints in features.
        For example a feature "a" with values [1, 2, 3] will be transformed 
        into a set of features with columns ["a <= 1.5", "a <= 2.5"]. The
        transformed features will be binary, with 1 indicating that the
        original feature was less than or equal to the midpoint, and 0
        otherwise.

        Parameters
        ----------
        X : array-like, shape [n_samples, n_features]
            The data to encode.

        Returns
        -------
        Xt : ndarray of shape [n_samples, n_features_out_]
            Transformed input.
        """
        # Check if fit had been called
        check_is_fitted(self, ['n_features_in_', 'n_features_out_', 'column_values_', 'feature_names_in_'])

        # input validation
        X = check_array(X, accept_sparse=False)

        # Check that the input is of the same shape as the one passed
        # during fit.
        if X.shape[1] != self.n_features_in_:
            raise ValueError('Shape of input is different form what was seen'
                             'in `fit`')

        # Return a binarization of the input samples based on halfway point splits
        Xb_lists = [[X[:, i] <= val for val in _halfway_points(self.column_values_[i])] for i in range(self.n_features_in_)]
        Xb = [col for sublist in Xb_lists for col in sublist]
        return np.column_stack(Xb).astype(float)

    def inverse_transform(self, Xt):
        """
        Convert the data back to the original representation.

        Parameters
        ----------
        Xt : array-like, shape [n_samples, n_features_out_]
            The transformed data.

        Returns
        -------
        X : ndarray of shape (n_samples, n_features_in_)
            Inverse transformed array.
        """
        # Check if fit had been called
        check_is_fitted(self, ['n_features_in_', 'n_features_out_', 'column_values_', 'feature_names_in_'])

        # input validation
        Xt = check_array(Xt, accept_sparse=False)

        # Check that the input is of the same shape as the one passed
        # during fit.
        if Xt.shape[1] != self.n_features_out_:
            raise ValueError('Shape of input is different from what was seen '
                             'in `fit`')

        # Initialize an empty array for inverse transformed data
        X = np.empty((Xt.shape[0], self.n_features_in_))

        # Iterate over the columns and reconstruct the original values
        column_idx = 0
        for i, feature_values in enumerate(self.column_values_):
            num_unique_values = len(feature_values)
            num_midpoints = num_unique_values - 1
            if num_unique_values == 1:
                # If there is only one unique value, then there is no need to
                # transform the feature.
                X[:, i] = feature_values[0]
                continue
            
            # Calculate the indices of columns for each feature
            col_indices = np.arange(column_idx, column_idx + num_midpoints)

            # Find the indices where the one-hot encoding is True
            indices = np.argmax(Xt[:, col_indices], axis=1)
            # Fix the case where all values in a row are False (This occurs for
            # the maximal value in a column).
            indices = np.where(
                np.any(Xt[:, col_indices], axis=1), indices, -1)

            # Calculate the values based on the indices and original feature values
            X[:, i] = np.array([feature_values[idx] for idx in indices])
            # Update the column index for the next feature
            column_idx += num_midpoints

        return X
    
    def feature_map(self):
        """
        Extract the feature map from the encoder. This is a dictionary that can be passed to
        the GOSDTClassifier to generate N-ary trees.
        
        Returns
        -------
        ret : dict
            A dictionary where the keys are the original feature indices and the values are the indices
            of the transformed features that correspond to the original feature.
        """
        
        # Check if fit had been called
        check_is_fitted(self, ['n_features_in_', 'n_features_out_', 'column_values_', 'feature_names_in_'])
        
        # Create the feature map
        ret = {}
        idx = 0
        for i, col in enumerate(self.column_values_):
            ret[i] = list(range(idx, idx + len(col) - 1))
            idx += len(col) - 1
        return ret
    
