"""Tests for TileDB Tensorflow Keras model save and load."""

import glob
import io
import os
import pickle
import platform
import shutil

import numpy as np
import pytest
import tensorflow as tf

import tiledb
from tiledb.ml import __version__ as tiledb_ml_version
from tiledb.ml.models.tensorflow_keras import TensorflowKerasTileDBModel

try:
    from keras.testing_infra.test_utils import (
        get_small_functional_mlp,
        get_small_sequential_mlp,
    )
except ImportError:
    from keras.testing_utils import get_small_functional_mlp, get_small_sequential_mlp

# Suppress all Tensorflow messages
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

batch_get_value = tf.keras.backend.batch_get_value


def read_files(dirpath):
    files = {}
    for path in glob.glob(f"{dirpath}/*"):
        with open(path, "rb") as f:
            files[path] = f.read()
    return files


class ConfigSubclassModel(tf.keras.Model):
    def __init__(self, hidden_units):
        super(ConfigSubclassModel, self).__init__()
        self.hidden_units = hidden_units
        self.dense_layers = [tf.keras.layers.Dense(u) for u in hidden_units]

    def call(self, inputs):
        x = inputs
        for layer in self.dense_layers:
            x = layer(x)
        return x

    def get_config(self):
        return {"hidden_units": self.hidden_units}

    @classmethod
    def from_config(cls, config):
        return cls(**config)


def test_load_tiledb_error_with_wrong_uri():
    tiledb_model_obj = TensorflowKerasTileDBModel(uri="dummy_uri")
    with pytest.raises(tiledb.TileDBError):
        tiledb_model_obj.load(compile_model=False)


loss_optimizer_metrics = pytest.mark.parametrize(
    "loss,optimizer,metrics",
    [
        (None, None, None),
        ("binary_crossentropy", "rmsprop", "accuracy"),
        ("mean_squared_error", "rmsprop", "categorical_accuracy"),
        ("sparse_categorical_crossentropy", "sgd", "sparse_categorical_crossentropy"),
    ],
)

api = pytest.mark.parametrize(
    "api",
    [get_small_sequential_mlp, get_small_functional_mlp],
)


class TestTensorflowKerasModel:
    @loss_optimizer_metrics
    def test_exception_on_subclassed_models(self, tmpdir, loss, optimizer, metrics):
        model = ConfigSubclassModel(hidden_units=[16, 16, 10])
        tiledb_uri = os.path.join(tmpdir, "model_array")

        # Compiles the model if optimizer is present
        if optimizer:
            model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])

        if not model.built:
            model.build(tuple(np.random.randint(20, size=2)))
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        with pytest.raises(RuntimeError):
            tiledb_model_obj.save(include_optimizer=True if optimizer else False)

    @api
    @loss_optimizer_metrics
    def test_save_model_to_tiledb_array(self, tmpdir, api, loss, optimizer, metrics):
        model = api(num_hidden=1, num_classes=2, input_dim=3)
        tiledb_uri = os.path.join(tmpdir, "model_array")

        # Compiles the model if optimizer is present
        if optimizer:
            model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])

        if not model.built:
            model.build(tuple(np.random.randint(20, size=2)))
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=True if optimizer else False)
        assert tiledb.array_exists(tiledb_uri)

        with pytest.raises(RuntimeError) as ex:
            tiledb_obj = TensorflowKerasTileDBModel(uri="")
            tiledb_obj.save()

        assert "Model is not initialized" in str(ex.value)

    @api
    @loss_optimizer_metrics
    def test_save_model_to_tiledb_array_predictions(
        self, tmpdir, api, loss, optimizer, metrics
    ):
        model = api(num_hidden=1, num_classes=2, input_dim=3)
        tiledb_uri = os.path.join(tmpdir, "model_array")

        # Compiles the model if optimizer is present
        if optimizer:
            model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])

        input_shape = tuple(np.random.randint(20, size=2))
        if not model.built:
            model.build(input_shape)
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)

        tiledb_model_obj.save(include_optimizer=True if optimizer else False)

        loaded_model = tiledb_model_obj.load(compile_model=False)

        data = np.random.rand(100, 3)

        # Assert model predictions are equal
        np.testing.assert_array_equal(loaded_model.predict(data), model.predict(data))

    @api
    @loss_optimizer_metrics
    def test_save_model_to_tiledb_array_weights(
        self, tmpdir, api, loss, optimizer, metrics
    ):
        model = api(num_hidden=1, num_classes=2, input_dim=3)
        tiledb_uri = os.path.join(tmpdir, "model_array")

        # Compiles the model if optimizer is present
        if optimizer:
            model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])

        input_shape = tuple(np.random.randint(20, size=2))
        if not model.built:
            model.build(input_shape)
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)

        tiledb_model_obj.save(include_optimizer=True if optimizer else False)
        loaded_model = tiledb_model_obj.load(compile_model=True if optimizer else False)

        data = np.random.rand(100, 3)

        if optimizer:
            model_opt_weights = batch_get_value(model.optimizer.weights)
            loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights)

            # Assert optimizer weights are equal
            for weight_model, weight_loaded_model in zip(
                model_opt_weights, loaded_opt_weights
            ):
                np.testing.assert_array_equal(weight_model, weight_loaded_model)

            # Assert model predictions are equal
            np.testing.assert_array_equal(
                loaded_model.predict(data), model.predict(data)
            )

    @loss_optimizer_metrics
    def test_save_load_with_dense_features(self, tmpdir, loss, optimizer, metrics):
        if optimizer is None:
            pytest.skip()
        cols = [
            tf.feature_column.numeric_column("a"),
            tf.feature_column.indicator_column(
                tf.feature_column.categorical_column_with_vocabulary_list(
                    "b", ["one", "two"]
                )
            ),
        ]
        input_layers = {
            "a": tf.keras.layers.Input(shape=(1,), name="a"),
            "b": tf.keras.layers.Input(shape=(1,), name="b", dtype="string"),
        }

        fc_layer = tf.keras.layers.DenseFeatures(cols)(input_layers)
        output = tf.keras.layers.Dense(10)(fc_layer)

        model = tf.keras.Model(input_layers, output)

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

        tiledb_uri = os.path.join(tmpdir, "model_array")
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=True)
        loaded_model = tiledb_model_obj.load(compile_model=True)

        model_opt_weights = batch_get_value(model.optimizer.weights)
        loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights)

        # Assert optimizer weights are equal
        for weight_model, weight_loaded_model in zip(
            model_opt_weights, loaded_opt_weights
        ):
            np.testing.assert_array_equal(weight_model, weight_loaded_model)

        inputs_a = np.arange(10).reshape(10, 1)
        inputs_b = np.arange(10).reshape(10, 1).astype("str")

        # Assert model predictions are equal
        np.testing.assert_array_equal(
            loaded_model.predict({"a": inputs_a, "b": inputs_b}),
            model.predict({"a": inputs_a, "b": inputs_b}),
        )

    @loss_optimizer_metrics
    def test_save_load_with_sequence_features(self, tmpdir, loss, optimizer, metrics):
        if optimizer is None:
            pytest.skip()

        cols = [
            tf.feature_column.sequence_numeric_column("a"),
            tf.feature_column.indicator_column(
                tf.feature_column.sequence_categorical_column_with_vocabulary_list(
                    "b", ["one", "two"]
                )
            ),
        ]
        input_layers = {
            "a": tf.keras.layers.Input(shape=(None, 1), sparse=True, name="a"),
            "b": tf.keras.layers.Input(
                shape=(None, 1), sparse=True, name="b", dtype="string"
            ),
        }

        fc_layer, _ = tf.keras.experimental.SequenceFeatures(cols)(input_layers)
        x = tf.keras.layers.GRU(32)(fc_layer)
        output = tf.keras.layers.Dense(10)(x)

        model = tf.keras.Model(input_layers, output)

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

        tiledb_uri = os.path.join(tmpdir, "model_array")
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=True)
        loaded_model = tiledb_model_obj.load(compile_model=True)

        model_opt_weights = batch_get_value(model.optimizer.weights)
        loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights)

        # Assert optimizer weights are equal
        for weight_model, weight_loaded_model in zip(
            model_opt_weights, loaded_opt_weights
        ):
            np.testing.assert_array_equal(weight_model, weight_loaded_model)

        batch_size = 10
        timesteps = 1

        values_a = np.arange(10, dtype=np.float32)
        indices_a = np.zeros((10, 3), dtype=np.int64)
        indices_a[:, 0] = np.arange(10)
        inputs_a = tf.SparseTensor(indices_a, values_a, (batch_size, timesteps, 1))

        values_b = np.zeros(10, dtype=np.str)
        indices_b = np.zeros((10, 3), dtype=np.int64)
        indices_b[:, 0] = np.arange(10)
        inputs_b = tf.SparseTensor(indices_b, values_b, (batch_size, timesteps, 1))

        # Assert model predictions are equal
        np.testing.assert_array_equal(
            loaded_model.predict({"a": inputs_a, "b": inputs_b}, steps=1),
            model.predict({"a": inputs_a, "b": inputs_b}, steps=1),
        )

    def test_functional_model_save_load_with_custom_loss_and_metric(self, tmpdir):
        inputs = tf.keras.layers.Input(shape=(4,))
        x = tf.keras.layers.Dense(8, activation="relu")(inputs)
        outputs = tf.keras.layers.Dense(3, activation="softmax")(x)
        model = tf.keras.Model(inputs=inputs, outputs=outputs)
        custom_loss = tf.keras.layers.Lambda(lambda x: tf.keras.backend.sum(x * x))(x)
        model.add_loss(custom_loss)
        model.add_metric(custom_loss, aggregation="mean", name="custom_loss")

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

        data_x = np.random.normal(size=(32, 4))
        data_y = np.random.randint(0, 3, size=32)
        model.train_on_batch(data_x, data_y)

        tiledb_uri = os.path.join(tmpdir, "model_array")
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=True)
        loaded_model = tiledb_model_obj.load(compile_model=True)

        # Assert all evaluation results are the same.
        assert all(
            [
                a == pytest.approx(b, 1e-9)
                for a, b in zip(
                    model.evaluate(data_x, data_y),
                    loaded_model.evaluate(data_x, data_y),
                )
            ]
        )

        # Assert model predictions are equal
        np.testing.assert_array_equal(
            loaded_model.predict(data_x), model.predict(data_x)
        )

    def test_save_load_for_rnn_layers(self, tmpdir):
        inputs = tf.keras.layers.Input([10, 10], name="train_input")
        rnn_layers = [
            tf.keras.layers.LSTMCell(size, recurrent_dropout=0, name="rnn_cell%d" % i)
            for i, size in enumerate([32, 32])
        ]
        rnn_output = tf.keras.layers.RNN(
            rnn_layers, return_sequences=True, name="rnn_layer"
        )(inputs)
        pred_feat = tf.keras.layers.Dense(10, name="prediction_features")(rnn_output)
        pred = tf.keras.layers.Softmax()(pred_feat)
        model = tf.keras.Model(inputs=[inputs], outputs=[pred, pred_feat])

        tiledb_uri = os.path.join(tmpdir, "model_array")
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=False)
        loaded_model = tiledb_model_obj.load(compile_model=False)

        data = np.random.rand(50, 10, 10)

        # Assert model predictions are equal
        np.testing.assert_array_equal(loaded_model.predict(data), model.predict(data))

    def test_sequential_model_save_load_without_input_shape(self, tmpdir):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Dense(2))
        model.add(tf.keras.layers.RepeatVector(3))
        model.add(tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(3)))
        model.compile(
            loss="mean_squared_error",
            optimizer="rmsprop",
            metrics="categorical_accuracy",
            weighted_metrics="categorical_accuracy",
            sample_weight_mode="temporal",
        )
        data_x = np.random.random((1, 3))
        data_y = np.random.random((1, 3, 3))
        model.train_on_batch(data_x, data_y)

        tiledb_uri = os.path.join(tmpdir, "model_array")
        tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
        tiledb_model_obj.save(include_optimizer=True)
        loaded_model = tiledb_model_obj.load(compile_model=True)

        # Assert model predictions are equal
        np.testing.assert_array_equal(
            loaded_model.predict(data_x), model.predict(data_x)
        )

    @api
    @loss_optimizer_metrics
    def test_preview(self, tmpdir, api, loss, optimizer, metrics):
        model = api(num_hidden=1, num_classes=2, input_dim=3)
        tiledb_uri = os.path.join(tmpdir, "model_array")

        # Compiles the model if optimizer is present
        if optimizer:
            model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])

        # With model given as argument
        if model.built:
            tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model)
            s = io.StringIO()
            model.summary(print_fn=lambda x: s.write(x + "\n"))
            model_summary = s.getvalue()
            assert tiledb_model_obj.preview() == model_summary
        else:
            # Model should be built before preview it
            with pytest.raises(ValueError):
                tiledb_model_obj = TensorflowKerasTileDBModel(
                    uri=tiledb_uri, model=model
                )
                tiledb_model_obj.preview()

        # When model is None then preview returns empty string
        tiledb_model_obj_none = TensorflowKerasTileDBModel(uri=tiledb_uri, model=None)
        assert tiledb_model_obj_none.preview() == ""


class TestTensorflowKerasModelCloud:
    def test_file_properties(self, tmpdir):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))

        # Get model summary in a string
        s = io.StringIO()
        model.summary(print_fn=lambda x: s.write(x + "\n"))
        model_summary = s.getvalue()

        uri = os.path.join(tmpdir, "model_array")
        tiledb_obj = TensorflowKerasTileDBModel(uri=uri, model=model)

        assert (
            tiledb_obj._file_properties["TILEDB_ML_MODEL_ML_FRAMEWORK"]
            == "TENSORFLOW KERAS"
        )
        assert tiledb_obj._file_properties["TILEDB_ML_MODEL_STAGE"] == "STAGING"
        assert (
            tiledb_obj._file_properties["TILEDB_ML_MODEL_PYTHON_VERSION"]
            == platform.python_version()
        )
        assert (
            tiledb_obj._file_properties["TILEDB_ML_MODEL_ML_FRAMEWORK_VERSION"]
            == tf.__version__
        )
        assert tiledb_obj._file_properties["TILEDB_ML_MODEL_PREVIEW"] == model_summary
        assert (
            tiledb_obj._file_properties["TILEDB_ML_MODEL_VERSION"] == tiledb_ml_version
        )

    def test_get_cloud_uri_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))
        uri = os.path.join(tmpdir, "model_array")

        mock_get_cloud_uri = mocker.patch(
            "tiledb.ml.models._base.get_cloud_uri", return_value=uri
        )

        _ = TensorflowKerasTileDBModel(uri=uri, namespace="test_namespace", model=model)

        mock_get_cloud_uri.assert_called_once_with(uri, "test_namespace")

    def test_get_s3_prefix_call_for_models_on_tiledb_cloud(self, tmpdir, mocker):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))
        uri = os.path.join(tmpdir, "model_array")

        mock_get_s3_prefix = mocker.patch(
            "tiledb.ml.models._cloud_utils.get_s3_prefix", return_value="s3 prefix"
        )

        _ = TensorflowKerasTileDBModel(uri=uri, namespace="test_namespace", model=model)

        mock_get_s3_prefix.assert_called_once_with("test_namespace")

    def test_update_file_properties_call(self, tmpdir, mocker):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))

        # Get model summary in a string
        s = io.StringIO()
        model.summary(print_fn=lambda x: s.write(x + "\n"))
        model_summary = s.getvalue()

        uri = os.path.join(tmpdir, "model_array")

        mocker.patch("tiledb.ml.models._base.get_cloud_uri", return_value=uri)

        tiledb_obj = TensorflowKerasTileDBModel(
            uri=uri, namespace="test_namespace", model=model
        )

        mock_update_file_properties = mocker.patch(
            "tiledb.ml.models._base.update_file_properties",
            return_value=None,
        )
        mocker.patch(
            "tiledb.ml.models.tensorflow_keras.TensorflowKerasTileDBModel._write_array"
        )

        tiledb_obj.save()

        file_properties_dict = {
            "TILEDB_ML_MODEL_ML_FRAMEWORK": "TENSORFLOW KERAS",
            "TILEDB_ML_MODEL_ML_FRAMEWORK_VERSION": tf.__version__,
            "TILEDB_ML_MODEL_STAGE": "STAGING",
            "TILEDB_ML_MODEL_PYTHON_VERSION": platform.python_version(),
            "TILEDB_ML_MODEL_PREVIEW": model_summary,
            "TILEDB_ML_MODEL_VERSION": tiledb_ml_version,
        }

        mock_update_file_properties.assert_called_once_with(uri, file_properties_dict)

    def test_exception_raise_file_property_in_meta_error(self, tmpdir):
        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))
        tiledb_array = os.path.join(tmpdir, "model_array")
        tiledb_obj = TensorflowKerasTileDBModel(uri=tiledb_array, model=model)
        with pytest.raises(ValueError) as ex:
            tiledb_obj.save(
                meta={"TILEDB_ML_MODEL_ML_FRAMEWORK": "TILEDB_ML_MODEL_ML_FRAMEWORK"},
            )

        assert "Please avoid using file property key names as metadata keys!" in str(
            ex.value
        )

    def test_tensorboard_callback(self, mocker, tmpdir):
        mocker.patch(
            "tiledb.ml.models.tensorflow_keras.TensorflowKerasTileDBModel._use_legacy_schema",
            return_value=False,
        )

        model = tf.keras.Sequential()
        model.add(tf.keras.layers.Flatten(input_shape=(10, 10)))
        tiledb_array = os.path.join(tmpdir, "model_array")
        tiledb_obj = TensorflowKerasTileDBModel(uri=tiledb_array, model=model)

        cb = [tf.keras.callbacks.TensorBoard(log_dir=tmpdir)]

        os.mkdir(os.path.join(tmpdir, "train"))
        with open(os.path.join(tmpdir, "train", "foo_tfevents_1"), "wb") as f:
            f.write(b"test_bytes_1")
        with open(os.path.join(tmpdir, "train", "bar_tfevents_2"), "wb") as f:
            f.write(b"test_bytes_2")

        log_files = read_files(os.path.join(tmpdir, "train"))
        assert log_files

        tiledb_obj.save(callbacks=cb)
        with tiledb.open(tiledb_array) as A:
            tb_size = A.meta["tensorboard_size"]
            assert pickle.loads(A[0:tb_size]["tensorboard"]) == {
                os.path.join(tmpdir, "train", "foo_tfevents_1"): b"test_bytes_1",
                os.path.join(tmpdir, "train", "bar_tfevents_2"): b"test_bytes_2",
            }
        shutil.rmtree(os.path.join(tmpdir, "train"))

        tiledb_obj.load(compile_model=False, callback=True)
        with open(os.path.join(tmpdir, "train", "foo_tfevents_1"), "rb") as f:
            assert f.read() == b"test_bytes_1"
        with open(os.path.join(tmpdir, "train", "bar_tfevents_2"), "rb") as f:
            assert f.read() == b"test_bytes_2"

        shutil.rmtree(os.path.join(tmpdir, "train"))
