# Source code for dipy.nn.tf.model

    from dipy.testing.decorators import warning_for_keywords
    from dipy.utils.optpkg import optional_package

    tf, have_tf, _ = optional_package("tensorflow", min_version="2.18.0")

    class SingleLayerPerceptron:
        @warning_for_keywords()
        def __init__(
            self,
            *,
            input_shape=(28, 28),
            num_hidden=128,
            act_hidden="relu",
            dropout=0.2,
            num_out=10,
            act_out="softmax",
            optimizer="adam",
            loss="sparse_categorical_crossentropy",
        ):
            """Single Layer Perceptron with Dropout.

            Parameters
            ----------
            input_shape : tuple
                Shape of data to be trained
            num_hidden : int
                Number of nodes in hidden layer
            act_hidden : string
                Activation function used in hidden layer
            dropout : float
                Dropout ratio
            num_out : 10
                Number of nodes in output layer
            act_out : string
                Activation function used in output layer
            optimizer :  string
                Select optimizer. Default adam.
            loss : string
                Select loss function for measuring accuracy.
                Default sparse_categorical_crossentropy.

            """
            self.accuracy = None
            self.loss = None

            model = tf.keras.models.Sequential(
                [
                    tf.keras.layers.Input(shape=input_shape),
                    tf.keras.layers.Flatten(),
                    tf.keras.layers.Dense(num_hidden, activation=act_hidden),
                    tf.keras.layers.Dropout(dropout),
                    tf.keras.layers.Dense(num_out, activation=act_out),
                ]
            )

            model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])

            self.model = model

        def summary(self):
            """Get the summary of the model.

            The summary is textual and includes information about:
            The layers and their order in the model.
            The output shape of each layer.

            Returns
            -------
            summary : NoneType
                the summary of the model

            """
            return self.model.summary()

        @warning_for_keywords()
        def fit(self, x_train, y_train, *, epochs=5):
            """Train the model on train dataset.

            The fit method will train the model for a fixed
            number of epochs (iterations) on a dataset.

            Parameters
            ----------
            x_train : ndarray
                the x_train is the train dataset
            y_train : ndarray shape=(BatchSize,)
                the y_train is the labels of the train dataset
            epochs : int (Default = 5)
                the number of epochs

            Returns
            -------
            hist : object
                A History object. Its History.history attribute is a record of
                training loss values and metrics values at successive epochs

            """
            hist = self.model.fit(x_train, y_train, epochs=epochs)
            self.accuracy = hist.history["accuracy"][-1]
            self.loss = hist.history["loss"][-1]
            return hist

        @warning_for_keywords()
        def evaluate(self, x_test, y_test, *, verbose=2):
            """Evaluate the model on test dataset.

            The evaluate method will evaluate the model on a test
            dataset.

            Parameters
            ----------
            x_test : ndarray
                the x_test is the test dataset
            y_test : ndarray shape=(BatchSize,)
                the y_test is the labels of the test dataset
            verbose : int (Default = 2)
                By setting verbose 0, 1 or 2 you just say how do you want to
                'see' the training progress for each epoch.

            Returns
            -------
            evaluate : List
                return list of loss value and accuracy value on test dataset

            """
            return self.model.evaluate(x_test, y_test, verbose=verbose)

        def predict(self, x_test):
            """Predict the output from input samples.

            The predict method will generates output predictions
            for the input samples.

            Parameters
            ----------
            x_train : ndarray
                the x_test is the test dataset or input samples

            Returns
            -------
            predict : ndarray shape(TestSize,OutputSize)
                Numpy array(s) of predictions.

            """
            return self.model.predict(x_test)

    class MultipleLayerPercepton:
        @warning_for_keywords()
        def __init__(
            self,
            *,
            input_shape=(28, 28),
            num_hidden=(128,),
            act_hidden="relu",
            dropout=0.2,
            num_out=10,
            act_out="softmax",
            loss="sparse_categorical_crossentropy",
            optimizer="adam",
        ):
            """Multiple Layer Perceptron with Dropout.

            Parameters
            ----------
            input_shape : tuple
                Shape of data to be trained
            num_hidden : array-like
                List of number of nodes in hidden layers
            act_hidden : string
                Activation function used in hidden layer
            dropout : float
                Dropout ratio
            num_out : 10
                Number of nodes in output layer
            act_out : string
                Activation function used in output layer
            optimizer :  string
                Select optimizer. Default adam.
            loss : string
                Select loss function for measuring accuracy.
                Default sparse_categorical_crossentropy.

            """
            self.input_shape = input_shape
            self.num_hidden = num_hidden
            self.act_hidden = act_hidden
            self.dropout = dropout
            self.num_out = num_out
            self.act_out = act_out
            self.loss = loss
            self.optimizer = optimizer
            self.accuracy = None

            # model building

            inp = tf.keras.layers.Input(shape=self.input_shape)
            x = tf.keras.layers.Flatten()(inp)

            for i in range(len(self.num_hidden)):
                x = tf.keras.layers.Dense(self.num_hidden[i])(x)

            x = tf.keras.layers.Dropout(self.dropout)(x)
            out = tf.keras.layers.Dense(self.num_out, activation=self.act_out)(x)

            self.model = tf.keras.models.Model(inputs=inp, outputs=out)

            # compiling the model
            self.model.compile(
                optimizer=self.optimizer, loss=self.loss, metrics=["accuracy"]
            )

        def summary(self):
            """Get the summary of the model.

            The summary is textual and includes information about:
            The layers and their order in the model.
            The output shape of each layer.

            Returns
            -------
            summary :  NoneType
                the summary of the model

            """
            return self.model.summary()

        @warning_for_keywords()
        def fit(self, x_train, y_train, *, epochs=5):
            """Train the model on train dataset.

            The fit method will train the model for a fixed
            number of epochs (iterations) on a dataset.

            Parameters
            ----------
            x_train : ndarray
                the x_train is the train dataset
            y_train : ndarray shape=(BatchSize,)
                the y_train is the labels of the train dataset
            epochs : int (Default = 5)
                the number of epochs

            Returns
            -------
            hist : object
                A History object. Its History.history attribute is a record of
                training loss values and metrics values at successive epochs

            """
            hist = self.model.fit(x_train, y_train, epochs=epochs)
            self.accuracy = hist.history["accuracy"][-1]
            self.loss = hist.history["loss"][-1]
            return hist

        @warning_for_keywords()
        def evaluate(self, x_test, y_test, *, verbose=2):
            """Evaluate the model on test dataset.

            The evaluate method will evaluate the model on a test
            dataset.

            Parameters
            ----------
            x_test : ndarray
                the x_test is the test dataset
            y_test : ndarray shape=(BatchSize,)
                the y_test is the labels of the test dataset
            verbose : int (Default = 2)
                By setting verbose 0, 1 or 2 you just say how do you want to
                'see' the training progress for each epoch.

            Returns
            -------
            evaluate : List
                return list of loss value and accuracy value on test dataset

            """
            return self.model.evaluate(x_test, y_test, verbose=verbose)

        def predict(self, x_test):
            """Predict the output from input samples.

            The predict method will generates output predictions
            for the input samples.

            Parameters
            ----------
            x_train : ndarray
                the x_test is the test dataset or input samples

            Returns
            -------
            predict : ndarray shape(TestSize,OutputSize)
                Numpy array(s) of predictions.

            """
            return self.model.predict(x_test)