# 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 tempfile
import warnings
import pdb

import grid2op
from grid2op.Agent import OneChangeThenNothing, RandomAgent
from grid2op.tests.helper_path_test import *
from grid2op.Chronics import Multifolder
from grid2op.Reward import L2RPNReward
from grid2op.Backend import PandaPowerBackend
from grid2op.Runner import Runner
from grid2op.Episode import EpisodeData
from grid2op.dtypes import dt_float
from grid2op.Agent import BaseAgent
from grid2op.Action import TopologyAction
from grid2op.Parameters import Parameters
from grid2op.MakeEnv import make
from grid2op.Opponent.BaseActionBudget import BaseActionBudget
from grid2op.Opponent import RandomLineOpponent


DEBUG = True
PATH_ADN_CHRONICS_FOLDER = os.path.abspath(os.path.join(PATH_CHRONICS, "test_multi_chronics"))

import warnings
warnings.simplefilter("error")


class TestEpisodeData(unittest.TestCase):
    def setUp(self):
        """
        The case file is a representation of the case14 as found in the ieee14 powergrid.
        :return:
        """
        self.tolvect = dt_float(1e-2)
        self.tol_one = dt_float(1e-5)
        self.max_iter = 10
        # self.real_reward = dt_float(199.99800)
        self.real_reward = dt_float(179.99818)

        self.init_grid_path = os.path.join(
            PATH_DATA_TEST_PP, "test_case14.json")
        self.path_chron = PATH_ADN_CHRONICS_FOLDER
        self.parameters_path = None
        self.names_chronics_to_backend = {"loads": {"2_C-10.61": 'load_1_0', "3_C151.15": 'load_2_1',
                                                    "14_C63.6": 'load_13_2', "4_C-9.47": 'load_3_3',
                                                    "5_C201.84": 'load_4_4',
                                                    "6_C-6.27": 'load_5_5', "9_C130.49": 'load_8_6',
                                                    "10_C228.66": 'load_9_7',
                                                    "11_C-138.89": 'load_10_8', "12_C-27.88": 'load_11_9',
                                                    "13_C-13.33": 'load_12_10'},
                                          "lines": {'1_2_1': '0_1_0', '1_5_2': '0_4_1', '9_10_16': '8_9_2',
                                                    '9_14_17': '8_13_3',
                                                    '10_11_18': '9_10_4', '12_13_19': '11_12_5', '13_14_20': '12_13_6',
                                                    '2_3_3': '1_2_7', '2_4_4': '1_3_8', '2_5_5': '1_4_9',
                                                    '3_4_6': '2_3_10',
                                                    '4_5_7': '3_4_11', '6_11_11': '5_10_12', '6_12_12': '5_11_13',
                                                    '6_13_13': '5_12_14', '4_7_8': '3_6_15', '4_9_9': '3_8_16',
                                                    '5_6_10': '4_5_17',
                                                    '7_8_14': '6_7_18', '7_9_15': '6_8_19'},
                                          "prods": {"1_G137.1": 'gen_0_4', "3_G36.31": "gen_2_1", "6_G63.29": "gen_5_2",
                                                    "2_G-56.47": "gen_1_0", "8_G40.43": "gen_7_3"},
                                          }
        self.gridStateclass = Multifolder
        self.backendClass = PandaPowerBackend
        self.runner = Runner(init_grid_path=self.init_grid_path,
                             path_chron=self.path_chron,
                             parameters_path=self.parameters_path,
                             names_chronics_to_backend=self.names_chronics_to_backend,
                             gridStateclass=self.gridStateclass,
                             backendClass=self.backendClass,
                             rewardClass=L2RPNReward,
                             other_rewards={"test": L2RPNReward},
                             max_iter=self.max_iter,
                             name_env="test_episodedata_env")

    def test_load_ambiguous(self):
        f = tempfile.mkdtemp()

        class TestSuitAgent(BaseAgent):
            def __init__(self, *args, **kwargs):
                BaseAgent.__init__(self, *args, **kwargs)

            def act(self, observation, reward, done=False):
                # do a ambiguous action
                return self.action_space({"set_line_status": [(0, 1)],
                                          "change_line_status": [0]}
                                         )

        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            with grid2op.make("rte_case14_test", test=True) as env:
                my_agent = TestSuitAgent(env.action_space)
                runner = Runner(**env.get_params_for_runner(),
                                agentClass=None,
                                agentInstance=my_agent)

                # test that the right seeds are assigned to the agent
                res = runner.run(nb_episode=1,
                                 max_iter=self.max_iter,
                                 path_save=f)
            episode_data = EpisodeData.from_disk(agent_path=f, name=res[0][1])
        assert int(episode_data.meta["chronics_max_timestep"]) == self.max_iter
        assert len(episode_data.actions) == self.max_iter
        assert len(episode_data.observations) == self.max_iter + 1
        assert len(episode_data.env_actions) == self.max_iter
        assert len(episode_data.attacks) == self.max_iter

    def test_one_episode_with_saving(self):
        f = tempfile.mkdtemp()
        episode_name, cum_reward, timestep, episode_data_cached = self.runner.run_one_episode(path_save=f)
        episode_data = EpisodeData.from_disk(agent_path=f, name=episode_name)
        assert int(episode_data.meta["chronics_max_timestep"]) == self.max_iter
        assert len(episode_data.other_rewards) == self.max_iter
        for other, real in zip(episode_data.other_rewards, episode_data.rewards):
            assert dt_float(np.abs(other["test"] - real)) <= self.tol_one
        assert np.abs(dt_float(episode_data.meta["cumulative_reward"]) - self.real_reward) <= self.tol_one

    def test_collection_wrapper_after_run(self):
        OneChange = OneChangeThenNothing.gen_next({"set_bus": {"lines_or_id": [(1,-1)]}})
        runner = Runner(init_grid_path=self.init_grid_path,
                        path_chron=self.path_chron,
                        parameters_path=self.parameters_path,
                        names_chronics_to_backend=self.names_chronics_to_backend,
                        gridStateclass=self.gridStateclass,
                        backendClass=self.backendClass,
                        rewardClass=L2RPNReward,
                        other_rewards={"test": L2RPNReward},
                        max_iter=self.max_iter,
                        name_env="test_episodedata_env",
                        agentClass=OneChange)
        _, cum_reward, timestep, episode_data = runner.run_one_episode(max_iter=self.max_iter, detailed_output=True)
        # Check that the type of first action is set bus
        assert episode_data.actions[0].get_types()[2]

    def test_len(self):
        """test i can use the function "len" of the episode data"""
        f = tempfile.mkdtemp()
        episode_name, cum_reward, timestep, episode_data_cached = self.runner.run_one_episode(path_save=f)
        episode_data = EpisodeData.from_disk(agent_path=f, name=episode_name)
        len(episode_data)

    def test_3_episode_with_saving(self):
        f = tempfile.mkdtemp()
        res = self.runner._run_sequential(nb_episode=3, path_save=f)
        for i, episode_name, cum_reward, timestep, total_ts in res:
            episode_data = EpisodeData.from_disk(agent_path=f, name=episode_name)
            assert int(episode_data.meta["chronics_max_timestep"]) == self.max_iter
            assert np.abs(
                dt_float(episode_data.meta["cumulative_reward"]) - self.real_reward) <= self.tol_one

    def test_3_episode_3process_with_saving(self):
        f = tempfile.mkdtemp()
        nb_episode = 2
        res = self.runner._run_parrallel(nb_episode=nb_episode, nb_process=2, path_save=f)
        assert len(res) == nb_episode
        for i, episode_name, cum_reward, timestep, total_ts in res:
            episode_data = EpisodeData.from_disk(agent_path=f, name=episode_name)
            assert int(episode_data.meta["chronics_max_timestep"]) == self.max_iter
            assert np.abs(
                dt_float(episode_data.meta["cumulative_reward"]) - self.real_reward) <= self.tol_one

    def test_with_opponent(self):
        init_budget = 1000
        opponent_attack_duration = 15
        opponent_attack_cooldown = 30
        opponent_budget_per_ts = 0.
        opponent_action_class = TopologyAction

        LINES_ATTACKED = ["1_3_3", "1_4_4", "3_6_15", "9_10_12", "11_12_13", "12_13_14"]

        p = Parameters()
        p.NO_OVERFLOW_DISCONNECTION = True
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            env = make("rte_case14_realistic",
                       test=True,
                       param=p,
                       opponent_init_budget=init_budget,
                       opponent_budget_per_ts=opponent_budget_per_ts,
                       opponent_attack_cooldown=opponent_attack_cooldown,
                       opponent_attack_duration=opponent_attack_duration,
                       opponent_action_class=opponent_action_class,
                       opponent_budget_class=BaseActionBudget,
                       opponent_class=RandomLineOpponent,
                       kwargs_opponent={
                           "lines_attacked": LINES_ATTACKED
                       })
        env.seed(0)
        runner = Runner(**env.get_params_for_runner())

        f = tempfile.mkdtemp()
        res = runner.run(nb_episode=1,
                         env_seeds=[4], agent_seeds=[0],
                         max_iter=opponent_attack_cooldown - 1,
                         path_save=f)

        episode_data = EpisodeData.from_disk(agent_path=f, name=res[0][1])
        lines_impacted, subs_impacted = episode_data.attacks[0].get_topological_impact()
        assert lines_impacted[3]


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