# Copyright (c) 2019-2020, RTE (https://www.rte-france.com)
# See AUTHORS.txt
# This Source Code Form is subject to the terms of the Mozilla Public License, version 2.0.
# If a copy of the Mozilla Public License, version 2.0 was not distributed with this file,
# you can obtain one at http://mozilla.org/MPL/2.0/.
# SPDX-License-Identifier: MPL-2.0
# This file is part of Grid2Op, Grid2Op a testbed platform to model sequential decision making in power systems.

import grid2op
import numpy as np
#!/usr/bin/env python3

import grid2op
import unittest
import numpy as np
import warnings
from grid2op.Parameters import Parameters
import pdb


class Issue140Tester(unittest.TestCase):
    def test_issue_140(self):

        # TODO change that
        envs = grid2op.list_available_local_env()
        env_name = "l2rpn_neurips_2020_track1_small"
        if env_name not in envs:
            # the environment is not downloaded, I skip this test
            self.skipTest("{} is not downloaded".format(env_name))

        param = Parameters()
        # test was originally designed with these values
        param.init_from_dict({
            "NO_OVERFLOW_DISCONNECTION":  False,
            "NB_TIMESTEP_OVERFLOW_ALLOWED": 3,
            "NB_TIMESTEP_COOLDOWN_SUB": 3,
            "NB_TIMESTEP_COOLDOWN_LINE": 3,
            "HARD_OVERFLOW_THRESHOLD": 200.,
            "NB_TIMESTEP_RECONNECTION": 12,
            "IGNORE_MIN_UP_DOWN_TIME": True,
            "ALLOW_DISPATCH_GEN_SWITCH_OFF": True,
            "ENV_DC": False,
            "FORECAST_DC": False,
            "MAX_SUB_CHANGED": 1,
            "MAX_LINE_STATUS_CHANGED": 1
        })
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            env = grid2op.make(env_name, param=param)
        ts_per_chronics = 2016

        seed = 725
        np.random.seed(seed)
        env.seed(seed)

        env.set_id(np.random.randint(100000))
        _ = env.reset()

        max_day = (env.chronics_handler.max_timestep() - 2016) // 288
        start_timestep = np.random.randint(max_day) * 288 - 1  # start at 00:00
        if start_timestep > 0:
            env.fast_forward_chronics(start_timestep)

        obs = env.get_obs()

        done = False
        steps = 0
        do_nothing = env.action_space({})

        for i in range(119):
            obs, reward, done, info = env.step(do_nothing)
            assert not done

        act0 = env.action_space({"set_bus": {"lines_ex_id": [(41, 2), (42, 1), (57, 1)],
                                             "lines_or_id": [(44, 2)],
                                             "generators_id": [(16, 2)]}})
        act1 = env.action_space({"set_bus": {"lines_ex_id": [(17, 2), (18, 1), (19, 1), (20, 2), (21, 2),
                                                             ],
                                             "lines_or_id": [(22, 2), (23, 2), (27, 1), (28, 1), (48, 2),
                                                             (49, 2), (54, 1)],
                                             "generators_id": [(5, 2), (6, 1), (7, 1), (8, 1)],
                                             "loads_id": [(17, 2)]}})
        act2 = env.action_space({"set_bus": {"lines_ex_id": [(36, 2), (37, 2), (38, 2), (39, 2)],
                                             "lines_or_id": [(40, 1), (41, 1)],
                                             "generators_id": [(14, 1)],
                                             "loads_id": [(27, 1)]}})
        obs, reward, done, info = env.step(act0)
        obs, reward, done, info = env.step(act1)
        assert not done

        simulate_obs0, simulate_reward0, simulate_done0, simulate_info0 = obs.simulate(do_nothing)
        simulate_obs1, simulate_reward1, simulate_done1, simulate_info1 = obs.simulate(act2)
        obs, reward, done, info = env.step(act2)

        assert simulate_done0 is False
        assert simulate_done1 is True
        assert done is True

        assert info["is_illegal"] == False
        assert simulate_info1["is_illegal"] == False


if __name__ == "__main__":
    unittest.main()
