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

from grid2op.tests.helper_path_test import *

import grid2op
from grid2op.Parameters import Parameters
from grid2op.dtypes import dt_float
from grid2op.Action import PlayableAction, CompleteAction

import warnings

# TODO check when there is also redispatching


class TestCurtailmentEnv(HelperTests):
    """test the env part of the storage functionality"""
    def setUp(self) -> None:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore")
            param = Parameters()
            param.NO_OVERFLOW_DISCONNECTION = True
            self.env1 = grid2op.make("educ_case14_storage", test=True, action_class=PlayableAction, param=param)
            self.env2 = self.env1.copy()

    def tearDown(self) -> None:
        self.env1.close()
        self.env2.close()

    def test_reset(self):
        """test curtailment is ok when env is reset"""
        obs = self.env1.reset()
        assert np.all(obs.curtailment >= 0.)
        assert np.all(obs.curtailment <= 1.)
        assert np.all(obs.curtailment[~obs.gen_renewable] == 0.)

    def test_action_ok(self):
        """test a storage action is supported (basic test)"""
        act = self.env1.action_space({"curtail": [(2, 0.5)]})
        str_ = act.__str__()
        real_str = 'This action will:\n' \
                   '\t - NOT change anything to the injections\n' \
                   '\t - NOT perform any redispatching action\n' \
                   '\t - NOT modify any storage capacity\n' \
                   '\t - Perform the following curtailment:\n' \
                   '\t \t - Limit unit "gen_5_2" to 50.0% of its Pmax (setpoint: 0.500)\n' \
                   '\t - NOT force any line status\n' \
                   '\t - NOT switch any line status\n\t - NOT switch anything in the topology\n' \
                   '\t - NOT force any particular bus configuration'
        assert str_ == real_str

    def test_curtailment_no_effect(self):
        act = self.env1.action_space({"curtail": [(2, 5)]})
        obs1, *_ = self.env1.step(act)
        obs2, *_ = self.env2.step(self.env2.action_space())
        assert obs1 == obs2

    def test_curtailment_effect(self):
        """
        test different kind of equations when curtailment has been applied
        in this example the 23% of curtailment is just the right amount so that, during the
        """
        gen_id = 2
        ratio_max = 0.23
        act = self.env1.action_space({"curtail": [(gen_id, ratio_max)]})
        self.env1.step(act)
        self.env2.step(self.env2.action_space())
        tested = False
        for step in range(100):
            obs1_3, reward1_3, done1_3, info1_3 = self.env1.step(self.env1.action_space())
            obs2_3, reward2_3, done2_3, info2_3 = self.env2.step(self.env2.action_space())
            if done1_3 or done2_3:
                raise RuntimeError("Environment did game over, this is not normal")
            if info1_3["exception"]:
                raise RuntimeError("Exception for info1_3")
            if info2_3["exception"]:
                raise RuntimeError("Exception for info2_3")
            assert abs(np.sum(obs1_3.actual_dispatch) + (obs1_3.gen_p[gen_id] - obs2_3.gen_p[gen_id])) <= 1e-4, \
                f"for step {step}"
            this_prod_th = obs1_3.gen_p_before_curtail[self.env1.gen_renewable]
            th_prod = obs2_3.gen_p[self.env1.gen_renewable]
            assert np.max(np.abs(this_prod_th - th_prod)) <= 1e-4, \
                f"for step {step}"
            assert obs1_3.gen_p[gen_id] <= ratio_max * obs1_3.gen_pmax[gen_id] + 1e-5, \
                f"for step {step}"
            amount_curtailed = (np.sum(this_prod_th - obs1_3.gen_p[self.env1.gen_renewable]))
            assert abs(np.sum(obs1_3.actual_dispatch) - amount_curtailed) <= 1e-4, \
                f"for step {step}"
            if amount_curtailed == 0.:
                # test that redispatch goes back to all O when no curtailment is done
                assert np.max(np.abs(obs1_3.actual_dispatch)) <= 1e-4,  f"for step {step}"
                tested = True
            assert np.all(obs1_3.curtailment[~obs1_3.gen_renewable] == 0.)

        if not tested:
            raise RuntimeError("the case where no curtailment have been used could not be tested. Please consider "
                               "making another test")

    def test_change_curtailment(self):
        """test that i can change the curtailment, cancel it, redo it etc."""
        gen_id = 2
        ratios_max = [0.25, 0.24, 0.23, 0.18, 0.29, 0.19, 0.05, 0.2, 0.3, 1.0]  # don't change it
        n_steps = 100  # don't change it
        acts = []
        for ratio in ratios_max:
            acts.append(self.env1.action_space({"curtail": [(gen_id, ratio)]}))
        self.env1.step(self.env2.action_space())
        self.env2.step(self.env2.action_space())
        tested = False
        for step in range(n_steps):
            li_id = step // 10
            ratio_max = ratios_max[li_id]
            obs1_3, reward1_3, done1_3, info1_3 = self.env1.step(acts[li_id])
            obs2_3, reward2_3, done2_3, info2_3 = self.env2.step(self.env2.action_space())
            if done1_3 or done2_3:
                raise RuntimeError("Environment did game over, this is not normal")
            if info1_3["exception"]:
                raise RuntimeError("Exception for info1_3")
            if info2_3["exception"]:
                raise RuntimeError("Exception for info2_3")

            assert abs(np.sum(obs1_3.actual_dispatch) + (obs1_3.gen_p[gen_id] - obs2_3.gen_p[gen_id])) <= 1e-4, \
                f"for step {step}"
            this_prod_th = obs1_3.gen_p_before_curtail[self.env1.gen_renewable]
            th_prod = obs2_3.gen_p[self.env1.gen_renewable]
            assert np.max(np.abs(this_prod_th - th_prod)) <= 1e-4, \
                f"for step {step}"
            assert obs1_3.gen_p[gen_id] <= ratio_max * obs1_3.gen_pmax[gen_id] + 1e-5, \
                f"for step {step}"
            amount_curtailed = (np.sum(this_prod_th - obs1_3.gen_p[self.env1.gen_renewable]))
            assert abs(np.sum(obs1_3.actual_dispatch) - amount_curtailed) <= 1e-4, \
                f"for step {step}"
            if amount_curtailed == 0.:
                # test that redispatch goes back to all O when no curtailment is done
                assert np.max(np.abs(obs1_3.actual_dispatch)) <= 1e-4,  f"for step {step}"
                tested = True
            assert np.all(obs1_3.curtailment[~obs1_3.gen_renewable] == 0.)
        if not tested:
            raise RuntimeError("the case where no curtailment have been used could not be tested. Please consider "
                               "making another test")

    def test_curtailment_mw_to_ratio(self):
        act = self.env1.action_space()
        gen_id = 2
        gen_id2 = 3
        amount_max_mw = 20.
        # normal test
        act.curtail = act.curtailment_mw_to_ratio([(gen_id, amount_max_mw)])
        assert np.array_equal(act.curtail, np.array([-1., -1., 0.2857143, -1., -1., -1.], dtype=dt_float))
        # test i can modify the feature a second time
        act.curtail = act.curtailment_mw_to_ratio([(gen_id, 2.0 * amount_max_mw)])
        assert np.array_equal(act.curtail, np.array([-1., -1., 2.0 * 0.2857143, -1., -1., -1.], dtype=dt_float))
        # test i modify another gen, it does not erase the first one
        act.curtail = act.curtailment_mw_to_ratio([(gen_id2, 0.5 * amount_max_mw)])
        assert np.array_equal(act.curtail, np.array([-1., -1., 2.0 * 0.2857143, 0.14285715, -1., -1.], dtype=dt_float))
        # test i cannot end up with results above 1
        act.curtail = act.curtailment_mw_to_ratio([(gen_id2, act.gen_pmax[gen_id2] + 10.)])
        assert np.array_equal(act.curtail, np.array([-1., -1., 2.0 * 0.2857143, 1.0, -1., -1.], dtype=dt_float))

        # test the property "curtail_mw"
        act2 = self.env1.action_space()
        act2.curtail_mw = [(gen_id, amount_max_mw)]
        assert np.array_equal(act2.curtail, np.array([-1., -1., 0.2857143, -1., -1., -1.], dtype=dt_float))
        assert np.array_equal(act2.curtail_mw, np.array([-1., -1., amount_max_mw, -1., -1., -1.], dtype=dt_float))
        # test i can modify the feature a second time
        act2.curtail_mw = [(gen_id, 2.0 * amount_max_mw)]
        assert np.array_equal(act2.curtail, np.array([-1., -1., 2.0 * 0.2857143, -1., -1., -1.], dtype=dt_float))
        assert np.array_equal(act2.curtail_mw, np.array([-1., -1., 2.0 * amount_max_mw, -1., -1., -1.], dtype=dt_float))
        # test i modify another gen, it does not erase the first one
        act2.curtail_mw = [(gen_id2, 0.5 * amount_max_mw)]
        assert np.array_equal(act2.curtail, np.array([-1., -1., 2.0 * 0.2857143, 0.14285715, -1., -1.], dtype=dt_float))
        assert np.array_equal(act2.curtail_mw, np.array([-1., -1., 2.0 * amount_max_mw, 0.5 * amount_max_mw, -1., -1.],
                                                        dtype=dt_float))
        # test i cannot end up with results above 1
        act2.curtail_mw = [(gen_id2, act2.gen_pmax[gen_id2] + 10.)]
        assert np.array_equal(act2.curtail, np.array([-1., -1., 2.0 * 0.2857143, 1.0, -1., -1.], dtype=dt_float))
        assert np.array_equal(act2.curtail_mw, np.array([-1., -1., 2.0 * amount_max_mw, act2.gen_pmax[gen_id2], -1., -1.],
                                                        dtype=dt_float))


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