import os
if "tf" in os.environ["DELIRA_BACKEND"]:
    from delira.io import tf_load_checkpoint, tf_save_checkpoint
    from delira.models import AbstractTfNetwork
    from delira.training.train_utils import initialize_uninitialized
    import tensorflow as tf
    import numpy as np

    class DummyNetwork(AbstractTfNetwork):
        def __init__(self, in_channels, n_outputs):
            super().__init__(in_channels=in_channels, n_outputs=n_outputs)
            self.net = self._build_model(in_channels, n_outputs)

        @staticmethod
        def _build_model(in_channels, n_outputs):
            return tf.keras.models.Sequential(
                layers=[
                    tf.keras.layers.Dense(64, input_shape=in_channels, bias_initializer='glorot_uniform'),
                    tf.keras.layers.ReLU(),
                    tf.keras.layers.Dense(n_outputs, bias_initializer='glorot_uniform')]
            )

    def test_load_save():
        
        net = DummyNetwork((32,), 1)
        initialize_uninitialized(net._sess)

        vars_1 = net._sess.run(tf.global_variables())

        tf_save_checkpoint("./model", model=net)

        net._sess.run(tf.initializers.global_variables())

        vars_2 = net._sess.run(tf.global_variables())

        tf_load_checkpoint("./model", model=net)

        vars_3 = net._sess.run(tf.global_variables())

        for var_1, var_2 in zip(vars_1, vars_2):
            assert np.all(var_1 != var_2)

        for var_1, var_3 in zip(vars_1, vars_3):
            assert np.all(var_1 == var_3)

    if __name__ == '__main__':
        test_load_save()