# -*- coding: utf-8 -*-
# Copyright (c) 2019 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# This file is copied and adapted from
# https://github.com/ludwig-ai/ludwig/blob/master/tests/integration_tests/test_ray.py

import contextlib
import os
import sys
import tempfile

import pytest

import ray

ludwig_installed = True
tf_installed = True

try:
    import ludwig  # noqa: F401
except (ImportError, ModuleNotFoundError):
    ludwig_installed = False

try:
    import tensorflow as tf  # noqa: F401
except (ImportError, ModuleNotFoundError):
    tf_installed = False

skip = not ludwig_installed or not tf_installed

# These tests are written for versions of Modin that require python 3.7+
pytestmark = pytest.mark.skipif(skip, reason="Missing Ludwig dependency")

if not skip:
    from ludwig.backend.ray import RayBackend, get_horovod_kwargs

    from ray.tests.ludwig.ludwig_test_utils import (
        bag_feature,
        binary_feature,
        category_feature,
        create_data_set_to_use,
        date_feature,
        generate_data,
        h3_feature,
        numerical_feature,
        sequence_feature,
        set_feature,
        spawn,
        train_with_backend,
        vector_feature,
    )

else:

    def spawn(func):
        return func


@contextlib.contextmanager
def ray_start_2_cpus():
    is_ray_initialized = ray.is_initialized()
    with tempfile.TemporaryDirectory() as tmpdir:
        if not is_ray_initialized:
            res = ray.init(
                num_cpus=2,
                include_dashboard=False,
                object_store_memory=150 * 1024 * 1024,
                _temp_dir=tmpdir,
            )
        else:
            res = None
        try:
            yield res
        finally:
            if not is_ray_initialized:
                ray.shutdown()


def run_api_experiment(config, data_parquet):
    # Sanity check that we get 4 slots over 1 host
    kwargs = get_horovod_kwargs()
    assert kwargs.get("num_workers") == 2

    # Train on Parquet
    dask_backend = RayBackend()
    assert train_with_backend(
        dask_backend, config, dataset=data_parquet, evaluate=False
    )


@spawn
def run_test_parquet(
    input_features,
    output_features,
    num_examples=100,
    run_fn=run_api_experiment,
    expect_error=False,
):
    tf.config.experimental_run_functions_eagerly(True)
    with ray_start_2_cpus():
        config = {
            "input_features": input_features,
            "output_features": output_features,
            "combiner": {"type": "concat", "fc_size": 14},
            "training": {"epochs": 2, "batch_size": 8},
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            csv_filename = os.path.join(tmpdir, "dataset.csv")
            dataset_csv = generate_data(
                input_features, output_features, csv_filename, num_examples=num_examples
            )
            dataset_parquet = create_data_set_to_use("parquet", dataset_csv)

            if expect_error:
                with pytest.raises(ValueError):
                    run_fn(config, data_parquet=dataset_parquet)
            else:
                run_fn(config, data_parquet=dataset_parquet)


def test_ray_tabular():
    input_features = [
        sequence_feature(reduce_output="sum"),
        numerical_feature(normalization="zscore"),
        set_feature(),
        binary_feature(),
        bag_feature(),
        vector_feature(),
        h3_feature(),
        date_feature(),
    ]
    output_features = [
        category_feature(vocab_size=2, reduce_input="sum"),
        binary_feature(),
        set_feature(max_len=3, vocab_size=5),
        numerical_feature(normalization="zscore"),
        vector_feature(),
    ]
    run_test_parquet(input_features, output_features)


def test_ray_tabular_client():
    from ray.util.client.ray_client_helpers import ray_start_client_server

    with ray_start_2_cpus():
        assert not ray.util.client.ray.is_connected()
        with ray_start_client_server():
            assert ray.util.client.ray.is_connected()

            test_ray_tabular()


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "-x", __file__]))