import numpy as np
from gym import spaces
from gym import Env


class PointEnv(Env):
    """
    point mass on a 2-D plane
    goals are sampled randomly from a square
    """

    def __init__(self, num_tasks=1):
        self.reset_task()
        self.reset()
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(2,))
        self.action_space = spaces.Box(low=-0.1, high=0.1, shape=(2,))

    def reset_task(self, is_evaluation=False):
        '''
        sample a new task randomly

        Problem 3: make training and evaluation goals disjoint sets
        if `is_evaluation` is true, sample from the evaluation set,
        otherwise sample from the training set
        '''
        #====================================================================================#
        #                           ----------PROBLEM 3----------
        #====================================================================================#
        # YOUR CODE HERE

        ## Granularity 1

        # odd_even = np.random.randint(2)
        # x_interval = np.random.randint(-5, high=5)
        # y_interval = np.random.randint(-5, high=5)
        # x = 2 * x_interval + np.random.uniform()
        # y = 2 * y_interval + np.random.uniform()

        # if is_evaluation:
        #     x = x + odd_even
        #     y = y + odd_even
        # else:
        #     x = x + odd_even
        #     y = y + (1 - odd_even)

        ## Granularity 2

        # zero_two = np.random.choice([0, 2])
        # x_interval = np.random.choice( range(-10, 10, 4) )
        # y_interval = np.random.choice( range(-10, 10, 4) )
        # x = x_interval + np.random.uniform(low=0, high=2)
        # y = y_interval + np.random.uniform(low=0, high=2)

        # if is_evaluation:
        #     x = x + zero_two
        #     y = y + zero_two
        # else:
        #     x = x + zero_two
        #     y = y + (2 - zero_two)

        ## Granularity 5

        # zero_five = np.random.choice([0, 5])
        # x_interval = np.random.choice( [-10, 0] )
        # y_interval = np.random.choice( [-10, 0] )
        # x = x_interval + np.random.uniform(low=0, high=5)
        # y = y_interval + np.random.uniform(low=0, high=5)

        # if is_evaluation:
        #     x = x + zero_five
        #     y = y + zero_five
        # else:
        #     x = x + zero_five
        #     y = y + (5 - zero_five)

        ## Granularity 10

        zero_ten = np.random.choice([0, 10])
        x = -10 + np.random.uniform(low=0, high=10)
        y = -10 + np.random.uniform(low=0, high=10)

        if is_evaluation:
            x = x + zero_ten
            y = y + zero_ten
        else:
            x = x + zero_ten
            y = y + (10 - zero_ten)

        # x = np.random.uniform(-10, 10)
        # y = np.random.uniform(-10, 10)

        self._goal = np.array([x, y])

    def reset(self):
        self._state = np.array([0, 0], dtype=np.float32)
        return self._get_obs()

    def _get_obs(self):
        return np.copy(self._state)

    def reward_function(self, x, y):
        return - (x ** 2 + y ** 2) ** 0.5

    def step(self, action):
        x, y = self._state
        # compute reward, add penalty for large actions instead of clipping them
        x -= self._goal[0]
        y -= self._goal[1]
        # check if task is complete
        done = abs(x) < .01 and abs(y) < .01
        reward = self.reward_function(x, y)
        # move to next state
        self._state = self._state + action
        ob = self._get_obs()
        return ob, reward, done, dict()

    def viewer_setup(self):
        print('no viewer')
        pass

    def render(self):
        print('current state:', self._state)

    def seed(self, seed):
        np.random.seed = seed
