import numpy as np


def affine_forward(x, w, b):
  """
  Computes the forward pass for an affine (fully-connected) layer.

  The input x has shape (N, d_1, ..., d_k) and contains a minibatch of N
  examples, where each example x[i] has shape (d_1, ..., d_k). We will
  reshape each input into a vector of dimension D = d_1 * ... * d_k, and
  then transform it to an output vector of dimension M.

  Inputs:
  - x: A numpy array containing input data, of shape (N, d_1, ..., d_k)
  - w: A numpy array of weights, of shape (D, M)
  - b: A numpy array of biases, of shape (M,)
  
  Returns a tuple of:
  - out: output, of shape (N, M)
  - cache: (x, w, b)
  """
  out = None
  #############################################################################
  # Done: Implement the affine forward pass. Store the result in out. You     #
  # will need to reshape the input into rows.                                 #
  #############################################################################
  x_row = np.reshape(x, (len(x), -1))
  out = np.dot(x_row, w) + b
  cache = (x, w, b)
  
  return out, cache


def affine_backward(dout, cache):
  """
  Computes the backward pass for an affine layer.

  Inputs:
  - dout: Upstream derivative, of shape (N, M)
  - cache: Tuple of:
    - x: Input data, of shape (N, d_1, ... d_k)
    - w: Weights, of shape (D, M)

  Returns a tuple of:
  - dx: Gradient with respect to x, of shape (N, d1, ..., d_k)
  - dw: Gradient with respect to w, of shape (D, M)
  - db: Gradient with respect to b, of shape (M,)
  """
  x, w, b = cache
  dx, dw, db = None, None, None
  #############################################################################
  # Done: Implement the affine backward pass.                                 #
  #############################################################################
  dx = np.dot(dout, w.T).reshape(x.shape)
  dw = np.dot(x.reshape(x.shape[0], np.prod(x.shape[1:])).T, dout)
  db = np.sum(dout, axis=0)

  return dx, dw, db


def relu_forward(x):
  """
  Computes the forward pass for a layer of rectified linear units (ReLUs).

  Input:
  - x: Inputs, of any shape

  Returns a tuple of:
  - out: Output, of the same shape as x
  - cache: x
  """
  out = None
  #############################################################################
  # Done: Implement the ReLU forward pass.                                    #
  #############################################################################
  out = np.maximum(0, x)
  cache = x
  
  return out, cache


def relu_backward(dout, cache):
  """
  Computes the backward pass for a layer of rectified linear units (ReLUs).

  Input:
  - dout: Upstream derivatives, of any shape
  - cache: Input x, of same shape as dout

  Returns:
  - dx: Gradient with respect to x
  """
  dx, x = None, cache
  #############################################################################
  # Done: Implement the ReLU backward pass.                                   #
  #############################################################################
  dx = dout * (x >= 0)
  
  return dx


def batchnorm_forward(x, gamma, beta, bn_param):
  """
  Forward pass for batch normalization.
  
  During training the sample mean and (uncorrected) sample variance are
  computed from minibatch statistics and used to normalize the incoming data.
  During training we also keep an exponentially decaying running mean of the mean
  and variance of each feature, and these averages are used to normalize data
  at test-time.

  At each timestep we update the running averages for mean and variance using
  an exponential decay based on the momentum parameter:

  running_mean = momentum * running_mean + (1 - momentum) * sample_mean
  running_var = momentum * running_var + (1 - momentum) * sample_var

  Note that the batch normalization paper suggests a different test-time
  behavior: they compute sample mean and variance for each feature using a
  large number of training images rather than using a running average. For
  this implementation we have chosen to use running averages instead since
  they do not require an additional estimation step; the torch7 implementation
  of batch normalization also uses running averages.

  Input:
  - x: Data of shape (N, D)
  - gamma: Scale parameter of shape (D,)
  - beta: Shift paremeter of shape (D,)
  - bn_param: Dictionary with the following keys:
    - mode: 'train' or 'test'; required
    - eps: Constant for numeric stability
    - momentum: Constant for running mean / variance.
    - running_mean: Array of shape (D,) giving running mean of features
    - running_var Array of shape (D,) giving running variance of features

  Returns a tuple of:
  - out: of shape (N, D)
  - cache: A tuple of values needed in the backward pass
  """
  mode = bn_param['mode']
  eps = bn_param.get('eps', 1e-5)
  momentum = bn_param.get('momentum', 0.9)

  N, D = x.shape
  running_mean = bn_param.get('running_mean', np.zeros(D, dtype=x.dtype))
  running_var = bn_param.get('running_var', np.zeros(D, dtype=x.dtype))

  out, cache = None, None
  if mode == 'train':
    #############################################################################
    # Done: Implement the training-time forward pass for batch normalization.   #
    # Use minibatch statistics to compute the mean and variance, use these      #
    # statistics to normalize the incoming data, and scale and shift the        #
    # normalized data using gamma and beta.                                     #
    #                                                                           #
    # You should store the output in the variable out. Any intermediates that   #
    # you need for the backward pass should be stored in the cache variable.    #
    #                                                                           #
    # You should also use your computed sample mean and variance together with  #
    # the momentum variable to update the running mean and running variance,    #
    # storing your result in the running_mean and running_var variables.        #
    #############################################################################
    
    # 1) compute mean    
    mu = 1.0/N * np.sum(x, axis = 0)
    
    # 2) helper to compute the variance
    x_mu = x - mu
    
    # 3) helper to compute the variance
    x_mu_square = x_mu ** 2
    
    # 4) compute the variance
    var = 1.0/N * np.sum(x_mu_square, axis = 0)
    
    # 5) compute the square root of the variance
    sqrt_var = np.sqrt(var + eps)
    
    # 6) compute its inverse (easier to compute the derivatives this way)
    sqrt_var_inv = 1.0/sqrt_var
    
    # 7) do the normalization
    x_hat = x_mu * sqrt_var_inv 
    
    # 8) do the scaling
    x_scaled = gamma * x_hat
    
    # 9) do the shifting (which is also the final output)
    out = x_scaled + beta
    
    # cache everything
    cache = (x_scaled, x_hat, sqrt_var_inv, sqrt_var, var, x_mu_square, \
        x_mu, mu, x, gamma, beta, bn_param)
    
    # compute the running mean and the running variance
    running_mean = momentum * running_mean + (1.0 - momentum) * mu
    running_var = momentum * running_var + (1.0 - momentum) * var
    
  elif mode == 'test':
    #############################################################################
    # Done: Implement the test-time forward pass for batch normalization. Use   #
    # the running mean and variance to normalize the incoming data, then scale  #
    # and shift the normalized data using gamma and beta. Store the result in   #
    # the out variable.                                                         #
    #############################################################################
    mu = running_mean
    var = running_var
    x_hat = (x - mu) / np.sqrt(var + eps)
    out = gamma * x_hat + beta
    cache = (mu, var, gamma, beta, bn_param)

  else:
    raise ValueError('Invalid forward batchnorm mode "%s"' % mode)

  # Store the updated running means back into bn_param
  bn_param['running_mean'] = running_mean
  bn_param['running_var'] = running_var

  return out, cache


def batchnorm_backward(dout, cache):
  """
  Backward pass for batch normalization.
  
  For this implementation, you should write out a computation graph for
  batch normalization on paper and propagate gradients backward through
  intermediate nodes.
  
  Inputs:
  - dout: Upstream derivatives, of shape (N, D)
  - cache: Variable of intermediates from batchnorm_forward.
  
  Returns a tuple of:
  - dx: Gradient with respect to inputs x, of shape (N, D)
  - dgamma: Gradient with respect to scale parameter gamma, of shape (D,)
  - dbeta: Gradient with respect to shift parameter beta, of shape (D,)
  """
  dx, dgamma, dbeta = None, None, None
  #############################################################################
  # Done: Implement the backward pass for batch normalization. Store the      #
  # results in the dx, dgamma, and dbeta variables.                           #
  #############################################################################
  x_scaled, x_hat, sqrt_var_inv, sqrt_var, var, x_mu_square, \
        x_mu, mu, x, gamma, beta, bn_param = cache

  N, D = dout.shape        
  eps = bn_param.get('eps', 1e-5)      
        
  # back-prop (9)   
  dx_scaled = dout
  dbeta = np.sum(dout, axis = 0)  

  # back-prop (8)   
  dx_hat = gamma * dx_scaled
  dgamma = np.sum(x_hat * dx_scaled, axis = 0)
  
  # back-prop (7)
  dx_mu = sqrt_var_inv * dx_hat
  dsqrt_var_inv = np.sum(x_mu * dx_hat, axis = 0)
  
  # back-prop (6)
  dsqrt_var = -1.0/(sqrt_var ** 2) * dsqrt_var_inv
  
  # back-prop (5)
  dvar = 0.5 * (var + eps) ** (-0.5) * dsqrt_var
  
  # back-prop (4)
  dx_mu_square = 1.0/N * np.ones((x_mu_square.shape)) * dvar
  
  # back-prop (3)
  dx_mu += 2 * x_mu * dx_mu_square
  
  # back-prop (2)
  dx = dx_mu
  dmu = -np.sum(dx_mu, axis = 0)
  
  # back-prop (1)
  dx += 1.0/N * np.ones((dmu.shape)) * dmu
  #############################################################################
  #                             END OF YOUR CODE                              #
  #############################################################################

  return dx, dgamma, dbeta


def batchnorm_backward_alt(dout, cache):
  """
  Alternative backward pass for batch normalization.
  
  For this implementation you should work out the derivatives for the batch
  normalizaton backward pass on paper and simplify as much as possible. You
  should be able to derive a simple expression for the backward pass.
  
  Note: This implementation should expect to receive the same cache variable
  as batchnorm_backward, but might not use all of the values in the cache.
  
  Inputs / outputs: Same as batchnorm_backward
  """
  dx, dgamma, dbeta = None, None, None
  x_scaled, x_hat, sqrt_var_inv, sqrt_var, var, x_mu_square, \
        x_mu, mu, x, gamma, beta, bn_param = cache

  N, D = dout.shape        
  eps = bn_param.get('eps', 1e-5)  
  #############################################################################
  # Done: Implement the backward pass for batch normalization. Store the      #
  # results in the dx, dgamma, and dbeta variables.                           #
  #                                                                           #
  # After computing the gradient with respect to the centered inputs, you     #
  # should be able to compute gradients with respect to the inputs in a       #
  # single statement; our implementation fits on a single 80-character line.  #
  #############################################################################
  dbeta = np.sum(dout, axis=0)
  dgamma = np.sum((x - mu) * (var + eps)**(-1. / 2.) * dout, axis=0)
  dx = (1. / N) * gamma * (var + eps)**(-1. / 2.) * \
       (N * dout - np.sum(dout, axis=0) - (x - mu) * \
       (var + eps)**(-1.0) * np.sum(dout * (x - mu), axis=0))
  #############################################################################
  #                             END OF YOUR CODE                              #
  #############################################################################
  
  return dx, dgamma, dbeta


def dropout_forward(x, dropout_param):
  """
  Performs the forward pass for (inverted) dropout.

  Inputs:
  - x: Input data, of any shape
  - dropout_param: A dictionary with the following keys:
    - p: Dropout parameter. We drop each neuron output with probability p.
    - mode: 'test' or 'train'. If the mode is train, then perform dropout;
      if the mode is test, then just return the input.
    - seed: Seed for the random number generator. Passing seed makes this
      function deterministic, which is needed for gradient checking but not in
      real networks.

  Outputs:
  - out: Array of the same shape as x.
  - cache: A tuple (dropout_param, mask). In training mode, mask is the dropout
    mask that was used to multiply the input; in test mode, mask is None.
  """
  p, mode = dropout_param['p'], dropout_param['mode']
  if 'seed' in dropout_param:
    np.random.seed(dropout_param['seed'])

  mask = None
  out = None

  if mode == 'train':
    ###########################################################################
    # Done: Implement the training phase forward pass for inverted dropout.   #
    # Store the dropout mask in the mask variable.                            #
    ###########################################################################
    mask = (np.random.rand(*x.shape) < p) / p
    out = x * mask

  elif mode == 'test':
    ###########################################################################
    # Done: Implement the test phase forward pass for inverted dropout.       #
    ###########################################################################
    out = x
  
  cache = (dropout_param, mask)
  out = out.astype(x.dtype, copy=False)

  return out, cache


def dropout_backward(dout, cache):
  """
  Perform the backward pass for (inverted) dropout.

  Inputs:
  - dout: Upstream derivatives, of any shape
  - cache: (dropout_param, mask) from dropout_forward.
  """
  dropout_param, mask = cache
  mode = dropout_param['mode']
  
  dx = None
  if mode == 'train':
    ###########################################################################
    # Done: Implement the training phase backward pass for inverted dropout.  #
    ###########################################################################
    dx = dout * mask
    
  elif mode == 'test':
    dx = dout
  return dx


def conv_forward_naive(x, w, b, conv_param):
  """
  A naive implementation of the forward pass for a convolutional layer.

  The input consists of N data points, each with C channels, height H and width
  W. We convolve each input with F different filters, where each filter spans
  all C channels and has height HH and width HH.

  Input:
  - x: Input data of shape (N, C, H, W)
  - w: Filter weights of shape (F, C, HH, WW)
  - b: Biases, of shape (F,)
  - conv_param: A dictionary with the following keys:
    - 'stride': The number of pixels between adjacent receptive fields in the
      horizontal and vertical directions.
    - 'pad': The number of pixels that will be used to zero-pad the input.

  Returns a tuple of:
  - out: Output data, of shape (N, F, H', W') where H' and W' are given by
    H' = 1 + (H + 2 * pad - HH) / stride
    W' = 1 + (W + 2 * pad - WW) / stride
  - cache: (x, w, b, conv_param)
  """
  out = None
  #############################################################################
  # TODO: Implement the convolutional forward pass.                           #
  # Hint: you can use the function np.pad for padding.                        #
  #############################################################################
  
  # extract the dimensions of the data structures
  N, C, H, W = x.shape
  F, C, HH, WW = w.shape
  stride = conv_param['stride']
  pad = conv_param['pad']
  
  # add padding
  x_pad = np.pad(x, [(0,0), (0,0), (pad,pad), (pad,pad)], 'constant')
  
  # calculate the dimensions for the next layer
  H_prime = 1 + (H + 2 * pad - HH)/stride
  W_prime = 1 + (W + 2 * pad - WW)/stride
  
  # build the new data structure which will contain the info after 
  # convolution happens
  out = np.zeros((N, F, H_prime, W_prime))
  
  # looooooop :S
  for i in range(N): #iterate over all the images
    for j in range(F): # iterate over all filters
      for k in range(H_prime): # iterate over all possibilities in height
        # see how much we have to jump each time in horizontal direction
        window_horizontal = k * stride  
        for l in range(W_prime): # iterate over all possibilities in width
          # see how much we have to jump each time in vertical direction
          window_vertical = l * stride 
          # build the window
          window = x_pad[i, :, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW]
          # do the sot product    
          out[i, j, k, l] = np.sum(window * w[j]) + b[j]

  cache = (x, w, b, conv_param)
  return out, cache


def conv_backward_naive(dout, cache):
  """
  A naive implementation of the backward pass for a convolutional layer.

  Inputs:
  - dout: Upstream derivatives.
  - cache: A tuple of (x, w, b, conv_param) as in conv_forward_naive

  Returns a tuple of:
  - dx: Gradient with respect to x
  - dw: Gradient with respect to w
  - db: Gradient with respect to b
  """
  dx, dw, db = None, None, None
  x, w, b, conv_param = cache # extract the data structures
  
  # extract the dimensions of the data structures and initialize the arrays
  # which will contain the derivatives of them
  N, C, H, W = x.shape
  dx = np.zeros((N, C, H, W))
  F, C, HH, WW = w.shape
  dw = np.zeros((F, C, HH, WW))
  db = np.zeros(F)
  stride = conv_param['stride']
  pad = conv_param['pad']
  
  # calculate the dimensions for the next layer
  H_prime = 1 + (H + 2 * pad - HH)/stride
  W_prime = 1 + (W + 2 * pad - WW)/stride
  
  # do the padding for both x and dx
  x_pad = np.pad(x, [(0,0), (0,0), (pad,pad), (pad,pad)], 'constant')
  dx_pad = np.pad(dx, [(0,0), (0,0), (pad,pad), (pad,pad)], 'constant')

  # backprop and loop a lot
  for i in xrange(N): #iterate over all the images
    for j in xrange(F): # iterate over all filters
      for k in xrange(H_prime): # iterate over all possibilities in height
        # see how much we have to jump each time in horizontal direction
        window_horizontal = k * stride  
        for l in xrange(W_prime): # iterate over all possibilities in width
          # see how much we have to jump each time in vertical direction
          window_vertical = l * stride 
          # build the window
          window = x_pad[i, :, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW]
          # compute the gradients
          db[j] += dout[i, j, k, l]
          dw[j] += window * dout[i, j, k, l]
          dx_pad[i, :, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW] += w[j] * dout[i, j, k, l]
          
  # unpad
  dx = dx_pad[:, :, pad:pad+H, pad:pad+W]
  
  return dx, dw, db


def max_pool_forward_naive(x, pool_param):
  """
  A naive implementation of the forward pass for a max pooling layer.

  Inputs:
  - x: Input data, of shape (N, C, H, W)
  - pool_param: dictionary with the following keys:
    - 'pool_height': The height of each pooling region
    - 'pool_width': The width of each pooling region
    - 'stride': The distance between adjacent pooling regions

  Returns a tuple of:
  - out: Output data
  - cache: (x, pool_param)
  """
  
  # extract the dimensions of the data structures
  N, C, H, W = x.shape
  HH = pool_param['pool_height']
  WW = pool_param['pool_width']
  stride = pool_param['stride']

  # calculate the dimensions for the next layer
  H_prime = 1 + (H - HH)/stride
  W_prime = 1 + (W - WW)/stride
  
  # build the new data structure which will contain the info after 
  # convolution happens
  out = np.zeros((N, C, H_prime, W_prime))

  # loooooooop
  for i in xrange(N): # iterate over each image
    for j in xrange(C): # iterate ober each channel
      for k in range(H_prime): # iterate over all possibilities in height
        # see how much we have to jump each time in horizontal direction
        window_horizontal = k * stride  
        for l in range(W_prime): # iterate over all possibilities in width
          # see how much we have to jump each time in vertical direction
          window_vertical = l * stride 
          # build the window
          window = x[i, j, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW] 
          out[i, j, k, l] = np.max(window)  

  cache = (x, pool_param)
  return out, cache


def max_pool_backward_naive(dout, cache):
  """
  A naive implementation of the backward pass for a max pooling layer.

  Inputs:
  - dout: Upstream derivatives
  - cache: A tuple of (x, pool_param) as in the forward pass.

  Returns:
  - dx: Gradient with respect to x
  """
  dx = None
  x, pool_param = cache # extract the data structures
  
  # extract the dimensions of the data structures and initialize the arrays
  # which will contain the derivatives of them
  N, C, H, W = x.shape
  HH = pool_param['pool_height']
  WW = pool_param['pool_width']
  stride = pool_param['stride']
  
  # calculate the dimensions for the next layer
  H_prime = 1 + (H - HH)/stride
  W_prime = 1 + (W - WW)/stride  
  
  dx = np.zeros((N, C, H, W))
  
  # loooooooop
  for i in xrange(N): # iterate over each image
    for j in xrange(C): # iterate ober each channel
      for k in range(H_prime): # iterate over all possibilities in height
        # see how much we have to jump each time in horizontal direction
        window_horizontal = k * stride  
        for l in range(W_prime): # iterate over all possibilities in width
          # see how much we have to jump each time in vertical direction
          window_vertical = l * stride 
          # build the window
          window = x[i, j, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW]
          # get the maximum value in the window  
          m = np.max(window) 
          # compute the derivative (all values which are equal to max multiply by the previous derivatives)
          dx[i, j, window_horizontal:window_horizontal + HH, \
            window_vertical:window_vertical + WW] += (window == m) * dout[i, j, k, l]
          
  return dx


def spatial_batchnorm_forward(x, gamma, beta, bn_param):
  """
  Computes the forward pass for spatial batch normalization.
  
  Inputs:
  - x: Input data of shape (N, C, H, W)
  - gamma: Scale parameter, of shape (C,)
  - beta: Shift parameter, of shape (C,)
  - bn_param: Dictionary with the following keys:
    - mode: 'train' or 'test'; required
    - eps: Constant for numeric stability
    - momentum: Constant for running mean / variance. momentum=0 means that
      old information is discarded completely at every time step, while
      momentum=1 means that new information is never incorporated. The
      default of momentum=0.9 should work well in most situations.
    - running_mean: Array of shape (D,) giving running mean of features
    - running_var Array of shape (D,) giving running variance of features
    
  Returns a tuple of:
  - out: Output data, of shape (N, C, H, W)
  - cache: Values needed for the backward pass
  """
  N, C, H, W = x.shape
  mode = bn_param['mode']
  eps = bn_param.get('eps', 1e-5)
  momentum = bn_param.get('momentum', 0.9)

  running_mean = bn_param.get('running_mean', np.zeros(C, dtype=x.dtype))
  running_var = bn_param.get('running_var', np.zeros(C, dtype=x.dtype))

  if mode == 'train':
  # Step 1 , calcul the average for each channel
    mu = (1. / (N * H * W) * np.sum(x, axis=(0, 2, 3))).reshape(1, C, 1, 1)
    var = (1. / (N * H * W) * np.sum((x - mu)**2,
                                         axis=(0, 2, 3))).reshape(1, C, 1, 1)
    xhat = (x - mu) / (np.sqrt(eps + var))
    out = gamma.reshape(1, C, 1, 1) * xhat + beta.reshape(1, C, 1, 1)

    running_mean = momentum * running_mean + (1.0 - momentum) * np.squeeze(mu)
    running_var = momentum * running_var + (1.0 - momentum) * np.squeeze(var)

    cache = (mu, var, x, xhat, gamma, beta, bn_param)

    # Store the updated running means back into bn_param
    bn_param['running_mean'] = running_mean
    bn_param['running_var'] = running_var

  elif mode == 'test':
    mu = running_mean.reshape(1, C, 1, 1)
    var = running_var.reshape(1, C, 1, 1)

    xhat = (x - mu) / (np.sqrt(eps + var))
    out = gamma.reshape(1, C, 1, 1) * xhat + beta.reshape(1, C, 1, 1)
    cache = (mu, var, x, xhat, gamma, beta, bn_param)
    
  else:
    raise ValueError('Invalid forward batchnorm mode "%s"' % mode)

  return out, cache


def spatial_batchnorm_backward(dout, cache):
  """
  Computes the backward pass for spatial batch normalization.
  Inputs:
  - dout: Upstream derivatives, of shape (N, C, H, W)
  - cache: Values from the forward pass
  Returns a tuple of:
  - dx: Gradient with respect to inputs, of shape (N, C, H, W)
  - dgamma: Gradient with respect to scale parameter, of shape (C,)
  - dbeta: Gradient with respect to shift parameter, of shape (C,)
  """
  dx, dgamma, dbeta = None, None, None

  mu, var, x, xhat, gamma, beta, bn_param = cache
  N, C, H, W = x.shape
  mode = bn_param['mode']
  eps = bn_param.get('eps', 1e-5)

  gamma = gamma.reshape(1, C, 1, 1)
  beta = beta.reshape(1, C, 1, 1)

  dbeta = np.sum(dout, axis=(0, 2, 3))
  dgamma = np.sum(dout * xhat, axis=(0, 2, 3))

  Nt = N * H * W
  dx = (1. / Nt) * gamma * (var + eps)**(-1. / 2.) * (
        Nt * dout - np.sum(dout, axis=(0, 2, 3)).reshape(1, C, 1, 1)
        - (x - mu) * (var + eps)**(-1.0) * np.sum(dout * (x - mu), axis=(0, 2, 3)).reshape(1, C, 1, 1))

  return dx, dgamma, dbeta
  

def svm_loss(x, y):
  """
  Computes the loss and gradient using for multiclass SVM classification.

  Inputs:
  - x: Input data, of shape (N, C) where x[i, j] is the score for the jth class
    for the ith input.
  - y: Vector of labels, of shape (N,) where y[i] is the label for x[i] and
    0 <= y[i] < C

  Returns a tuple of:
  - loss: Scalar giving the loss
  - dx: Gradient of the loss with respect to x
  """
  N = x.shape[0]
  correct_class_scores = x[np.arange(N), y]
  margins = np.maximum(0, x - correct_class_scores[:, np.newaxis] + 1.0)
  margins[np.arange(N), y] = 0
  loss = np.sum(margins) / N
  num_pos = np.sum(margins > 0, axis=1)
  dx = np.zeros_like(x)
  dx[margins > 0] = 1
  dx[np.arange(N), y] -= num_pos
  dx /= N
  return loss, dx


def softmax_loss(x, y):
  """
  Computes the loss and gradient for softmax classification.

  Inputs:
  - x: Input data, of shape (N, C) where x[i, j] is the score for the jth class
    for the ith input.
  - y: Vector of labels, of shape (N,) where y[i] is the label for x[i] and
    0 <= y[i] < C

  Returns a tuple of:
  - loss: Scalar giving the loss
  - dx: Gradient of the loss with respect to x
  """
  probs = np.exp(x - np.max(x, axis=1, keepdims=True))
  probs /= np.sum(probs, axis=1, keepdims=True)
  N = x.shape[0]
  loss = -np.sum(np.log(probs[np.arange(N), y])) / N
  dx = probs.copy()
  dx[np.arange(N), y] -= 1
  dx /= N
  return loss, dx
