import numpy as np
from gym.spaces import Box, Discrete
import tensorflow as tf

def space_n_to_shape_n(space_n):
    return np.array([space_to_shape(space) for space in space_n])

def space_to_shape(space):
    if type(space) is Box:
        return space.shape
    if type(space) is Discrete:
        return [space.n]

class Schedule(object):
    def value(self, step):
        """
        Value of the schedule for a given timestep
        FROM STABLE BASELINES
        :param step: (int) the timestep
        :return: (float) the output value for the given timestep
        """
        raise NotImplementedError


class LinearSchedule(Schedule):
    """
    Linear interpolation between initial_p and final_p over
    schedule_timesteps. After this many timesteps pass final_p is
    returned.
    FROM STABLE BASELINES
    :param schedule_timesteps: (int) Number of timesteps for which to linearly anneal initial_p to final_p
    :param initial_p: (float) initial output value
    :param final_p: (float) final output value
    """

    def __init__(self, schedule_timesteps, final_p, initial_p=1.0):
        self.schedule_timesteps = schedule_timesteps
        self.current_step = 0
        self.final_p = final_p
        self.initial_p = initial_p

    def value(self, step):
        fraction = min(float(step) / self.schedule_timesteps, 1.0)
        return self.initial_p + fraction * (self.final_p - self.initial_p)


def clip_by_local_norm(gradients, norm):
    """
    Clips gradients by their own norm, NOT by the global norm as it should be done (according to TF documentation).
    This here is the way MADDPG does it.
    """
    for idx, grad in enumerate(gradients):
        gradients[idx] = tf.clip_by_norm(grad, norm)
    return gradients

