# 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 unittest
import warnings
import numpy as np

import grid2op
from grid2op.dtypes import dt_float
from grid2op.Reward import RedispReward
from grid2op.Runner import Runner


class Issue187Tester(unittest.TestCase):
    """
    this test ensure that every "test" environment can be converted to gym

    this test suit goes beyond the simple error raised in the github issue.
    """
    def setUp(self) -> None:
        self.tol = 1e-6

    def test_issue_187(self):
        """test the range of the reward class"""
        for env_name in grid2op.list_available_test_env():
            if env_name == "blank":
                continue
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                with grid2op.make(env_name, test=True, reward_class=RedispReward) as env:
                    obs = env.reset()
                    obs, reward, done, info = env.step(env.action_space())
                    assert reward <= env.reward_range[1], f"error for reward_max for {env_name}"
                    assert reward >= env.reward_range[0], f"error for reward_min for {env_name}"

    def test_custom_reward(self):
        """test i can generate the reward and use it in the envs"""
        reward_cls = RedispReward.generate_class_custom_params(alpha_redisph=2,
                                                               min_load_ratio=0.15,
                                                               worst_losses_ratio=0.05,
                                                               min_reward=-10.,
                                                               reward_illegal_ambiguous=0.,
                                                               least_losses_ratio=0.015)

        for env_name in grid2op.list_available_test_env():
            if env_name == "blank":
                continue
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                with grid2op.make(env_name, test=True, reward_class=reward_cls) as env:
                    obs = env.reset()
                    obs, reward, done, info = env.step(env.action_space())
                    # test that reward is in the correct range
                    assert reward <= env.reward_range[1], f"error reward > reward_max for {env_name}"
                    assert reward >= env.reward_range[0], f"error reward < reward_min for {env_name}"

                    # test the parameters are effectively changed

                    # what should be computed
                    _alpha_redisph = dt_float(2)
                    _min_load_ratio = dt_float(0.15)
                    _worst_losses_ratio = dt_float(0.05)
                    _min_reward = dt_float(-10.)
                    _reward_illegal_ambiguous = dt_float(0.)
                    _least_losses_ratio = dt_float(0.015)

                    worst_marginal_cost = np.max(env.gen_cost_per_MW)
                    worst_load = dt_float(np.sum(env.gen_pmax))
                    # it's not the worst, but definitely an upper bound
                    worst_losses = dt_float(_worst_losses_ratio) * worst_load
                    worst_redisp = _alpha_redisph * np.sum(env.gen_pmax)  # not realistic, but an upper bound
                    max_regret = (worst_losses + worst_redisp) * worst_marginal_cost
                    reward_min = dt_float(_min_reward)

                    least_loads = dt_float(worst_load * _min_load_ratio)  # half the capacity of the grid
                    least_losses = dt_float(_least_losses_ratio * least_loads)  # 1.5% of losses
                    least_redisp = dt_float(0.0)  # lower_bound is 0
                    base_marginal_cost = np.min(env.gen_cost_per_MW[env.gen_cost_per_MW > 0.])
                    min_regret = (least_losses + least_redisp) * base_marginal_cost
                    reward_max = dt_float((max_regret - min_regret) / least_loads)
                    assert abs(env.reward_range[1] - reward_max) <= self.tol, \
                            f"wrong reward max computed for {env_name}"
                    assert abs(env.reward_range[0] - reward_min) <= self.tol, \
                            f"wrong reward min computed for {env_name}"

    def test_custom_reward_runner(self):
        """test i can generate the reward and use it in the envs"""
        reward_cls = RedispReward.generate_class_custom_params(alpha_redisph=2,
                                                               min_load_ratio=0.15,
                                                               worst_losses_ratio=0.05,
                                                               min_reward=-10.,
                                                               reward_illegal_ambiguous=0.,
                                                               least_losses_ratio=0.015)
        env_name = "l2rpn_case14_sandbox"
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            with grid2op.make(env_name, test=True, reward_class=reward_cls) as env:
                obs = env.reset()
                runner = Runner(**env.get_params_for_runner())
                res = runner.run(nb_episode=2, nb_process=2)


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