/**
 * @file methods/ann/rnn_impl.hpp
 * @author Marcus Edel
 *
 * Definition of the RNN class, which implements recurrent neural networks.
 *
 * 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_METHODS_ANN_RNN_IMPL_HPP
#define MLPACK_METHODS_ANN_RNN_IMPL_HPP

// In case it hasn't been included yet.
#include "rnn.hpp"
#include "layer/recurrent_layer.hpp"

namespace mlpack {

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::RNN(
    const size_t bpttSteps,
    const bool single,
    OutputLayerType outputLayer,
    InitializationRuleType initializeRule) :
    bpttSteps(bpttSteps),
    single(single),
    network(std::move(outputLayer), std::move(initializeRule))
{
  /* Nothing to do here */
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::RNN(
    const RNN& network) :
    bpttSteps(network.bpttSteps),
    single(network.single),
    network(network.network)
{
  // Nothing else to do.
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::RNN(
    RNN&& network) :
    bpttSteps(std::move(network.bpttSteps)),
    single(std::move(network.single)),
    network(std::move(network.network))
{
  // Nothing to do here.
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>&
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::operator=(const RNN& other)
{
  if (this != &other)
  {
    bpttSteps = other.bpttSteps;
    single = other.single;
    network = other.network;
    predictors.clear();
    responses.clear();
  }

  return *this;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>&
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::operator=(RNN&& other)
{
  if (this != &other)
  {
    bpttSteps = std::move(other.bpttSteps);
    single = std::move(other.single);
    network = std::move(other.network);
    predictors.clear();
    responses.clear();
  }

  return *this;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::~RNN()
{
  // Nothing special to do.
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename OptimizerType, typename... CallbackTypes>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Train(
    CubeType predictors,
    CubeType responses,
    OptimizerType& optimizer,
    CallbackTypes&&... callbacks)
{
  ResetData(std::move(predictors), std::move(responses), URowType());

  network.WarnMessageMaxIterations(optimizer, this->predictors.n_cols);

  // Ensure that the network can be used.
  network.CheckNetwork("RNN::Train()", this->predictors.n_rows, true, true);

  // Train the model.
  Timer::Start("rnn_optimization");
  const ElemType out = optimizer.Optimize(*this, network.Parameters(),
      callbacks...);
  Timer::Stop("rnn_optimization");

  Log::Info << "RNN::Train(): final objective of trained model is " << out
      << "." << std::endl;
  return out;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename OptimizerType, typename... CallbackTypes>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Train(
    CubeType predictors,
    CubeType responses,
    CallbackTypes&&... callbacks)
{
  OptimizerType optimizer;
  return Train(std::move(predictors), std::move(responses), optimizer,
      callbacks...);
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename OptimizerType, typename... CallbackTypes>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Train(
    CubeType predictors,
    CubeType responses,
    URowType sequenceLengths,
    OptimizerType& optimizer,
    CallbackTypes&&... callbacks)
{
  ResetData(std::move(predictors), std::move(responses),
      std::move(sequenceLengths));

  network.WarnMessageMaxIterations(optimizer, this->predictors.n_cols);

  // Ensure that the network can be used.
  network.CheckNetwork("RNN::Train()", this->predictors.n_rows, true, true);

  // Train the model.
  Timer::Start("rnn_optimization");
  const ElemType out = optimizer.Optimize(*this, network.Parameters(),
      callbacks...);
  Timer::Stop("rnn_optimization");

  Log::Info << "RNN::Train(): final objective of trained model is " << out
      << "." << std::endl;
  return out;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename OptimizerType, typename... CallbackTypes>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Train(
    CubeType predictors,
    CubeType responses,
    URowType sequenceLengths,
    CallbackTypes&&... callbacks)
{
  OptimizerType optimizer;
  return Train(std::move(predictors), std::move(responses),
      std::move(sequenceLengths), optimizer, callbacks...);
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Predict(
    const CubeType& predictors,
    CubeType& results,
    const size_t batchSize)
{
  // Ensure that the network is configured correctly.
  network.CheckNetwork("RNN::Predict()", predictors.n_rows, true, false);

  results.set_size(network.network.OutputSize(), predictors.n_cols,
      single ? 1 : predictors.n_slices);

  MatType inputAlias, outputAlias;
  for (size_t i = 0; i < predictors.n_cols; i += batchSize)
  {
    const size_t effectiveBatchSize = std::min(batchSize,
        size_t(predictors.n_cols) - i);

    // Since we aren't doing a backward pass, we don't actually need to store
    // the state for each time step---we can fit it all in one buffer.
    ResetMemoryState(0, effectiveBatchSize);

    // Iterate over all time steps.
    for (size_t t = 0; t < predictors.n_slices; ++t)
    {
      SetCurrentStep(t, (t == predictors.n_slices - 1), effectiveBatchSize,
          effectiveBatchSize);

      // Create aliases for the input and output.  If we are in single mode, we
      // always output into the same slice.
      MakeAlias(inputAlias, predictors.slice(t), predictors.n_rows,
          effectiveBatchSize, i * predictors.n_rows);
      MakeAlias(outputAlias, results.slice(single ? 0 : t), results.n_rows,
          effectiveBatchSize, i * results.n_rows);

      network.Forward(inputAlias, outputAlias);
    }
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Predict(
    const CubeType& predictors,
    CubeType& results,
    const URowType& sequenceLengths)
{
  // Ensure that the network is configured correctly.
  network.CheckNetwork("RNN::Predict()", predictors.n_rows, true, false);

  results.set_size(network.network.OutputSize(), predictors.n_cols,
      single ? 1 : predictors.n_slices);

  MatType inputAlias, outputAlias;
  for (size_t i = 0; i < predictors.n_cols; i++)
  {
    // Since we aren't doing a backward pass, we don't actually need to store
    // the state for each time step---we can fit it all in one buffer.
    ResetMemoryState(0, 1);

    // Iterate over all time steps.
    const size_t steps = sequenceLengths[i];
    for (size_t t = 0; t < steps; ++t)
    {
      SetCurrentStep(t, (t == steps - 1), 1, 1);

      // Create aliases for the input and output.  If we are in single mode, we
      // always output into the same slice.
      MakeAlias(inputAlias, predictors.slice(t), predictors.n_rows, 1,
          i * predictors.n_rows);
      MakeAlias(outputAlias, results.slice(single ? 0 : t), results.n_rows, 1,
          i * results.n_rows);

      network.Forward(inputAlias, outputAlias);
    }
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Reset(const size_t inputDimensionality)
{
  // This is a reimplementation of FFN::Reset() that correctly prints
  // "RNN::Reset()".
  network.Parameters().clear();

  if (inputDimensionality != 0)
  {
    network.CheckNetwork("RNN::Reset()", inputDimensionality, true, false);
  }
  else
  {
    const size_t inputDims = std::accumulate(network.InputDimensions().begin(),
        network.InputDimensions().end(), 0);
    network.CheckNetwork("RNN::Reset()", inputDims, true, false);
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Evaluate(
    const CubeType& predictors,
    const CubeType& responses)
{
  // Ensure that the network is configured correctly.
  network.CheckNetwork("RNN::Evaluate()", predictors.n_rows);

  // Add the loss of the network unrelated to output.
  ElemType lossSum = ElemType(network.network.Loss());

  // Reset recurrent memory state.
  ResetMemoryState(0, predictors.n_cols);

  // Iterate over all time slices.
  MatType forwardOutput;
  for (size_t t = 0; t < predictors.n_slices; t++)
  {
    SetCurrentStep(t, (t == predictors.n_slices - 1), predictors.n_cols,
        predictors.n_cols);
    // Do a forward pass and calculate the loss.
    network.Forward(predictors.slice(t), forwardOutput);
    if (!single || t == predictors.n_slices - 1)
      lossSum += network.outputLayer.Forward(forwardOutput,
          responses.slice(single ? 0 : t));
  }

  return lossSum;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Evaluate(
    const CubeType& predictors,
    const CubeType& responses,
    const URowType& sequenceLengths,
    const size_t batchSize)
{
  // Ensure that the network is configured correctly.
  network.CheckNetwork("RNN::Evaluate()", predictors.n_rows);

  if (sequenceLengths.n_elem > 0 && batchSize != 1 && single)
    throw std::invalid_argument(
        "Batch size must be 1 for ragged sequences in single mode!");

  // Add the loss of the network unrelated to output.
  ElemType lossSum = ElemType(network.network.Loss());

  CubeType reordPredictors, reordResponses;
  URowType reordLengths;
  if (batchSize > 1)
  {
    // Make copies of the arguments so they can be reordered while the orignal
    // arguments stay constant.
    reordPredictors = predictors;
    reordResponses = responses;
    reordLengths = sequenceLengths;
  }
  else
  {
    // Reordering isn't actually needed so we just make aliases.
    MakeAlias(reordPredictors, predictors, predictors.n_rows,
        predictors.n_cols, predictors.n_slices);
    MakeAlias(reordResponses, responses, responses.n_rows,
        responses.n_cols, responses.n_slices);
    MakeAlias(reordLengths, sequenceLengths, sequenceLengths.n_elem);
  }

  for (size_t i = 0; i < predictors.n_cols; i += batchSize)
  {
    size_t effectiveBatchSize = std::min(batchSize,
        size_t(predictors.n_cols) - i);

    // Reset recurrent memory state.
    ResetMemoryState(0, effectiveBatchSize);

    // Reorder the data so the sequence lengths are in descending order.
    if (batchSize > 1)
      ReorderBatch(i, effectiveBatchSize, reordPredictors, reordResponses,
          reordLengths);

    MatType forwardOutput, inputAlias, responseAlias;
    // Iterate over all time slices.
    size_t slices = reordLengths
        .subvec(i, i + effectiveBatchSize - 1).max();
    size_t activeBatchSize = effectiveBatchSize;
    for (size_t t = 0; t < slices; t++)
    {
      // Calculate the number of active points.
      CalculateActivePoints(activeBatchSize, i, reordLengths, t);

      SetCurrentStep(t, (t == slices - 1), effectiveBatchSize, activeBatchSize);

      // Get the input and response data.
      MakeAlias(inputAlias, reordPredictors.slice(t), predictors.n_rows,
          activeBatchSize, i * predictors.n_rows);

      // Do a forward pass and calculate the loss.
      network.Forward(inputAlias, forwardOutput);
      if (!single || t == slices - 1)
      {
        MakeAlias(responseAlias, reordResponses.slice((single ? 0 : t)),
            responses.n_rows, activeBatchSize, i * responses.n_rows);
        lossSum += network.outputLayer.Forward(forwardOutput, responseAlias);
      }
    }
  }

  return lossSum;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename Archive>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::serialize(Archive& ar, const uint32_t /* version */)
{
  // Note: if you define MLPACK_ANN_IGNORE_SERIALIZATION_WARNING, you had
  // better ensure that every layer you are serializing has had
  // CEREAL_REGISTER_TYPE() called somewhere.  See layer/serialization.hpp for
  // more information.
  #if !defined(MLPACK_ENABLE_ANN_SERIALIZATION) && \
      !defined(MLPACK_ANN_IGNORE_SERIALIZATION_WARNING)
    if (std::is_same_v<MatType, arma::mat>)
    {
      throw std::runtime_error("RNN::serialize(): Cannot serialize"
        " a neural network with type `arma::mat` if "
        "`MLPACK_ENABLE_ANN_SERIALIZATION` is not defined.");
    }

    (void) ar;
  #endif

  #if !defined(MLPACK_ENABLE_ANN_SERIALIZATION_FMAT) && \
      !defined(MLPACK_ANN_IGNORE_SERIALIZATION_WARNING)
    if (std::is_same_v<MatType, arma::fmat>)
    {
      throw std::runtime_error("RNN::serialize(): Cannot serialize"
        " a neural network with type `arma::fmat` if "
        "`MLPACK_ENABLE_ANN_SERIALIZATION_FMAT` is not defined.");
    }

    (void) ar;
  #endif

  #if defined(MLPACK_ENABLE_ANN_SERIALIZATION) || \
      defined(MLPACK_ENABLE_ANN_SERIALIZATION_FMAT) || \
      defined(MLPACK_ANN_IGNORE_SERIALIZATION_WARNING)

    ar(CEREAL_NVP(bpttSteps));
    ar(CEREAL_NVP(single));
    ar(CEREAL_NVP(network));

    if (Archive::is_loading::value)
    {
      // We can clear these members, since it's not possible to serialize in the
      // middle of training and resume.
      predictors.clear();
      responses.clear();
      sequenceLengths.clear();
    }
  #endif
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Evaluate(
    const MatType& /* parameters */,
    const size_t begin,
    const size_t batchSize)
{
  // Ensure the network is valid.
  network.CheckNetwork("RNN::Evaluate()", predictors.n_rows);

  if (sequenceLengths.n_elem > 0 && batchSize != 1 && single)
    throw std::invalid_argument(
        "Batch size must be 1 for ragged sequences in single mode!");

  // The core of the computation here is to pass through each step.  Since we
  // are not computing the gradient, we can be "clever" and use only one memory
  // cell---we don't need to know about the past.
  ResetMemoryState(0, batchSize);
  MatType output(network.network.OutputSize(), batchSize);

  // Reorder the data so the sequence lengths are in descending order.
  if (sequenceLengths.n_elem > 0 && batchSize > 1)
    ReorderBatch(begin, batchSize, predictors, responses, sequenceLengths);

  ElemType loss = 0;
  MatType stepData, responseData;
  const size_t steps = (sequenceLengths.n_elem == 0) ? predictors.n_slices :
      sequenceLengths.subvec(begin, begin + batchSize - 1).max();
  size_t activeBatchSize = batchSize;
  for (size_t t = 0; t < steps; ++t)
  {
    // Calculate the number of active points.
    if (sequenceLengths.n_elem > 0)
      CalculateActivePoints(activeBatchSize, begin, sequenceLengths, t);

    // Manually reset the data of the network to be an alias of the current time
    // step.
    SetCurrentStep(t, (t == steps), batchSize, activeBatchSize);

    MakeAlias(network.predictors, predictors.slice(t), predictors.n_rows,
        activeBatchSize, begin * predictors.slice(t).n_rows);
    const size_t responseStep = (single) ? 0 : t;
    MakeAlias(network.responses, responses.slice(responseStep),
        responses.n_rows, activeBatchSize,
        begin * responses.slice(responseStep).n_rows);

    loss += network.Evaluate(output, begin, activeBatchSize);
  }

  return loss;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename GradType>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::EvaluateWithGradient(
    const MatType& parameters,
    GradType& gradient)
{
  return EvaluateWithGradient(parameters, 0, gradient, predictors.n_cols);
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename GradType>
typename MatType::elem_type RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::EvaluateWithGradient(
    const MatType& /* parameters */,
    const size_t begin,
    GradType& gradient,
    const size_t batchSize)
{
  network.CheckNetwork("RNN::EvaluateWithGradient()", predictors.n_rows);

  if (sequenceLengths.n_elem > 0 && batchSize != 1 && single)
    throw std::invalid_argument(
        "Batch size must be 1 for ragged sequences in single mode!");

  ElemType loss = 0;

  // We must save anywhere between 1 and `bpttSteps` states, but we are limited
  // by `predictors.n_slices`.
  const size_t effectiveBPTTSteps = std::max(size_t(1),
      std::min(bpttSteps, size_t(predictors.n_slices)));

  ResetMemoryState(effectiveBPTTSteps, batchSize);

  // This will store the outputs of the network at each time step.  Note that we
  // only need to store `effectiveBPTTSteps` of output.  We will treat `outputs`
  // as a circular buffer.
  CubeType outputs(
      network.network.OutputSize(), batchSize, effectiveBPTTSteps);

  MatType stepData, outputData, responseData;

  // Initialize gradient.
  gradient.zeros(network.Parameters().n_rows, network.Parameters().n_cols);

  // Add loss (this is not dependent on time steps, and should only be added
  // once).  This is, e.g., regularizer loss, and other additive losses not
  // having to do with the output layer.
  loss += ElemType(network.network.Loss());

  // Reorder the data so the sequence lengths are in descending order.
  if (sequenceLengths.n_elem > 0 && batchSize > 1)
    ReorderBatch(begin, batchSize, predictors, responses, sequenceLengths);

  // For backpropagation through time, we must backpropagate for every
  // subsequence of length `bpttSteps`.  Before we've taken `bpttSteps` though,
  // we will be backpropagating shorter sequences.
  const size_t steps = (sequenceLengths.n_elem == 0) ? predictors.n_slices :
      sequenceLengths.subvec(begin, begin + batchSize - 1).max();
  size_t activeBatchSize = batchSize;
  for (size_t t = 0; t < steps; ++t)
  {
    // Calculate the number of active points.
    if (sequenceLengths.n_elem > 0)
      CalculateActivePoints(activeBatchSize, begin, sequenceLengths, t);

    SetCurrentStep(t, (t == (steps - 1)), batchSize, activeBatchSize);

    // Make an alias of the step's data for the forward pass.
    MakeAlias(stepData, predictors.slice(t), predictors.n_rows, activeBatchSize,
        begin * predictors.n_rows);
    MakeAlias(outputData, outputs.slice(t % effectiveBPTTSteps), outputs.n_rows,
        activeBatchSize);
    network.network.Forward(stepData, outputData);

    // Determine what the response should be.  If we are in single mode but not
    // at the end of the sequence, we don't do a backwards pass.
    if (single && t != steps - 1)
    {
      continue;
    }

    // Now backpropagate through time, starting with the current time step and
    // moving backwards.
    MatType error;
    for (size_t step = 0; step < std::min(t + 1, effectiveBPTTSteps); ++step)
    {
      SetCurrentStep(t - step, (step == 0), batchSize, activeBatchSize, true);

      if (step > 0)
      {
        // Past the first step, the error is zero; only recurrent terms matter.
        error.zeros();

        MakeAlias(stepData, predictors.slice(t - step), predictors.n_rows,
            activeBatchSize, begin * predictors.slice(t - step).n_rows);
        MakeAlias(outputData, outputs.slice((t - step) % effectiveBPTTSteps),
            outputs.n_rows, activeBatchSize);
      }
      else
      {
        // Otherwise, use the backward pass on the output layer to compute the
        // error.
        const size_t responseStep = (single) ? 0 : t - step;
        MakeAlias(stepData, predictors.slice(t - step), predictors.n_rows,
            activeBatchSize, begin * predictors.n_rows);
        MakeAlias(responseData, responses.slice(responseStep), responses.n_rows,
            activeBatchSize, begin * responses.n_rows);
        MakeAlias(outputData, outputs.slice((t - step) % effectiveBPTTSteps),
            outputs.n_rows, activeBatchSize);

        // We only need to do this on the first time step of BPTT.
        loss += network.outputLayer.Forward(outputData, responseData);

        // Compute the output error.
        network.outputLayer.Backward(outputData, responseData, error);
      }

      // Now backpropagate that error through the network, and compute the
      // gradient.
      //
      // TODO: note that we could avoid the copy of currentGradient by having
      // each layer *add* its gradient to `gradient`.  However that would
      // require some amount of refactoring.
      MatType networkDelta(predictors.n_rows, activeBatchSize,
          GetFillType<MatType>::none);
      GradType currentGradient(gradient.n_rows, gradient.n_cols,
          GetFillType<MatType>::zeros);
      network.network.Backward(stepData, outputData, error, networkDelta);
      network.network.Gradient(stepData, error, currentGradient);

      gradient += currentGradient;
    }
  }

  return loss;
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
template<typename GradType>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Gradient(
    const MatType& parameters,
    const size_t begin,
    GradType& gradient,
    const size_t batchSize)
{
  this->EvaluateWithGradient(parameters, begin, gradient, batchSize);
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::Shuffle()
{
  if (sequenceLengths.n_elem > 0)
  {
    ShuffleData(predictors, responses, sequenceLengths,
        predictors, responses, sequenceLengths);
  }
  else
  {
    ShuffleData(predictors, responses, predictors, responses);
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::ResetData(
    CubeType predictors,
    CubeType responses,
    URowType sequenceLengths)
{
  this->predictors = std::move(predictors);
  this->responses = std::move(responses);
  this->sequenceLengths = std::move(sequenceLengths);
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::ResetMemoryState(const size_t memorySize, const size_t batchSize)
{
  // Iterate over all layers and set the memory size.
  for (Layer<MatType>* l : network.Network())
  {
    // We can only call ClearRecurrentState() on RecurrentLayers.
    RecurrentLayer<MatType>* r =
        dynamic_cast<RecurrentLayer<MatType>*>(l);
    if (r != nullptr)
      r->ClearRecurrentState(memorySize, batchSize);
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::SetCurrentStep(const size_t step,
                  const bool end,
                  size_t batchSize,
                  size_t activeBatchSize,
                  bool backwards)
{
  // Iterate over all layers and set the current step.
  for (Layer<MatType>* l : network.Network())
  {
    // We can only call CurrentStep() on RecurrentLayers.
    RecurrentLayer<MatType>* r =
        dynamic_cast<RecurrentLayer<MatType>*>(l);
    if (r != nullptr)
    {
      r->CurrentStep(step, end);
      r->OnStepChanged(step, batchSize, activeBatchSize, backwards);
    }
  }
}

template<
    typename OutputLayerType,
    typename InitializationRuleType,
    typename MatType
>
void RNN<
    OutputLayerType,
    InitializationRuleType,
    MatType
>::ReorderBatch(const size_t begin,
                const size_t batchSize,
                CubeType& predictors,
                CubeType& responses,
                URowType& sequenceLengths)
{
  using UColType = typename GetUColType<MatType>::type;
  URowType batchLengths;
  MakeAlias(batchLengths, sequenceLengths, batchSize, begin);

  // Get the new ordering of this batch.
  UColType ordering = sort_index(batchLengths, "descending");

  // Reorder all slices to use the new ordering.
  MatType batchPredictors, batchResponses;
  for (size_t i = 0; i < predictors.n_slices; i++)
  {
    MakeAlias(batchPredictors, predictors.slice(i), predictors.n_rows,
        batchSize, begin * predictors.n_rows);
    MakeAlias(batchResponses, responses.slice(i), responses.n_rows,
        batchSize, begin * responses.n_rows);
    batchPredictors = batchPredictors.cols(ordering);
    batchResponses = batchResponses.cols(ordering);
  }
  batchLengths = batchLengths.cols(ordering);
}

} // namespace mlpack

#endif
