# Path: astroNN/nn/losses.py
# def zeros_loss(y_true, y_pred, sample_weight=None):
#     """
#     Always return zeros
# 
#     :param y_true: Ground Truth
#     :type y_true: Union(tf.Tensor, tf.Variable)
#     :param y_pred: Prediction
#     :type y_pred: Union(tf.Tensor, tf.Variable)
#     :param sample_weight: Sample weights
#     :type sample_weight: Union(tf.Tensor, tf.Variable, list)
#     :return: Zeros
#     :rtype: tf.Tensor
#     :History: 2018-May-24 - Written - Henry Leung (University of Toronto)
#     """
#     losses = tf.reduce_mean(0. * y_true + 0. * y_pred, axis=-1)
#     return weighted_loss(losses, sample_weight)
# 
# Path: astroNN/shared/nn_tools.py
# def gpu_memory_manage(ratio=True, log_device_placement=False):
#     """
#     To manage GPU memory usage, prevent Tensorflow preoccupied all the video RAM
# 
#     :param ratio: Optional, ratio of GPU memory pre-allocating to astroNN
#     :type ratio: Union[NoneType, float]
#     :param log_device_placement: whether or not log the device placement
#     :type log_device_placement: bool
#     :History:
#         | 2017-Nov-25 - Written - Henry Leung (University of Toronto)
#         | 2020-Jun-1 - Updated for tf v2
#     """
#     gpu_phy_devices = tf.config.list_physical_devices('GPU')
# 
#     general_warning_msg = f"Tensorflow has already been initialized, {inspect.currentframe().f_code.co_name}() needs " \
#                           f"to be called before any Tensorflow operation, as a result this function will have no effect"
# 
#     try:
#         if ratio:
#             for gpu in gpu_phy_devices:
#                 tf.config.experimental.set_memory_growth(gpu, True)
#         else:
#             for gpu in gpu_phy_devices:
#                 tf.config.experimental.set_memory_growth(gpu, False)
# 
#         if log_device_placement:
#             tf.debugging.set_log_device_placement(True)
#         else:
#             tf.debugging.set_log_device_placement(False)
#     except RuntimeError:
#         warnings.warn(general_warning_msg)
# 
#     return None


import unittest
import numpy as np
import numpy.testing as npt
import tensorflow as tf
from tensorflow import keras as tfk
from astroNN.nn.losses import zeros_loss
from astroNN.shared.nn_tools import gpu_memory_manage
        from astroNN.nn.layers import MCDropout
        from astroNN.nn.layers import MCGaussianDropout
        from astroNN.nn.layers import MCConcreteDropout
        from astroNN.nn.layers import MCSpatialDropout1D
        from astroNN.nn.layers import MCSpatialDropout2D
        from astroNN.nn.layers import ErrorProp
        from astroNN.nn.layers import StopGrad
        from astroNN.nn.layers import BoolMask
        from astroNN.apogee import aspcap_mask
        from astroNN.nn.layers import FastMCInference
        from astroNN.nn.layers import TensorInput
        from astroNN.nn.layers import PolyFit
        print('==========StopGrad tests==========')

        # Data preparation
        random_xdata = np.random.normal(0, 1, (100, 7514))
        random_ydata = np.random.normal(0, 1, (100, 25))

        input = Input(shape=[7514])
        output = Dense(25)(input)
        stopped_output = StopGrad(name='stopgrad', always_on=True)(output)
        model = Model(inputs=input, outputs=output)
        model_stopped = Model(inputs=input, outputs=stopped_output)
        model.compile(optimizer='adam', loss='mse')

        model_stopped.compile(optimizer='adam', loss='mse')
        # assert error because of no gradient via this layer
        self.assertRaises(ValueError, model_stopped.fit, random_xdata, random_ydata, batch_size=128, epochs=1)
        # keras.backend.get_session().run(keras.backend.gradients(output, input), feed_dict={input: random_xdata, keras.backend.learning_phase(): 0})

        x = model.predict(random_xdata)
        y = model_stopped.predict(random_xdata)
        npt.assert_almost_equal(x, y)  # make sure StopGrad does not change any result

        # # =================test weight equals================= #
        input2 = Input(shape=[7514])
        dense1 = Dense(100, name='normaldense')(input2)
        dense2 = Dense(25, name='wanted_dense')(input2)
        dense2_stopped = StopGrad(name='stopgrad', always_on=True)(dense2)
        output2 = Dense(25, name='wanted_dense2')(concatenate([dense1, dense2_stopped]))
        model2 = Model(inputs=input2, outputs=[output2, dense2])
        model2.compile(optimizer=tfk.optimizers.SGD(lr=0.1),

# Path: astroNN/nn/losses.py
# def zeros_loss(y_true, y_pred, sample_weight=None):
#     """
#     Always return zeros
# 
#     :param y_true: Ground Truth
#     :type y_true: Union(tf.Tensor, tf.Variable)
#     :param y_pred: Prediction
#     :type y_pred: Union(tf.Tensor, tf.Variable)
#     :param sample_weight: Sample weights
#     :type sample_weight: Union(tf.Tensor, tf.Variable, list)
#     :return: Zeros
#     :rtype: tf.Tensor
#     :History: 2018-May-24 - Written - Henry Leung (University of Toronto)
#     """
#     losses = tf.reduce_mean(0. * y_true + 0. * y_pred, axis=-1)
#     return weighted_loss(losses, sample_weight)
# 
# Path: astroNN/shared/nn_tools.py
# def gpu_memory_manage(ratio=True, log_device_placement=False):
#     """
#     To manage GPU memory usage, prevent Tensorflow preoccupied all the video RAM
# 
#     :param ratio: Optional, ratio of GPU memory pre-allocating to astroNN
#     :type ratio: Union[NoneType, float]
#     :param log_device_placement: whether or not log the device placement
#     :type log_device_placement: bool
#     :History:
#         | 2017-Nov-25 - Written - Henry Leung (University of Toronto)
#         | 2020-Jun-1 - Updated for tf v2
#     """
#     gpu_phy_devices = tf.config.list_physical_devices('GPU')
# 
#     general_warning_msg = f"Tensorflow has already been initialized, {inspect.currentframe().f_code.co_name}() needs " \
#                           f"to be called before any Tensorflow operation, as a result this function will have no effect"
# 
#     try:
#         if ratio:
#             for gpu in gpu_phy_devices:
#                 tf.config.experimental.set_memory_growth(gpu, True)
#         else:
#             for gpu in gpu_phy_devices:
#                 tf.config.experimental.set_memory_growth(gpu, False)
# 
#         if log_device_placement:
#             tf.debugging.set_log_device_placement(True)
#         else:
#             tf.debugging.set_log_device_placement(False)
#     except RuntimeError:
#         warnings.warn(general_warning_msg)
# 
#     return None


# Path: tests/test_layers.py
import unittest
import numpy as np
import numpy.testing as npt
import tensorflow as tf
from tensorflow import keras as tfk
from astroNN.nn.losses import zeros_loss
from astroNN.shared.nn_tools import gpu_memory_manage
        from astroNN.nn.layers import MCDropout
        from astroNN.nn.layers import MCGaussianDropout
        from astroNN.nn.layers import MCConcreteDropout
        from astroNN.nn.layers import MCSpatialDropout1D
        from astroNN.nn.layers import MCSpatialDropout2D
        from astroNN.nn.layers import ErrorProp
        from astroNN.nn.layers import StopGrad
        from astroNN.nn.layers import BoolMask
        from astroNN.apogee import aspcap_mask
        from astroNN.nn.layers import FastMCInference
        from astroNN.nn.layers import TensorInput
        from astroNN.nn.layers import PolyFit

        print('==========StopGrad tests==========')

        # Data preparation
        random_xdata = np.random.normal(0, 1, (100, 7514))
        random_ydata = np.random.normal(0, 1, (100, 25))

        input = Input(shape=[7514])
        output = Dense(25)(input)
        stopped_output = StopGrad(name='stopgrad', always_on=True)(output)
        model = Model(inputs=input, outputs=output)
        model_stopped = Model(inputs=input, outputs=stopped_output)
        model.compile(optimizer='adam', loss='mse')

        model_stopped.compile(optimizer='adam', loss='mse')
        # assert error because of no gradient via this layer
        self.assertRaises(ValueError, model_stopped.fit, random_xdata, random_ydata, batch_size=128, epochs=1)
        # keras.backend.get_session().run(keras.backend.gradients(output, input), feed_dict={input: random_xdata, keras.backend.learning_phase(): 0})

        x = model.predict(random_xdata)
        y = model_stopped.predict(random_xdata)
        npt.assert_almost_equal(x, y)  # make sure StopGrad does not change any result

        # # =================test weight equals================= #
        input2 = Input(shape=[7514])
        dense1 = Dense(100, name='normaldense')(input2)
        dense2 = Dense(25, name='wanted_dense')(input2)
        dense2_stopped = StopGrad(name='stopgrad', always_on=True)(dense2)
        output2 = Dense(25, name='wanted_dense2')(concatenate([dense1, dense2_stopped]))
        model2 = Model(inputs=input2, outputs=[output2, dense2])
        model2.compile(optimizer=tfk.optimizers.SGD(lr=0.1),

                       loss={'wanted_dense2': 'mse', 'wanted_dense': zeros_loss})