"""
Main Author: Will LeVine
Corresponding Email: levinewill@icloud.com
"""
import keras
import numpy as np
from sklearn.tree import DecisionTreeClassifier
from sklearn.utils.validation import check_array, check_is_fitted, check_X_y

from .base import BaseTransformer


class NeuralClassificationTransformer(BaseTransformer):
    """
    A class used to transform data from a category to a specialized representation.

    Parameters
    ----------
    network : object
        A neural network used in the classification transformer.

    euclidean_layer_idx : int
        An integer to represent the final layer of the transformer.

    optimizer : str or keras.optimizers instance
        An optimizer used when compiling the neural network.

    loss : str, default="categorical_crossentropy"
        A loss function used when compiling the neural network.

    pretrained : bool, default=False
        A boolean used to identify if the network is pretrained.

    compile_kwargs : dict, default={"metrics": ["acc"]}
        A dictionary containing metrics for judging network performance.

    fit_kwargs : dict, default={
                "epochs": 100,
                "callbacks": [keras.callbacks.EarlyStopping(patience=5, monitor="val_acc")],
                "verbose": False,
                "validation_split": 0.33,
            },
        A dictionary to hold epochs, callbacks, verbose, and validation split for the network.

    Attributes
    ----------
    encoder_ : object
        A Keras model with inputs and outputs based on the network attribute.
        Output layers are determined by the euclidean_layer_idx parameter.
    """

    def __init__(
        self,
        network,
        euclidean_layer_idx,
        optimizer,
        loss="categorical_crossentropy",
        pretrained=False,
        compile_kwargs={"metrics": ["acc"]},
        fit_kwargs={
            "epochs": 100,
            "callbacks": [keras.callbacks.EarlyStopping(patience=5, monitor="val_acc")],
            "verbose": False,
            "validation_split": 0.33,
        },
    ):
        self.network = keras.models.clone_model(network)
        self.encoder_ = keras.models.Model(
            inputs=self.network.inputs,
            outputs=self.network.layers[euclidean_layer_idx].output,
        )
        self.pretrained = pretrained
        self.optimizer = optimizer
        self.loss = loss
        self.compile_kwargs = compile_kwargs
        self.fit_kwargs = fit_kwargs

    def fit(self, X, y):
        """
        Fits the transformer to data X with labels y.

        Parameters
        ----------
        X : ndarray
            Input data matrix.
        y : ndarray
            Output (i.e. response data matrix).

        Returns
        -------
        self : NeuralClassificationTransformer
            The object itself.
        """
        check_X_y(X, y)
        _, y = np.unique(y, return_inverse=True)

        # more typechecking
        self.network.compile(
            loss=self.loss, optimizer=self.optimizer, **self.compile_kwargs
        )

        self.network.fit(X, keras.utils.to_categorical(y), **self.fit_kwargs)

        return self

    def transform(self, X):
        """
        Performs inference using the transformer.

        Parameters
        ----------
        X : ndarray
            Input data matrix.

        Returns
        -------
        X_transformed : ndarray
            The transformed input.

        Raises
        ------
        NotFittedError
            When the model is not fitted.
        """
        check_is_fitted(self)
        check_array(X)
        return self.encoder_.predict(X)


class TreeClassificationTransformer(BaseTransformer):
    """
    A class used to transform data from a category to a specialized representation.

    Parameters
    ----------
    kwargs : dict, default={}
        A dictionary to contain parameters of the tree.

    Attributes
    ----------
    transformer : sklearn.tree.DecisionTreeClassifier
        an internal sklearn DecisionTreeClassifier
    """

    def __init__(self, kwargs={}):
        self.kwargs = kwargs

    def fit(self, X, y):
        """
        Fits the transformer to data X with labels y.

        Parameters
        ----------
        X : ndarray
            Input data matrix.
        y : ndarray
            Output (i.e. response data matrix).

        Returns
        -------
        self : TreeClassificationTransformer
            The object itself.
        """
        X, y = check_X_y(X, y)
        self.transformer_ = DecisionTreeClassifier(**self.kwargs).fit(X, y)
        return self

    def transform(self, X):
        """
        Performs inference using the transformer.

        Parameters
        ----------
        X : ndarray
            Input data matrix.

        Returns
        -------
        X_transformed : ndarray
            The transformed input.

        Raises
        ------
        NotFittedError
            When the model is not fitted.
        """
        check_is_fitted(self)
        X = check_array(X)
        return self.transformer_.apply(X)

[((67, 23, 67, 56), 'keras.models.clone_model', 'keras.models.clone_model', ({(67, 48, 67, 55): 'network'}, {}), '(network)', False, 'import keras\n'), ((68, 24, 71, 9), 'keras.models.Model', 'keras.models.Model', (), '', False, 'import keras\n'), ((94, 8, 94, 23), 'sklearn.utils.validation.check_X_y', 'check_X_y', ({(94, 18, 94, 19): 'X', (94, 21, 94, 22): 'y'}, {}), '(X, y)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((95, 15, 95, 48), 'numpy.unique', 'np.unique', (), '', True, 'import numpy as np\n'), ((125, 8, 125, 29), 'sklearn.utils.validation.check_is_fitted', 'check_is_fitted', ({(125, 24, 125, 28): 'self'}, {}), '(self)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((126, 8, 126, 22), 'sklearn.utils.validation.check_array', 'check_array', ({(126, 20, 126, 21): 'X'}, {}), '(X)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((164, 15, 164, 30), 'sklearn.utils.validation.check_X_y', 'check_X_y', ({(164, 25, 164, 26): 'X', (164, 28, 164, 29): 'y'}, {}), '(X, y)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((187, 8, 187, 29), 'sklearn.utils.validation.check_is_fitted', 'check_is_fitted', ({(187, 24, 187, 28): 'self'}, {}), '(self)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((188, 12, 188, 26), 'sklearn.utils.validation.check_array', 'check_array', ({(188, 24, 188, 25): 'X'}, {}), '(X)', False, 'from sklearn.utils.validation import check_array, check_is_fitted, check_X_y\n'), ((102, 28, 102, 57), 'keras.utils.to_categorical', 'keras.utils.to_categorical', ({(102, 55, 102, 56): 'y'}, {}), '(y)', False, 'import keras\n'), ((62, 26, 62, 86), 'keras.callbacks.EarlyStopping', 'keras.callbacks.EarlyStopping', (), '', False, 'import keras\n'), ((165, 28, 165, 65), 'sklearn.tree.DecisionTreeClassifier', 'DecisionTreeClassifier', ({}, {}), '(**self.kwargs)', False, 'from sklearn.tree import DecisionTreeClassifier\n')]