"""
This files provides a set of 2D environments with increasing difficulty
of exploration.

The difficulty is ranked by the level.

Important:
    * To create instances, use the function get_benchmark_env(level).
    * The horizon H is also set as an attribute of the environment.
"""

import numpy as np

from rlberry.wrappers.autoreset import AutoResetWrapper
from rlberry.envs.benchmarks.ball_exploration.pball import PBall2D


def get_benchmark_env(level=1):
    if level == 1:
        env = _get_autoreset_env(BallLevel1())
        return env
    elif level == 2:
        env = _get_autoreset_env(BallLevel2())
        return env
    elif level == 3:
        env = _get_autoreset_env(BallLevel3())
        return env
    elif level == 4:
        env = _get_autoreset_env(BallLevel4())
        return env
    elif level == 5:
        env = _get_autoreset_env(BallLevel5())
        return env
    else:
        raise NotImplementedError("Invalid benchmark level.")


def _get_autoreset_env(env):
    horizon = env.horizon
    return AutoResetWrapper(env, horizon)


#
# Level 1
#

class BallLevel1(PBall2D):
    """
    Dense rewards
    """
    def __init__(self):
        self.horizon = 30
        #
        self.p = 2
        self.action_list = [np.array([0.0, 0.0]),
                            0.05 * np.array([1.0, 0.0]),
                            -0.05 * np.array([1.0, 0.0]),
                            0.05 * np.array([0.0, 1.0]),
                            -0.05 * np.array([0.0, 1.0])]

        self.reward_amplitudes = np.array([1.0])
        self.reward_smoothness = np.array([0.5*np.sqrt(2)])
        self.reward_centers = [np.array([0.5, 0.5])]
        self.A = np.eye(2)
        self.B = np.eye(2)
        self.sigma = 0.01
        self.sigma_init = 0.001
        self.mu_init = np.array([0.0, 0.0])

        PBall2D.__init__(self,
                         self.p,
                         self.action_list,
                         self.reward_amplitudes,
                         self.reward_smoothness,
                         self.reward_centers,
                         self.A,
                         self.B,
                         self.sigma,
                         self.sigma_init,
                         self.mu_init)
        self.name = "Ball Exploration Benchmark - Level 1"


#
# Level 2
#

class BallLevel2(BallLevel1):
    """
    Sparse rewards
    """
    def __init__(self):
        BallLevel1.__init__(self)
        self.reward_amplitudes = np.array([1.0])
        self.reward_smoothness = np.array([0.2])
        self.reward_centers = [np.array([0.5, 0.5])]
        self.name = "Ball Exploration Benchmark - Level 2"

#
# Level 3
#


class BallLevel3(BallLevel2):
    """
    Sparse rewards, noisier
    """
    def __init__(self):
        BallLevel2.__init__(self)
        self.sigma = 0.025
        self.name = "Ball Exploration Benchmark - Level 3"


#
# Level 4
#


class BallLevel4(BallLevel1):
    """
    Far sparse reward (as lvl 2) + dense suboptimal rewards
    """
    def __init__(self):
        BallLevel1.__init__(self)

        self.reward_amplitudes = np.array([1.0, 0.1])
        self.reward_smoothness = np.array([0.2, 0.5*np.sqrt(2)])
        self.reward_centers = [np.array([-0.5, -0.5]),     # far sparse
                               np.array([0.5, 0.5])]       # dense
        self.name = "Ball Exploration Benchmark - Level 4"


#
# Level 5
#

class BallLevel5(BallLevel4):
    """
    Far sparse reward (as lvl 2) + dense suboptimal rewards, noisier
    """
    def __init__(self):
        BallLevel4.__init__(self)
        self.sigma = 0.025
        self.name = "Ball Exploration Benchmark - Level 5"


# if __name__ == '__main__':
#     env = get_benchmark_env(1)
#     env.enable_rendering()
#     for ii in range(100):
#         # env.step(1)
#         # env.step(3)
#         # env.step(env.action_space.sample())
#         # env.step(0)
#         env.step(4)

#     env.render()
