# coding: utf-8
import os
import pickle
import shutil
import tempfile
import unittest

import ray
from ray import tune
from ray.rllib import _register_all
from ray.tune import Trainable
from ray.tune.utils import validate_save_restore


class SerialTuneRelativeLocalDirTest(unittest.TestCase):
    local_mode = True
    prefix = "Serial"

    class MockTrainable(Trainable):
        _name = "MockTrainable"

        def setup(self, config):
            self.state = {"hi": 1}

        def step(self):
            return {"timesteps_this_iter": 1, "done": True}

        def save_checkpoint(self, checkpoint_dir):
            checkpoint_path = os.path.join(
                checkpoint_dir, "checkpoint-{}".format(self._iteration)
            )
            with open(checkpoint_path, "wb") as f:
                pickle.dump(self.state, f)
            return checkpoint_path

        def load_checkpoint(self, checkpoint_path):
            with open(checkpoint_path, "rb") as f:
                extra_data = pickle.load(f)
            self.state.update(extra_data)

    def setUp(self):
        self.absolute_local_dir = None
        ray.init(num_cpus=1, num_gpus=0, local_mode=self.local_mode)

    def tearDown(self):
        if self.absolute_local_dir is not None:
            shutil.rmtree(self.absolute_local_dir, ignore_errors=True)
            self.absolute_local_dir = None
        ray.shutdown()
        # Without this line, test_tune_server.testAddTrial would fail.
        _register_all()

    def _get_trial_dir(self, absoulte_exp_dir):
        print("looking for", self.MockTrainable._name)
        print("in", os.listdir(absoulte_exp_dir))
        trial_dirname = next(
            (
                child_dir
                for child_dir in os.listdir(absoulte_exp_dir)
                if (
                    os.path.isdir(os.path.join(absoulte_exp_dir, child_dir))
                    and child_dir.startswith(self.MockTrainable._name)
                )
            )
        )

        trial_absolute_dir = os.path.join(absoulte_exp_dir, trial_dirname)

        return trial_dirname, trial_absolute_dir

    def _train(self, exp_name, local_dir, absolute_local_dir):
        (trial,) = tune.run(
            self.MockTrainable,
            name=exp_name,
            stop={"training_iteration": 1},
            checkpoint_freq=1,
            local_dir=local_dir,
            config={"env": "CartPole-v0", "log_level": "DEBUG"},
        ).trials

        exp_dir = os.path.join(absolute_local_dir, exp_name)
        _, abs_trial_dir = self._get_trial_dir(exp_dir)

        self.assertIsNone(trial.error_file)
        self.assertEqual(trial.local_dir, exp_dir)
        self.assertEqual(trial.logdir, abs_trial_dir)

        self.assertTrue(os.path.isdir(absolute_local_dir), absolute_local_dir)
        self.assertTrue(os.path.isdir(exp_dir))
        self.assertTrue(os.path.isdir(abs_trial_dir))
        self.assertTrue(
            os.path.isfile(
                os.path.join(abs_trial_dir, "checkpoint_000001/checkpoint-1")
            )
        )

    def _restore(self, exp_name, local_dir, absolute_local_dir):
        trial_name, abs_trial_dir = self._get_trial_dir(
            os.path.join(absolute_local_dir, exp_name)
        )

        checkpoint_path = os.path.join(
            local_dir, exp_name, trial_name, "checkpoint_000001/checkpoint-1"
        )  # Relative checkpoint path

        # The file tune would find. The absolute checkpoint path.
        tune_find_file = os.path.abspath(os.path.expanduser(checkpoint_path))
        self.assertTrue(
            os.path.isfile(tune_find_file), "{} is not exist!".format(tune_find_file)
        )

        (trial,) = tune.run(
            self.MockTrainable,
            name=exp_name,
            stop={"training_iteration": 2},  # train one more iteration.
            restore=checkpoint_path,  # Restore the checkpoint
            config={"env": "CartPole-v0", "log_level": "DEBUG"},
        ).trials
        self.assertIsNone(trial.error_file)

    def testDottedRelativePath(self):
        local_dir = "./test_dotted_relative_local_dir"
        exp_name = self.prefix + "DottedRelativeLocalDir"
        absolute_local_dir = os.path.abspath(local_dir)
        self.absolute_local_dir = absolute_local_dir
        self.assertFalse(os.path.exists(absolute_local_dir))
        self._train(exp_name, local_dir, absolute_local_dir)
        self._restore(exp_name, local_dir, absolute_local_dir)

    def testRelativePath(self):
        local_dir = "test_relative_local_dir"
        exp_name = self.prefix + "RelativePath"
        absolute_local_dir = os.path.abspath(local_dir)
        self.absolute_local_dir = absolute_local_dir
        self.assertFalse(os.path.exists(absolute_local_dir))
        self._train(exp_name, local_dir, absolute_local_dir)
        self._restore(exp_name, local_dir, absolute_local_dir)

    def testTildeAbsolutePath(self):
        local_dir = "~/test_tilde_absolute_local_dir"
        exp_name = self.prefix + "TildeAbsolutePath"
        absolute_local_dir = os.path.abspath(os.path.expanduser(local_dir))
        self.absolute_local_dir = absolute_local_dir
        self.assertFalse(os.path.exists(absolute_local_dir))
        self._train(exp_name, local_dir, absolute_local_dir)
        self._restore(exp_name, local_dir, absolute_local_dir)

    def testTempfile(self):
        local_dir = tempfile.mkdtemp()
        exp_name = self.prefix + "Tempfile"
        self.absolute_local_dir = local_dir
        self._train(exp_name, local_dir, local_dir)
        self._restore(exp_name, local_dir, local_dir)

    def testCheckpointWithNoop(self):
        """Tests that passing the checkpoint_dir right back works."""

        class MockTrainable(Trainable):
            def setup(self, config):
                pass

            def step(self):
                return {"score": 1}

            def save_checkpoint(self, checkpoint_dir):
                with open(os.path.join(checkpoint_dir, "test.txt"), "wb") as f:
                    pickle.dump("test", f)
                return checkpoint_dir

            def load_checkpoint(self, checkpoint_dir):
                with open(os.path.join(checkpoint_dir, "test.txt"), "rb") as f:
                    x = pickle.load(f)

                assert x == "test"
                return checkpoint_dir

        validate_save_restore(MockTrainable)
        validate_save_restore(MockTrainable, use_object_store=True)


if __name__ == "__main__":
    import pytest
    import sys

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

[((43, 8, 43, 68), 'ray.init', 'ray.init', (), '', False, 'import ray\n'), ((49, 8, 49, 22), 'ray.shutdown', 'ray.shutdown', ({}, {}), '()', False, 'import ray\n'), ((51, 8, 51, 23), 'ray.rllib._register_all', '_register_all', ({}, {}), '()', False, 'from ray.rllib import _register_all\n'), ((67, 29, 67, 74), 'os.path.join', 'os.path.join', ({(67, 42, 67, 58): 'absoulte_exp_dir', (67, 60, 67, 73): 'trial_dirname'}, {}), '(absoulte_exp_dir, trial_dirname)', False, 'import os\n'), ((81, 18, 81, 60), 'os.path.join', 'os.path.join', ({(81, 31, 81, 49): 'absolute_local_dir', (81, 51, 81, 59): 'exp_name'}, {}), '(absolute_local_dir, exp_name)', False, 'import os\n'), ((102, 26, 104, 9), 'os.path.join', 'os.path.join', ({(103, 12, 103, 21): 'local_dir', (103, 23, 103, 31): 'exp_name', (103, 33, 103, 43): 'trial_name', (103, 45, 103, 77): '"""checkpoint_000001/checkpoint-1"""'}, {}), "(local_dir, exp_name, trial_name, 'checkpoint_000001/checkpoint-1')", False, 'import os\n'), ((124, 29, 124, 55), 'os.path.abspath', 'os.path.abspath', ({(124, 45, 124, 54): 'local_dir'}, {}), '(local_dir)', False, 'import os\n'), ((133, 29, 133, 55), 'os.path.abspath', 'os.path.abspath', ({(133, 45, 133, 54): 'local_dir'}, {}), '(local_dir)', False, 'import os\n'), ((149, 20, 149, 38), 'tempfile.mkdtemp', 'tempfile.mkdtemp', ({}, {}), '()', False, 'import tempfile\n'), ((177, 8, 177, 44), 'ray.tune.utils.validate_save_restore', 'validate_save_restore', ({(177, 30, 177, 43): 'MockTrainable'}, {}), '(MockTrainable)', False, 'from ray.tune.utils import validate_save_restore\n'), ((178, 8, 178, 67), 'ray.tune.utils.validate_save_restore', 'validate_save_restore', (), '', False, 'from ray.tune.utils import validate_save_restore\n'), ((185, 13, 185, 42), 'pytest.main', 'pytest.main', ({(185, 25, 185, 41): "['-v', __file__]"}, {}), "(['-v', __file__])", False, 'import pytest\n'), ((47, 12, 47, 70), 'shutil.rmtree', 'shutil.rmtree', (), '', False, 'import shutil\n'), ((55, 20, 55, 48), 'os.listdir', 'os.listdir', ({(55, 31, 55, 47): 'absoulte_exp_dir'}, {}), '(absoulte_exp_dir)', False, 'import os\n'), ((72, 19, 79, 9), 'ray.tune.run', 'tune.run', (), '', False, 'from ray import tune\n'), ((88, 24, 88, 57), 'os.path.isdir', 'os.path.isdir', ({(88, 38, 88, 56): 'absolute_local_dir'}, {}), '(absolute_local_dir)', False, 'import os\n'), ((89, 24, 89, 46), 'os.path.isdir', 'os.path.isdir', ({(89, 38, 89, 45): 'exp_dir'}, {}), '(exp_dir)', False, 'import os\n'), ((90, 24, 90, 52), 'os.path.isdir', 'os.path.isdir', ({(90, 38, 90, 51): 'abs_trial_dir'}, {}), '(abs_trial_dir)', False, 'import os\n'), ((99, 12, 99, 54), 'os.path.join', 'os.path.join', ({(99, 25, 99, 43): 'absolute_local_dir', (99, 45, 99, 53): 'exp_name'}, {}), '(absolute_local_dir, exp_name)', False, 'import os\n'), ((107, 41, 107, 76), 'os.path.expanduser', 'os.path.expanduser', ({(107, 60, 107, 75): 'checkpoint_path'}, {}), '(checkpoint_path)', False, 'import os\n'), ((109, 12, 109, 42), 'os.path.isfile', 'os.path.isfile', ({(109, 27, 109, 41): 'tune_find_file'}, {}), '(tune_find_file)', False, 'import os\n'), ((112, 19, 118, 9), 'ray.tune.run', 'tune.run', (), '', False, 'from ray import tune\n'), ((126, 25, 126, 59), 'os.path.exists', 'os.path.exists', ({(126, 40, 126, 58): 'absolute_local_dir'}, {}), '(absolute_local_dir)', False, 'import os\n'), ((135, 25, 135, 59), 'os.path.exists', 'os.path.exists', ({(135, 40, 135, 58): 'absolute_local_dir'}, {}), '(absolute_local_dir)', False, 'import os\n'), ((142, 45, 142, 74), 'os.path.expanduser', 'os.path.expanduser', ({(142, 64, 142, 73): 'local_dir'}, {}), '(local_dir)', False, 'import os\n'), ((144, 25, 144, 59), 'os.path.exists', 'os.path.exists', ({(144, 40, 144, 58): 'absolute_local_dir'}, {}), '(absolute_local_dir)', False, 'import os\n'), ((33, 16, 33, 42), 'pickle.dump', 'pickle.dump', ({(33, 28, 33, 38): 'self.state', (33, 40, 33, 41): 'f'}, {}), '(self.state, f)', False, 'import pickle\n'), ((38, 29, 38, 43), 'pickle.load', 'pickle.load', ({(38, 41, 38, 42): 'f'}, {}), '(f)', False, 'import pickle\n'), ((93, 16, 93, 77), 'os.path.join', 'os.path.join', ({(93, 29, 93, 42): 'abs_trial_dir', (93, 44, 93, 76): '"""checkpoint_000001/checkpoint-1"""'}, {}), "(abs_trial_dir, 'checkpoint_000001/checkpoint-1')", False, 'import os\n'), ((59, 33, 59, 61), 'os.listdir', 'os.listdir', ({(59, 44, 59, 60): 'absoulte_exp_dir'}, {}), '(absoulte_exp_dir)', False, 'import os\n'), ((167, 20, 167, 42), 'pickle.dump', 'pickle.dump', ({(167, 32, 167, 38): '"""test"""', (167, 40, 167, 41): 'f'}, {}), "('test', f)", False, 'import pickle\n'), ((172, 24, 172, 38), 'pickle.load', 'pickle.load', ({(172, 36, 172, 37): 'f'}, {}), '(f)', False, 'import pickle\n'), ((166, 26, 166, 66), 'os.path.join', 'os.path.join', ({(166, 39, 166, 53): 'checkpoint_dir', (166, 55, 166, 65): '"""test.txt"""'}, {}), "(checkpoint_dir, 'test.txt')", False, 'import os\n'), ((171, 26, 171, 66), 'os.path.join', 'os.path.join', ({(171, 39, 171, 53): 'checkpoint_dir', (171, 55, 171, 65): '"""test.txt"""'}, {}), "(checkpoint_dir, 'test.txt')", False, 'import os\n'), ((61, 34, 61, 75), 'os.path.join', 'os.path.join', ({(61, 47, 61, 63): 'absoulte_exp_dir', (61, 65, 61, 74): 'child_dir'}, {}), '(absoulte_exp_dir, child_dir)', False, 'import os\n')]