import logging
import sys

import pytest

import ray
from ray import train, tune
from ray.air.constants import TRAINING_ITERATION
from ray.train._internal.worker_group import WorkerGroup
from ray.train.backend import Backend, BackendConfig
from ray.train.data_parallel_trainer import DataParallelTrainer
from ray.train.examples.pytorch.torch_fashion_mnist_example import (
    train_func_per_worker as fashion_mnist_train_func,
)
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
from ray.train.torch import TorchTrainer
from ray.tune.tune_config import TuneConfig
from ray.tune.tuner import Tuner


@pytest.fixture(scope="module")
def ray_start_4_cpus():
    address_info = ray.init(num_cpus=4)
    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()


@pytest.fixture
def ray_start_8_cpus():
    address_info = ray.init(num_cpus=8)
    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()


class TestConfig(BackendConfig):
    @property
    def backend_cls(self):
        return TestBackend


class TestBackend(Backend):
    def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
        pass

    def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
        pass


def torch_fashion_mnist(num_workers, use_gpu, num_samples):
    trainer = TorchTrainer(
        fashion_mnist_train_func,
        scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
    )
    tuner = Tuner(
        trainer,
        param_space={
            "train_loop_config": {
                "lr": tune.loguniform(1e-4, 1e-1),
                "batch_size_per_worker": tune.choice([32, 64, 128]),
                "epochs": 2,
            }
        },
        tune_config=TuneConfig(
            num_samples=num_samples,
        ),
    )
    analysis = tuner.fit()._experiment_analysis

    # Check that loss decreases in each trial.
    for df in analysis.trial_dataframes.values():
        assert df.loc[1, "loss"] < df.loc[0, "loss"]


def test_tune_torch_fashion_mnist(ray_start_8_cpus):
    torch_fashion_mnist(num_workers=2, use_gpu=False, num_samples=2)


@pytest.mark.skipif(
    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
)
def tune_tensorflow_mnist(num_workers, use_gpu, num_samples):
    from ray.train.examples.tf.tensorflow_mnist_example import (
        train_func as tensorflow_mnist_train_func,
    )
    from ray.train.tensorflow import TensorflowTrainer

    trainer = TensorflowTrainer(
        tensorflow_mnist_train_func,
        scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
    )
    tuner = Tuner(
        trainer,
        param_space={
            "train_loop_config": {
                "lr": tune.loguniform(1e-4, 1e-1),
                "batch_size": tune.choice([32, 64, 128]),
                "epochs": 2,
            }
        },
        tune_config=TuneConfig(
            num_samples=num_samples,
        ),
    )
    analysis = tuner.fit()._experiment_analysis

    # Check that loss decreases in each trial.
    for df in analysis.trial_dataframes.values():
        assert df.loc[1, "loss"] < df.loc[0, "loss"]


@pytest.mark.skipif(
    sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
)
def test_tune_tensorflow_mnist(ray_start_8_cpus):
    tune_tensorflow_mnist(num_workers=2, use_gpu=False, num_samples=2)


def test_tune_error(ray_start_4_cpus):
    def train_func(config):
        raise RuntimeError("Error in training function!")

    trainer = DataParallelTrainer(
        train_func,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
    )
    tuner = Tuner(
        trainer,
    )

    result_grid = tuner.fit()
    with pytest.raises(RuntimeError):
        raise result_grid[0].error


def test_tune_checkpoint(ray_start_4_cpus):
    def train_func():
        for i in range(9):
            train.report(dict(test=i))
        with create_dict_checkpoint(dict(hello="world")) as checkpoint:
            train.report(dict(test=i + 1), checkpoint=checkpoint)

    trainer = DataParallelTrainer(
        train_func,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
    )
    tuner = Tuner(
        trainer,
        param_space={"train_loop_config": {"max_iter": 5}},
    )

    result_grid = tuner.fit()
    assert len(result_grid) == 1
    result = result_grid[0]
    assert result.checkpoint
    assert load_dict_checkpoint(result.checkpoint)["hello"] == "world"


def test_reuse_checkpoint(ray_start_4_cpus):
    def train_func(config):
        itr = 0
        ckpt = train.get_checkpoint()
        if ckpt is not None:
            ckpt = load_dict_checkpoint(ckpt)
            itr = ckpt["iter"] + 1

        for i in range(itr, config["max_iter"]):
            with create_dict_checkpoint(dict(iter=i)) as checkpoint:
                train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)

    trainer = DataParallelTrainer(
        train_func,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
    )
    tuner = Tuner(
        trainer,
        param_space={"train_loop_config": {"max_iter": 5}},
    )
    result_grid = tuner.fit()
    assert len(result_grid) == 1
    result = result_grid[0]
    assert result.checkpoint
    assert load_dict_checkpoint(result.checkpoint)["iter"] == 4

    tuner = Tuner.restore(result_grid.experiment_path, trainable=trainer)
    result_grid = tuner.fit()
    assert len(result_grid) == 1
    assert len(result_grid[0].metrics_dataframe) == 5


def test_retry_with_max_failures(ray_start_4_cpus):
    """Tests trainer retry with max_failures > 0 when integrating with Tune."""

    def train_func():
        ckpt = train.get_checkpoint()
        restored = bool(ckpt)  # Does a previous checkpoint exist?
        itr = 0
        if ckpt:
            ckpt = load_dict_checkpoint(ckpt)
            itr = ckpt["iter"] + 1

        for i in range(itr, 4):
            if i == 2 and not restored:
                raise Exception("try to fail me")
            with create_dict_checkpoint(dict(iter=i)) as checkpoint:
                train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)

    trainer = DataParallelTrainer(
        train_func,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
    )
    tuner = Tuner(
        trainer,
        run_config=tune.RunConfig(failure_config=tune.FailureConfig(max_failures=3)),
    )

    result_grid = tuner.fit()
    checkpoint = load_dict_checkpoint(result_grid[0].checkpoint)
    assert checkpoint["iter"] == 3
    df = result_grid[0].metrics_dataframe
    assert len(df[TRAINING_ITERATION]) == 4


def test_restore_with_new_trainer(ray_start_4_cpus, tmpdir, propagate_logs, caplog):
    def train_func(config):
        raise RuntimeError("failing!")

    trainer = DataParallelTrainer(
        train_func,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
        run_config=train.RunConfig(
            name="restore_new_trainer", storage_path=str(tmpdir)
        ),
        datasets={"train": ray.data.from_items([{"a": i} for i in range(10)])},
    )
    results = Tuner(trainer).fit()
    assert results.errors

    def train_func(config):
        dataset = train.get_dataset_shard("train")
        assert train.get_context().get_world_size() == 2
        rows = 0
        for _ in dataset.iter_rows():
            rows += 1
        assert rows == 10

    trainer = DataParallelTrainer(
        # Training function can be modified
        train_func,
        backend_config=TestConfig(),
        # ScalingConfig can be modified
        scaling_config=train.ScalingConfig(num_workers=2),
        # New RunConfig will be ignored
        run_config=train.RunConfig(name="ignored"),
        # Datasets and preprocessors can be re-specified
        datasets={"train": ray.data.from_items([{"a": i} for i in range(20)])},
    )
    caplog.clear()
    with caplog.at_level(logging.WARNING, logger="ray.tune.impl.tuner_internal"):
        tuner = Tuner.restore(
            str(tmpdir / "restore_new_trainer"),
            trainable=trainer,
            resume_errored=True,
        )
        assert "they will be ignored in the resumed run" in caplog.text

    results = tuner.fit()
    assert not results.errors


@pytest.mark.parametrize("in_trainer", [True, False])
@pytest.mark.parametrize("in_tuner", [True, False])
def test_run_config_in_trainer_and_tuner(
    propagate_logs, tmp_path, caplog, in_trainer, in_tuner
):
    trainer_run_config = (
        train.RunConfig(name="trainer", storage_path=str(tmp_path))
        if in_trainer
        else None
    )
    tuner_run_config = (
        tune.RunConfig(name="tuner", storage_path=str(tmp_path)) if in_tuner else None
    )
    trainer = DataParallelTrainer(
        lambda config: None,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
        run_config=trainer_run_config,
    )
    with caplog.at_level(logging.INFO, logger="ray.tune.impl.tuner_internal"):
        tuner = Tuner(trainer, run_config=tuner_run_config)

    both_msg = (
        "`RunConfig` was passed to both the `Tuner` and the `DataParallelTrainer`"
    )
    run_config = tuner._local_tuner.get_run_config()
    if in_trainer and in_tuner:
        assert run_config.name == "tuner"
        assert both_msg in caplog.text
    elif in_trainer and not in_tuner:
        assert run_config.name == "trainer"
        assert both_msg not in caplog.text
    elif not in_trainer and in_tuner:
        assert run_config.name == "tuner"
        assert both_msg not in caplog.text
    else:
        assert both_msg not in caplog.text


def test_run_config_in_param_space():
    trainer = DataParallelTrainer(
        lambda config: None,
        backend_config=TestConfig(),
        scaling_config=train.ScalingConfig(num_workers=1),
    )
    with pytest.raises(ValueError):
        Tuner(trainer, param_space={"run_config": train.RunConfig(name="ignored")})


if __name__ == "__main__":
    import sys

    import pytest

    sys.exit(pytest.main(["-v", "-x", __file__]))