# 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.

# do some generic tests that can be implemented directly to test if a backend implementation can work out of the box
# with grid2op.
# see an example of test_Pandapower for how to use this suit.
import unittest
import numpy as np
import warnings

import grid2op
from grid2op.Backend.EducPandaPowerBackend import EducPandaPowerBackend
from grid2op.Exceptions import EnvError

import pdb


class TestAuxFunctions(unittest.TestCase):
    def setUp(self) -> None:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            self.envref = grid2op.make("rte_case14_realistic",
                                       test=True,
                                       _add_to_name="test_gridobjects_testauxfunctions")
        seed = 0
        self.nb_test = 10
        self.max_iter = 30

        self.envref.seed(seed)
        self.seeds = [i for i in range(self.nb_test)]  # used for seeding environment and agent

    def tearDown(self) -> None:
        self.envref.close()

    def test_auxilliary_func(self):
        """
        test the methods _compute_sub_pos works
        """
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            MyClass = EducPandaPowerBackend.init_grid(self.envref.backend)
            backend = MyClass()

        backend.n_sub = self.envref.backend.n_sub
        backend.n_load = self.envref.backend.n_load
        backend.n_gen = self.envref.backend.n_gen
        backend.n_line = self.envref.backend.n_line
        backend.gen_to_subid = self.envref.backend.gen_to_subid
        backend.load_to_subid = self.envref.backend.load_to_subid
        backend.line_or_to_subid = self.envref.backend.line_or_to_subid
        backend.line_ex_to_subid = self.envref.backend.line_ex_to_subid

        # now "hack" the class to check that the correct element are correctly implemented
        bk_cls = type(backend)

        # delete the attributes we want to test
        bk_cls.sub_info = None
        bk_cls.load_to_sub_pos = None
        bk_cls.gen_to_sub_pos = None
        bk_cls.line_or_to_sub_pos = None
        bk_cls.line_ex_to_sub_pos = None
        bk_cls.line_ex_to_sub_pos = None
        bk_cls.load_pos_topo_vect = None
        bk_cls.gen_pos_topo_vect = None
        bk_cls.line_or_pos_topo_vect = None
        bk_cls.line_ex_pos_topo_vect = None

        # test that the grid is not correct now
        with self.assertRaises(EnvError):
            bk_cls.assert_grid_correct_cls()

        # fill the _compute_sub_elements
        bk_cls._compute_sub_elements()
        assert np.sum(bk_cls.sub_info) == 56
        assert np.all(bk_cls.sub_info == [3, 6, 4, 6, 5, 6, 3, 2, 5, 3, 3, 3, 4, 3])

        # fill the *sub_pos
        bk_cls._compute_sub_pos()
        assert np.all(bk_cls.load_to_sub_pos == 0)
        assert np.all(bk_cls.gen_to_sub_pos == [1, 1, 1, 0, 0])
        assert np.all(bk_cls.line_or_to_sub_pos == [1, 2, 2, 3, 4, 2, 1, 2, 3, 4, 1, 2, 1, 1, 1, 2, 3, 1, 0, 3])
        assert np.all(bk_cls.line_ex_to_sub_pos == [5, 2, 3, 4, 3, 5, 4, 1, 2, 2, 2, 1, 2, 3, 2, 1, 4, 5, 1, 2])

        # fill the *pos_topo_vect
        backend._compute_pos_big_topo()  # i test the object class here
        assert np.all(bk_cls.load_pos_topo_vect == [3,  9, 13, 19, 24, 35, 40, 43, 46, 49, 53])
        assert np.all(bk_cls.gen_pos_topo_vect == [4, 10, 25, 33,  0])
        assert np.all(bk_cls.line_or_pos_topo_vect == [1,  2,  5,  6,  7, 11, 14, 26, 27, 28, 36, 37, 41, 47, 50, 15,
                                                       16, 20, 30, 38])
        assert np.all(bk_cls.line_ex_pos_topo_vect == [8, 21, 12, 17, 22, 18, 23, 44, 48, 51, 42, 54, 45, 52, 55, 31,
                                                       39, 29, 34, 32])
        # this should pass
        bk_cls.assert_grid_correct_cls()


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