import numpy as np
import matplotlib.pyplot as plt
import h5py
from sklearn.metrics import mean_squared_error


def sigmoid(Z):
  A = 1 / (1 + np.exp(-Z))
  cache = Z

  return A, cache


def relu(Z):
  A = np.maximum(0, Z)

  assert (A.shape == Z.shape)

  cache = Z
  return A, cache


def relu_backward(dA, cache):
  Z = cache
  dZ = np.array(dA, copy=True)  # just converting dz to a correct object.

  # When z <= 0, you should set dz to 0 as well.
  dZ[Z <= 0] = 0

  assert (dZ.shape == Z.shape)

  return dZ


def sigmoid_backward(dA, cache):
  Z = cache

  s = 1 / (1 + np.exp(-Z))
  dZ = dA * s * (1 - s)

  assert (dZ.shape == Z.shape)

  return dZ


def initialize_parameters(n_x, n_h, n_y):
  np.random.seed(1)

  W1 = np.random.randn(n_h, n_x) * 0.01
  b1 = np.zeros((n_h, 1))
  W2 = np.random.randn(n_y, n_h) * 0.01
  b2 = np.zeros((n_y, 1))

  assert (W1.shape == (n_h, n_x))
  assert (b1.shape == (n_h, 1))
  assert (W2.shape == (n_y, n_h))
  assert (b2.shape == (n_y, 1))

  parameters = {"W1": W1,
                "b1": b1,
                "W2": W2,
                "b2": b2}

  return parameters


def initialize_parameters_deep(layer_dims):
  np.random.seed(1)
  parameters = {}
  L = len(layer_dims)  # number of layers in the network

  for l in range(1, L):
    parameters['W' + str(l)] = np.random.randn(layer_dims[l], layer_dims[l - 1]) / np.sqrt(layer_dims[l - 1])  # *0.01
    parameters['b' + str(l)] = np.zeros((layer_dims[l], 1))

    assert (parameters['W' + str(l)].shape == (layer_dims[l], layer_dims[l - 1]))
    assert (parameters['b' + str(l)].shape == (layer_dims[l], 1))

  return parameters


def linear_forward(A, W, b):
  Z = W.dot(A) + b

  assert (Z.shape == (W.shape[0], A.shape[1]))
  cache = (A, W, b)

  return Z, cache


def linear_activation_forward(A_prev, W, b, activation):
  if activation == "sigmoid":
    # Inputs: "A_prev, W, b". Outputs: "A, activation_cache".
    Z, linear_cache = linear_forward(A_prev, W, b)
    A, activation_cache = sigmoid(Z)

  elif activation == "relu":
    # Inputs: "A_prev, W, b". Outputs: "A, activation_cache".
    Z, linear_cache = linear_forward(A_prev, W, b)
    A, activation_cache = relu(Z)

  assert (A.shape == (W.shape[0], A_prev.shape[1]))
  cache = (linear_cache, activation_cache)

  return A, cache


def compute_cost(AL, Y):
  m = Y.shape[1]

  # Compute loss from aL and y.
  cost = -1 / m * np.sum(Y * np.log(AL) + (1 - Y) * np.log(1 - AL))
  # (1. / m) * np.sum(-np.dot(Y, np.log(AL).T) - np.dot(1 - Y, np.log(1 - AL).T))

  cost = np.squeeze(cost)  # To make sure your cost's shape is what we expect (e.g. this turns [[17]] into 17).
  assert (cost.shape == ())

  return cost


def linear_backward(dZ, cache):
  A_prev, W, b = cache
  m = A_prev.shape[1]

  dW = 1. / m * np.dot(dZ, A_prev.T)
  db = 1. / m * np.sum(dZ, axis=1, keepdims=True)
  dA_prev = np.dot(W.T, dZ)

  assert (dA_prev.shape == A_prev.shape)
  assert (dW.shape == W.shape)
  assert (db.shape == b.shape)

  return dA_prev, dW, db


def linear_activation_backward(dA, cache, activation):
  linear_cache, activation_cache = cache

  if activation == "relu":
    dZ = relu_backward(dA, activation_cache)
    dA_prev, dW, db = linear_backward(dZ, linear_cache)

  elif activation == "sigmoid":
    dZ = sigmoid_backward(dA, activation_cache)
    dA_prev, dW, db = linear_backward(dZ, linear_cache)

  return dA_prev, dW, db


def update_parameters(parameters, grads, learning_rate):
  L = len(parameters) // 2  # number of layers in the neural network

  # Update rule for each parameter. Use a for loop.
  for l in range(L):
    parameters["W" + str(l + 1)] = parameters["W" + str(l + 1)] - learning_rate * grads["dW" + str(l + 1)]
    parameters["b" + str(l + 1)] = parameters["b" + str(l + 1)] - learning_rate * grads["db" + str(l + 1)]

  return parameters


def predict(X, Y, parameters):
  m = X.shape[1]
  n = len(parameters) // 2  # number of layers in the neural network
  p = np.zeros((28, m), dtype=int)

  # Forward propagation
  probas, caches = L_model_forward_predict(X, parameters)

  count = 0
  prediction = []
  true = []
  probas = probas.T
  Y = Y.T
  for x in range(0, probas.shape[0]):
    ind1 = np.unravel_index(np.argmax(probas[x], axis=None), probas.shape)
    ind2 = np.unravel_index(np.argmax(Y[x], axis=None), Y.shape)
    prediction.append(ind1[1])
    true.append(ind2[1])

  for i in range(len(prediction)):
    if prediction[i] == true[i]:
      count += 1
  # print results
  # print ("predictions: " + str(p))
  # print ("true labels: " + str(y))
  print("Accuracy: %f" % (count / m * 100))

def predict_one(X, parameters):

  # Forward propagation
  probas, caches = L_model_forward_predict(X, parameters)
  probas = probas.T

  ind1 = np.unravel_index(np.argmax(probas, axis=None), probas.shape)
  return ind1[1]


def L_model_forward(
  X, parameters, D, keep_prob, hidden_layers_activation_fn="relu"):
  A = X  # since input matrix A0
  A = np.multiply(A, D[str(0)])
  A /= keep_prob[0]
  caches = []  # initialize the caches list
  L = len(parameters) // 2  # number of layer in the network

  for l in range(1, L):
    A_prev = A
    A, cache = linear_activation_forward(
      A_prev, parameters["W" + str(l)], parameters["b" + str(l)],
      hidden_layers_activation_fn)
    # shut down some units
    A = np.multiply(A, D[str(l)])
    # scale that value of units to keep expected value the same
    A /= keep_prob[l]
    caches.append(cache)

  AL, cache = linear_activation_forward(
    A, parameters["W" + str(L)], parameters["b" + str(L)], "sigmoid")
  AL = np.multiply(AL, D[str(L)])
  AL /= keep_prob[L]
  caches.append(cache)
  assert (AL.shape == (28, X.shape[1]))

  return AL, caches


def L_model_backward(
  AL, Y, caches, D, keep_prob, hidden_layers_activation_fn="relu"):
  Y = Y.reshape(AL.shape)
  L = len(caches)
  grads = {}

  # dA for output layer
  dAL = np.divide(AL - Y, np.multiply(AL, 1 - AL))
  dAL = np.multiply(dAL, D[str(L)])
  dAL /= keep_prob[L]

  grads["dA" + str(L - 1)], grads["dW" + str(L)], grads[
    "db" + str(L)] = linear_activation_backward(
    dAL, caches[L - 1], "sigmoid")
  grads["dA" + str(L - 1)] = np.multiply(
    grads["dA" + str(L - 1)], D[str(L - 1)])
  grads["dA" + str(L - 1)] /= keep_prob[L - 1]

  for l in range(L - 1, 0, -1):
    current_cache = caches[l - 1]
    grads["dA" + str(l - 1)], grads["dW" + str(l)], grads[
      "db" + str(l)] = linear_activation_backward(
      grads["dA" + str(l)], current_cache,
      hidden_layers_activation_fn)

    grads["dA" + str(l - 1)] = np.multiply(
      grads["dA" + str(l - 1)], D[str(l - 1)])
    grads["dA" + str(l - 1)] /= keep_prob[l - 1]

  return grads


def drop_out_matrices(layers_dims, m, keep_prob):
  np.random.seed(1)
  D = {}
  L = len(layers_dims)

  for l in range(L):
    # initialize the random values for the dropout matrix
    D[str(l)] = np.random.rand(layers_dims[l], m)
    # Convert it to 0/1 to shut down neurons corresponding to each element
    D[str(l)] = D[str(l)] < keep_prob[l]
    assert (D[str(l)].shape == (layers_dims[l], m))
  return D


def L_model_forward_predict(X, parameters):
  caches = []
  A = X
  L = len(parameters) // 2  # number of layers in the neural network

  # Implement [LINEAR -> RELU]*(L-1). Add "cache" to the "caches" list.
  for l in range(1, L):
    A_prev = A
    A, cache = linear_activation_forward(A_prev, parameters['W' + str(l)], parameters['b' + str(l)], activation="relu")

    caches.append(cache)

  # Implement LINEAR -> SIGMOID. Add "cache" to the "caches" list.
  AL, cache = linear_activation_forward(A, parameters['W' + str(L)], parameters['b' + str(L)], activation="sigmoid")

  caches.append(cache)

  assert (AL.shape == (28, X.shape[1]))

  return AL, caches


'''


def L_model_backward(AL, Y, caches):

  grads = {}
  L = len(caches)  # the number of layers
  m = AL.shape[1]
  Y = Y.reshape(AL.shape)  # after this line, Y is the same shape as AL

  # Initializing the backpropagation
  dAL = - (np.divide(Y, AL) - np.divide(1 - Y, 1 - AL))

  # Lth layer (SIGMOID -> LINEAR) gradients. Inputs: "AL, Y, caches". Outputs: "grads["dAL"], grads["dWL"], grads["dbL"]
  current_cache = caches[L - 1]
  grads["dA" + str(L)], grads["dW" + str(L)], grads["db" + str(L)] = linear_activation_backward(dAL, current_cache,
                                                                                                activation="sigmoid")

  for l in reversed(range(L - 1)):
    # lth layer: (RELU -> LINEAR) gradients.
    current_cache = caches[l]
    dA_prev_temp, dW_temp, db_temp = linear_activation_backward(grads["dA" + str(l + 2)], current_cache,
                                                                activation="relu")
    grads["dA" + str(l + 1)] = dA_prev_temp
    grads["dW" + str(l + 1)] = dW_temp
    grads["db" + str(l + 1)] = db_temp

  return grads


'''
