/**
 * @file tests/ann_test_tools.hpp
 * @author Marcus Edel
 *
 * This file includes some useful functions for ann tests.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#ifndef MLPACK_TESTS_ANN_TEST_TOOLS_HPP
#define MLPACK_TESTS_ANN_TEST_TOOLS_HPP

#include <mlpack/core.hpp>
#include <mlpack/methods/ann/ffn.hpp>
#include <mlpack/methods/ann/init_rules/random_init.hpp>

using namespace mlpack;

// Approximate Jacobian and supposedly-true Jacobian, then compare them
// similarly to before.
template<typename ModuleType>
double JacobianTest(ModuleType& module,
                    arma::mat& input,
                    const double minValue = -2,
                    const double maxValue = -1,
                    const double perturbation = 1e-6)
{
  arma::mat output, outputA, outputB, jacobianA, jacobianB;
  output.set_size(module.OutputSize(), input.n_cols);
  outputA.set_size(module.OutputSize(), input.n_cols);
  outputB.set_size(module.OutputSize(), input.n_cols);

  // Initialize the input matrix.
  RandomInitialization init(minValue, maxValue);
  init.Initialize(input, input.n_rows, input.n_cols);

  // Initialize the jacobian matrix.
  module.Forward(input, output);
  jacobianA = arma::zeros(input.n_elem, output.n_elem);

  // Share the input paramter matrix.
  arma::mat sin = arma::mat(input.memptr(), input.n_rows, input.n_cols,
      false, false);

  for (size_t i = 0; i < input.n_elem; ++i)
  {
    double original = sin(i);
    sin(i) = original - perturbation;
    module.Forward(input, outputA);
    sin(i) = original + perturbation;
    module.Forward(input, outputB);
    sin(i) = original;

    outputB -= outputA;
    outputB /= 2 * perturbation;
    jacobianA.row(i) = outputB.t();
  }

  // Initialize the derivative parameter.
  arma::mat deriv = arma::zeros(output.n_rows, output.n_cols);

  // Share the derivative parameter.
  arma::mat derivTemp = arma::mat(deriv.memptr(), deriv.n_rows, deriv.n_cols,
      false, false);

  // Initialize the jacobian matrix.
  jacobianB = arma::zeros(input.n_elem, output.n_elem);

  for (size_t i = 0; i < derivTemp.n_elem; ++i)
  {
    deriv.zeros();
    derivTemp(i) = 1;

    arma::mat delta(input.n_rows, input.n_cols);
    module.Backward(input, output, deriv, delta);

    jacobianB.col(i) = delta;
  }

  return arma::max(arma::max(arma::abs(jacobianA - jacobianB)));
}

// Custom Jacobian Test where we get the input from outside of this function
// unlike the original Jacobian Test where input is generated inside that
// funcion.
template <typename ModuleType>
double CustomJacobianTest(ModuleType& module,
                          arma::mat& input,
                          const double perturbation = 1e-6)
{
  arma::mat output, outputA, outputB, jacobianA, jacobianB;
  output.set_size(module.OutputSize(), input.n_cols);
  outputA.set_size(module.OutputSize(), input.n_cols);
  outputB.set_size(module.OutputSize(), input.n_cols);

  // Initialize the jacobian matrix.
  module.Forward(input, output);
  jacobianA = arma::zeros(input.n_elem, output.n_elem);

  for (size_t i = 0; i < input.n_elem; ++i)
  {
    double original = input(i);
    input(i) = original - perturbation;
    module.Forward(input, outputA);
    input(i) = original + perturbation;
    module.Forward(input, outputB);
    input(i) = original;

    outputB -= outputA;
    outputB /= 2 * perturbation;
    jacobianA.row(i) = outputB.t();
  }

  // Initialize the derivative parameter.
  arma::mat deriv = arma::zeros(output.n_rows, output.n_cols);

  // Initialize the jacobian matrix.
  jacobianB = arma::zeros(input.n_elem, output.n_elem);

  for (size_t i = 0; i < deriv.n_elem; ++i)
  {
    deriv.zeros();
    deriv(i) = 1;

    arma::mat delta(input.n_rows, input.n_cols);
    module.Backward(input, output, deriv, delta);

    jacobianB.col(i) = delta;
  }

  return arma::max(arma::max(arma::abs(jacobianA - jacobianB)));
}

// Approximate Jacobian and supposedly-true Jacobian, then compare them
// similarly to before.
template<typename ModuleType>
double JacobianPerformanceTest(ModuleType& module,
                               arma::mat& input,
                               arma::mat& target,
                               const double eps = 1e-6)
{
  module.Forward(input, target);

  arma::mat delta;
  module.Backward(input, target, delta);

  arma::mat centralDifference = arma::zeros(delta.n_rows, delta.n_cols);
  arma::mat inputTemp = arma::mat(input.memptr(), input.n_rows, input.n_cols,
      false, false);

  arma::mat centralDifferenceTemp = arma::mat(centralDifference.memptr(),
      centralDifference.n_rows, centralDifference.n_cols, false, false);

  for (size_t i = 0; i < input.n_elem; ++i)
  {
    inputTemp(i) = inputTemp(i) + eps;
    double outputA = module.Forward(input, target);
    inputTemp(i) = inputTemp(i) - (2 * eps);
    double outputB = module.Forward(input, target);

    centralDifferenceTemp(i) = (outputA - outputB) / (2 * eps);
    inputTemp(i) = inputTemp(i) + eps;
  }

  return arma::max(arma::max(arma::abs(centralDifference - delta)));
}

template<typename ActivationFunction>
double ActivationJacobianTest(arma::mat& input,
                              const double minValue = -2,
                              const double maxValue = -1,
                              const double pertubation = 1e-6)
{
  arma::mat output, outputA, outputB, jacobianA, jacobianB;
  output.set_size(input.n_rows, input.n_cols);
  outputA.set_size(input.n_rows, input.n_cols);
  outputB.set_size(input.n_rows, input.n_cols);

  RandomInitialization init(minValue, maxValue);
  init.Initialize(input, input.n_rows, input.n_cols);

  ActivationFunction::Fn(input, output);
  jacobianA = arma::zeros(input.n_elem, output.n_elem);

  for (size_t i = 0; i < input.n_elem; ++i)
  {
    double original = input(i);
    input(i) = original - pertubation;
    ActivationFunction::Fn(input, outputA);
    input(i) = original + pertubation;
    ActivationFunction::Fn(input, outputB);
    input(i) = original;

    outputB -= outputA;
    outputB /= 2 * pertubation;
    jacobianA.row(i) = outputB.t();
  }

  arma::mat deriv = arma::zeros(output.n_rows, output.n_cols);
  jacobianB = arma::zeros(input.n_elem, output.n_elem);

  for (size_t i = 0; i < deriv.n_elem; ++i)
  {
    deriv.zeros();
    deriv(i) = 1;

    arma::mat delta(input.n_rows, input.n_cols);
    ActivationFunction::Deriv(input, output, delta);
    delta %= deriv;

    jacobianB.col(i) = delta;
  }

  return arma::max(arma::max(arma::abs(jacobianA - jacobianB)));
}

// Simple numerical gradient checker.
template<class FunctionType, typename MatType = arma::mat>
double CheckGradient(FunctionType& function, const double eps = 1e-7)
{
  // Get gradients for the current parameters.
  MatType orgGradient, gradient, estGradient;
  function.Gradient(orgGradient);

  estGradient = arma::zeros<MatType>(orgGradient.n_rows, orgGradient.n_cols);

  // Compute numeric approximations to gradient.
  for (size_t i = 0; i < orgGradient.n_elem; ++i)
  {
    double tmp = function.Parameters()(i);

    // Perturb parameter with a positive constant and get costs.
    function.Parameters()(i) += eps;
    double costPlus = function.Gradient(gradient);

    // Perturb parameter with a negative constant and get costs.
    function.Parameters()(i) -= (2 * eps);
    double costMinus = function.Gradient(gradient);

    // Restore the parameter value.
    function.Parameters()(i) = tmp;

    // Compute numerical gradients using the costs calculated above.
    estGradient(i) = (costPlus - costMinus) / (2 * eps);
  }

  // Estimate error of gradient.
  return arma::norm(orgGradient - estGradient) /
      arma::norm(orgGradient + estGradient);
}

// Simple numerical gradient checker for regularizers.
template<class FunctionType>
double CheckRegularizerGradient(FunctionType& function, const double eps = 1e-7)
{
  // Get gradients for the current parameters.
  arma::mat weight = arma::randu(10, 10);
  arma::mat orgGradient = arma::zeros(10 * 10, 1);
  function.Gradient(weight, orgGradient);

  arma::mat estGradient = arma::zeros(weight.n_rows, weight.n_cols);

  // Compute numeric approximations to gradient.
  for (size_t i = 0; i < weight.n_rows; ++i)
  {
    for (size_t j = 0; j < weight.n_cols; ++j)
    {
      double tmp = weight(i, j);

      weight(i, j) += eps;
      double costPlus = function.Output(weight, i, j);
      weight(i, j) -= (2 * eps);
      double costMinus = function.Output(weight, i, j);

      // Restore the weight value.
      weight(i, j) = tmp;
      estGradient(i, j) = (costPlus - costMinus) / (2 * eps);
    }
  }

  estGradient = vectorise(estGradient);
  // Estimate error of gradient.
  return arma::norm(orgGradient - estGradient) /
      arma::norm(orgGradient + estGradient);
}

/**
 * Train and evaluate a model with the specified structure, returning the
 * classification error.
 */
template<typename MatType = arma::mat, typename ModelType>
double TestClassificationNetwork(ModelType& model,
                                 MatType& trainData,
                                 MatType& trainLabels,
                                 MatType& testData,
                                 MatType& testLabels,
                                 const size_t maxEpochs)
{
  #if ENS_VERSION_MAJOR >= 3
  ens::RMSProp opt(0.32, 32, 0.88, 1e-8, trainData.n_cols * maxEpochs, -100);
  #else
  ens::RMSProp opt(0.01, 32, 0.88, 1e-8, trainData.n_cols * maxEpochs, -100);
  #endif
  model.Train(trainData, trainLabels, opt);

  MatType predictionTemp;
  model.Predict(testData, predictionTemp);
  MatType prediction = arma::zeros<MatType>(1, predictionTemp.n_cols);

  for (size_t i = 0; i < predictionTemp.n_cols; ++i)
  {
    prediction(i) = arma::as_scalar(arma::find(
        arma::max(predictionTemp.col(i)) == predictionTemp.col(i), 1));
  }

  size_t correct = accu(prediction == testLabels);

  return (1 - double(correct) / testData.n_cols);
}

// Utilities for RNN testing.

/**
 * Generate a super simple impulse whose response is a step function at the same
 * time step.  The impulse occurs at a random time in each dimension.
 *
 * Predicting this sequence is a super easy task for a recurrent network, but
 * not possible without a recurrent connection.
 */
template<typename ElemType = double>
void GenerateImpulseStepData(arma::Cube<ElemType>& data,
                             arma::Cube<ElemType>& responses,
                             const size_t dimensions,
                             const size_t numSequences,
                             const size_t seqLen)
{
  data.zeros(dimensions, numSequences, seqLen);
  responses.zeros(dimensions, numSequences, seqLen);

  for (size_t i = 0; i < numSequences; ++i)
  {
    for (size_t j = 0; j < dimensions; ++j)
    {
      const size_t impulseStep = RandInt(0, seqLen - 1);

      data(j, i, impulseStep) = ElemType(1);
      responses.subcube(j, i, impulseStep, j, i, seqLen - 1).fill(ElemType(1));
    }
  }
}

/**
 * Test that the recurrent layer is always able to learn to hold the output at 1
 * when the input impulse happens.
 */
template<typename RecurrentLayerType, typename ElemType = double>
ElemType ImpulseStepDataTest(const size_t dimensions, const size_t rho)
{
  arma::Cube<ElemType> data, responses;

  GenerateImpulseStepData(data, responses, dimensions, 1000, 50);

  arma::Cube<ElemType> trainData = data.cols(0, 699);
  arma::Cube<ElemType> trainResponses = responses.cols(0, 699);
  arma::Cube<ElemType> testData = data.cols(700, 999);
  arma::Cube<ElemType> testResponses = responses.cols(700, 999);

  RNN<MeanSquaredErrorType<arma::Mat<ElemType>>, ConstInitialization,
      arma::Mat<ElemType>> net(rho);
  net.template Add<RecurrentLayerType>(dimensions);

  const size_t numEpochs = 50;
  #if ENS_VERSION_MAJOR >= 3
  ens::RMSProp opt(0.096, 32, 0.9, 1e-08, 700 * numEpochs, 1e-5);
  #else
  ens::RMSProp opt(0.003, 32, 0.9, 1e-08, 700 * numEpochs, 1e-5);
  #endif

  net.Train(trainData, trainResponses, opt);

  arma::Cube<ElemType> testPreds;
  net.Predict(testData, testPreds);

  arma::Row<ElemType> testData1 = vectorise(testData.col(0)).t();
  arma::Row<ElemType> testPred1 = vectorise(testPreds.col(0)).t();
  arma::Row<ElemType> testResp1 = vectorise(testResponses.col(0)).t();

  // Compute the MSE of the test data.
  const ElemType error = std::sqrt(sum(square(
      vectorise(testPreds) - vectorise(testResponses)))) / testPreds.n_elem;

  return error;
}

#endif
