from cs231n.layers import *
from cs231n.fast_layers import *

def affine_relu_batch_foward(x, w , b, gamma, beta, bn_param):
  fc_out, fc_cache = affine_forward(x, w, b)
  batch_out, batch_cache = batchnorm_forward(fc_out, gamma, beta, bn_param)
  relu_out, relu_cache = relu_forward(batch_out)
  '''
  if dropout_param['p']>0
    dropout_out, dropout_cache = dropout_forward(relu_out,dropout_param)
    out = dropout
    cache = (fc_cache,batch_cache, relu_cache, dropout_cache)
  else:
  '''
  out = relu_out
  cache = (fc_cache,batch_cache, relu_cache)
  return out, cache

def affine_relu_batch_backward(dout, cache):
  """
  Backward pass for the affine-batch norm-relu convenience layer
  """
  '''
  if dropout_param['p']>0:
    fc_cache,batch_cache, relu_cache, dropout_cache = cache
    dout_relu = dropout_backward(dout, dropout_cache)
  else:
  '''
    
  fc_cache, batch_cache, relu_cache = cache
  dout_relu = dout
  
  dbatch = relu_backward(dout_relu, relu_cache)  
  dfc, dgamma, dbeta = batchnorm_backward(dbatch, batch_cache)
  dx, dw, db = affine_backward(dfc, fc_cache)
  return dx, dw, db, dgamma, dbeta
  

def affine_relu_forward(x, w, b):
  """
  Convenience layer that perorms an affine transform followed by a ReLU

  Inputs:
  - x: Input to the affine layer
  - w, b: Weights for the affine layer

  Returns a tuple of:
  - out: Output from the ReLU
  - cache: Object to give to the backward pass
  """
  a, fc_cache = affine_forward(x, w, b)
  relu_out, relu_cache = relu_forward(a)
  '''
  if dropout_param['p'] > 0:
    out, dropout_cache = dropout_foward(relu_out, dropout_param)
    cache = (fc_cache, relu_cache, dropout_cache)
  else:
  '''
  out = relu_out
  cache = (fc_cache, relu_cache)
    
  return out, cache


def affine_relu_backward(dout, cache):
  """
  Backward pass for the affine-relu convenience layer
  """
  '''
  if dropout_param['p'] > 0:
    fc_cache, relu_cache, dropout_cache = cache
    dout_relu = dropout_backward(dout, dropout_cache)
  else:
  '''
  fc_cache, relu_cache = cache
  dout_relu = dout
    
  da = relu_backward(dout_relu, relu_cache)
  dx, dw, db = affine_backward(da, fc_cache)
  return dx, dw, db


pass


def conv_relu_forward(x, w, b, conv_param):
  """
  A convenience layer that performs a convolution followed by a ReLU.

  Inputs:
  - x: Input to the convolutional layer
  - w, b, conv_param: Weights and parameters for the convolutional layer
  
  Returns a tuple of:
  - out: Output from the ReLU
  - cache: Object to give to the backward pass
  """
  a, conv_cache = conv_forward_fast(x, w, b, conv_param)
  out, relu_cache = relu_forward(a)
  cache = (conv_cache, relu_cache)
  return out, cache


def conv_relu_backward(dout, cache):
  """
  Backward pass for the conv-relu convenience layer.
  """
  conv_cache, relu_cache = cache
  da = relu_backward(dout, relu_cache)
  dx, dw, db = conv_backward_fast(da, conv_cache)
  return dx, dw, db


def conv_relu_pool_forward(x, w, b, conv_param, pool_param):
  """
  Convenience layer that performs a convolution, a ReLU, and a pool.

  Inputs:
  - x: Input to the convolutional layer
  - w, b, conv_param: Weights and parameters for the convolutional layer
  - pool_param: Parameters for the pooling layer

  Returns a tuple of:
  - out: Output from the pooling layer
  - cache: Object to give to the backward pass
  """
  a, conv_cache = conv_forward_fast(x, w, b, conv_param)
  s, relu_cache = relu_forward(a)
  out, pool_cache = max_pool_forward_fast(s, pool_param)
  cache = (conv_cache, relu_cache, pool_cache)
  return out, cache


def conv_relu_pool_backward(dout, cache):
  """
  Backward pass for the conv-relu-pool convenience layer
  """
  conv_cache, relu_cache, pool_cache = cache
  ds = max_pool_backward_fast(dout, pool_cache)
  da = relu_backward(ds, relu_cache)
  dx, dw, db = conv_backward_fast(da, conv_cache)
  return dx, dw, db

