import csv
import glob
import json
import os
import shutil
import sys
import tempfile
import unittest
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import pytest

import ray
from ray.air.constants import (
    EXPR_PARAM_FILE,
    EXPR_PARAM_PICKLE_FILE,
    EXPR_PROGRESS_FILE,
    EXPR_RESULT_FILE,
)
from ray.cloudpickle import cloudpickle
from ray.tune import Checkpoint
from ray.tune.logger import (
    CSVLoggerCallback,
    JsonLoggerCallback,
    TBXLoggerCallback,
)
from ray.tune.logger.aim import AimLoggerCallback
from ray.tune.utils import flatten_dict


@dataclass
class Trial:
    evaluated_params: dict
    trial_id: str
    logdir: str
    experiment_path: Optional[str] = None
    experiment_dir_name: Optional[str] = None
    path: Optional[str] = None
    checkpoint: Optional[Checkpoint] = None

    @property
    def config(self):
        return self.evaluated_params

    def init_local_path(self):
        return

    @property
    def local_path(self):
        if self.logdir:
            return self.logdir
        if not self.experiment_dir_name:
            return None
        return str(Path(self.experiment_path) / self.experiment_dir_name)

    @property
    def local_experiment_path(self):
        return self.experiment_path

    def __hash__(self):
        return hash(self.trial_id)

    def get_ray_actor_ip(self) -> str:
        return ray.util.get_node_ip_address()


def result(t, rew, **kwargs):
    results = dict(
        time_total_s=t,
        episode_reward_mean=rew,
        mean_accuracy=rew * 2,
        training_iteration=int(t),
    )
    results.update(kwargs)
    return results


class LoggerSuite(unittest.TestCase):
    """Test built-in loggers."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.test_dir, ignore_errors=True)

    def testCSV(self):
        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
        t = Trial(evaluated_params=config, trial_id="csv", logdir=self.test_dir)
        logger = CSVLoggerCallback()
        logger.on_trial_result(0, [], t, result(0, 4))
        logger.on_trial_result(1, [], t, result(1, 5))
        logger.on_trial_result(
            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
        )

        logger.on_trial_complete(3, [], t)
        self._validate_csv_result()

    def testCSVEmptyHeader(self):
        """Test that starting a trial twice does not lead to empty CSV headers.

        In a previous bug, the CSV header was sometimes missing when a trial
        crashed before reporting results. See
        https://github.com/ray-project/ray/issues/15106
        """
        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
        t = Trial(evaluated_params=config, trial_id="csv", logdir=self.test_dir)
        logger = CSVLoggerCallback()
        logger.on_trial_start(0, [], t)
        logger.on_trial_start(0, [], t)
        logger.on_trial_result(1, [], t, result(1, 5))

        with open(os.path.join(self.test_dir, "progress.csv"), "rt") as f:
            csv_contents = f.read()

        csv_lines = csv_contents.split("\n")

        # Assert header has been written to progress.csv
        assert "training_iteration" in csv_lines[0]

    def _validate_csv_result(self):
        results = []
        result_file = os.path.join(self.test_dir, EXPR_PROGRESS_FILE)
        with open(result_file, "rt") as fp:
            reader = csv.DictReader(fp)
            for row in reader:
                results.append(row)

        self.assertEqual(len(results), 3)
        self.assertSequenceEqual(
            [int(row["episode_reward_mean"]) for row in results], [4, 5, 6]
        )

    def testJSON(self):
        config = {"a": 2, "b": 5, "c": {"c": {"D": 123}, "e": None}}
        t = Trial(evaluated_params=config, trial_id="json", logdir=self.test_dir)
        logger = JsonLoggerCallback()
        logger.on_trial_result(0, [], t, result(0, 4))
        logger.on_trial_result(1, [], t, result(1, 5))
        logger.on_trial_result(
            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
        )

        logger.on_trial_complete(3, [], t)
        self._validate_json_result(config)

    def _validate_json_result(self, config):
        # Check result logs
        results = []
        result_file = os.path.join(self.test_dir, EXPR_RESULT_FILE)
        with open(result_file, "rt") as fp:
            for row in fp.readlines():
                results.append(json.loads(row))

        self.assertEqual(len(results), 3)
        self.assertSequenceEqual(
            [int(row["episode_reward_mean"]) for row in results], [4, 5, 6]
        )

        # Check json saved config file
        config_file = os.path.join(self.test_dir, EXPR_PARAM_FILE)
        with open(config_file, "rt") as fp:
            loaded_config = json.load(fp)

        self.assertEqual(loaded_config, config)

        # Check pickled config file
        config_file = os.path.join(self.test_dir, EXPR_PARAM_PICKLE_FILE)
        with open(config_file, "rb") as fp:
            loaded_config = cloudpickle.load(fp)

        self.assertEqual(loaded_config, config)

    def testTBX(self):
        config = {
            "a": 2,
            "b": [1, 2],
            "c": {"c": {"D": 123}},
            "int32": np.int32(1),
            "int64": np.int64(2),
            "bool8": np.bool_(True),
            "float32": np.float32(3),
            "float64": np.float64(4),
            "bad": np.float128(4),
        }
        t = Trial(evaluated_params=config, trial_id="tbx", logdir=self.test_dir)
        logger = TBXLoggerCallback()
        logger.on_trial_result(0, [], t, result(0, 4))
        logger.on_trial_result(1, [], t, result(1, 5))
        logger.on_trial_result(
            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
        )

        logger.on_trial_complete(3, [], t)

        self._validate_tbx_result(
            params=(b"float32", b"float64", b"int32", b"int64", b"bool8"),
            excluded_params=(b"bad",),
        )

    def _validate_tbx_result(self, params=None, excluded_params=None):
        try:
            from tensorflow.python.summary.summary_iterator import summary_iterator
        except ImportError:
            print("Skipping rest of test as tensorflow is not installed.")
            return

        events_file = list(glob.glob(f"{self.test_dir}/events*"))[0]
        results = []
        excluded_params = excluded_params or []
        for event in summary_iterator(events_file):
            for v in event.summary.value:
                if v.tag == "ray/tune/episode_reward_mean":
                    results.append(v.simple_value)
                elif v.tag == "_hparams_/experiment" and params:
                    for key in params:
                        self.assertIn(key, v.metadata.plugin_data.content)
                    for key in excluded_params:
                        self.assertNotIn(key, v.metadata.plugin_data.content)
                elif v.tag == "_hparams_/session_start_info" and params:
                    for key in params:
                        self.assertIn(key, v.metadata.plugin_data.content)
                    for key in excluded_params:
                        self.assertNotIn(key, v.metadata.plugin_data.content)

        self.assertEqual(len(results), 3)
        self.assertSequenceEqual([int(res) for res in results], [4, 5, 6])

    def testBadTBX(self):
        config = {"b": (1, 2, 3)}
        t = Trial(evaluated_params=config, trial_id="tbx", logdir=self.test_dir)
        logger = TBXLoggerCallback()
        logger.on_trial_result(0, [], t, result(0, 4))
        logger.on_trial_result(1, [], t, result(1, 5))
        logger.on_trial_result(
            2, [], t, result(2, 6, score=[1, 2, 3], hello={"world": 1})
        )
        with self.assertLogs("ray.tune.logger", level="INFO") as cm:
            logger.on_trial_complete(3, [], t)
        assert "INFO" in cm.output[0]


@pytest.mark.skipif(sys.version_info >= (3, 12), reason="Aim doesn't support py312")
class AimLoggerSuite(unittest.TestCase):
    """Test Aim integration."""

    def setUp(self):
        self.test_dir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.test_dir, ignore_errors=True)

    def initialize_logger(self, repo=None, experiment_name=None, metrics=None):
        try:
            from aim import Repo
        except ImportError:
            print("Skipping rest of test as aim is not installed.")
            return

        class Dummy:
            pass

        self.config = {
            "a": 2,
            "b": [1, 2],
            "c": {"d": {"e": 123}},
            "int32": np.int32(1),
            "int64": np.int64(2),
            "bool8": np.bool_(True),
            "float32": np.float32(3),
            "float64": np.float64(4),
            "bad": Dummy(),
        }
        trial_logdir = os.path.join(self.test_dir, "trial_logdir")
        trials = [
            Trial(
                evaluated_params=self.config,
                trial_id="aim_1",
                experiment_path=self.test_dir,
                logdir=trial_logdir,
                experiment_dir_name="aim_test",
                path="bucket/aim_test/trial_0_logdir",
            ),
            Trial(
                evaluated_params=self.config,
                trial_id="aim_2",
                experiment_path=self.test_dir,
                logdir=trial_logdir,
                experiment_dir_name="aim_test",
                path="bucket/aim_test/trial_1_logdir",
            ),
        ]

        # Test that aim repo is saved to the experiment directory
        # (one up from the trial directory) as the default.
        # In this example, this is `self.test_dir`.
        repo = repo or self.test_dir
        logger = AimLoggerCallback(
            repo=repo, experiment_name=experiment_name, metrics=metrics
        )

        for i, t in enumerate(trials):
            with self.assertLogs("ray.tune.logger", level="INFO") as cm:
                logger.log_trial_start(t)
                # Check that we log that the "bad" hparam gets thrown away
                assert "INFO" in cm.output[0]

            logger.on_trial_result(0, [], t, result(0, 3 * i + 1))
            logger.on_trial_result(1, [], t, result(1, 3 * i + 2))
            logger.on_trial_result(
                2, [], t, result(2, 3 * i + 3, score=[1, 2, 3], hello={"world": 1})
            )
            logger.on_trial_complete(3, [], t)

        aim_repo = Repo(repo)
        runs = list(aim_repo.iter_runs())
        assert len(runs) == 2
        runs.sort(key=lambda r: r["trial_id"])
        return runs

    def validateLogs(self, runs: list, metrics: list = None):
        expected_logged_hparams = set(flatten_dict(self.config)) - {"bad"}

        for i, run in enumerate(runs):
            assert set(run["hparams"]) == expected_logged_hparams
            assert run.get("trial_log_dir")
            assert run.get("trial_ip")

            results = None
            all_tune_metrics = set()
            for metric in run.metrics():
                if metric.name.startswith("ray/tune/"):
                    all_tune_metrics.add(metric.name.replace("ray/tune/", ""))
                if metric.name == "ray/tune/episode_reward_mean":
                    results = metric.values.values_list()

            assert results
            # Make sure that the set of reported metrics matches with the
            # set of metric names passed in
            # If None is passed in, then all Tune metrics get reported
            assert metrics is None or set(metrics) == all_tune_metrics

            results = [int(res) for res in results]
            if i == 0:
                self.assertSequenceEqual(results, [1, 2, 3])
            elif i == 1:
                self.assertSequenceEqual(results, [4, 5, 6])

    def testDefault(self):
        """Test AimLoggerCallback with default settings.
        - Req: a repo gets created at the experiment-level directory.
        - Req: the experiment param passed into each aim Run is the Tune experiment name
        """
        runs = self.initialize_logger()
        self.validateLogs(runs)
        for run in runs:
            assert run.repo.path == os.path.join(self.test_dir, ".aim")
            assert run.experiment == "aim_test"

    def testFilteredMetrics(self):
        """Test AimLoggerCallback, logging only a subset of metrics."""
        metrics_to_log = ("episode_reward_mean",)
        runs = self.initialize_logger(metrics=metrics_to_log)
        self.validateLogs(runs=runs, metrics=metrics_to_log)

    def testCustomConfigurations(self):
        """Test AimLoggerCallback, setting a custom repo and experiment name."""
        custom_repo = os.path.join(self.test_dir, "custom_repo")
        runs = self.initialize_logger(repo=custom_repo, experiment_name="custom")
        self.validateLogs(runs)
        for run in runs:
            assert run.repo.path == os.path.join(custom_repo, ".aim")
            assert run.experiment == "custom"


if __name__ == "__main__":
    import sys

    import pytest

    sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))