/*
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an "AS
 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 *  express or implied.  See the License for the specific language
 *  governing permissions and limitations under the License.
 *
 * For more about this software visit:
 *
 *      http://www.graphlab.ml.cmu.edu
 *
 */

/**
 *
 * \brief This application performs MAP inference on Markov Nets
 * provided in standard UAI file format via Dual-Decomposition.
 *
 *
 *  \author Dhruv Batra
 */

#ifndef __DD_GRLAB_HPP__
#define __DD_GRLAB_HPP__

#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_stl.hpp>

#include <Eigen/Dense>
#include "eigen_serialization.hpp"

#include <graphlab.hpp>
#include <graphlab/macros_def.hpp>

using namespace std;

/**
 * \brief Eigen library vectors are used to store the potentials (log-space)
 */
typedef Eigen::VectorXd factor_type;
typedef Eigen::VectorXd vec;
typedef Eigen::MatrixXd mat;

/**
 * \brief The convergence threshold for each message.  Smaller values
 * imply tighter convergence but slower execution.
 *
 */
double TOLERANCE = 0.01;

/////////////////////////////////////////////////////////////////////////
// Edge and Vertex data and Graph Type
/**
 * \brief There is a vertex for each factor in the graph AND each singleton
 */
struct vertex_data
{
    int nvars;              // Number of variables in this factor.
    int degree;             // Degree of this factor (same as nvars for higher-order factors).
    vector<int> cards;      // Cardinality of each variable.
    vector<int> neighbors;  // Vertex ids of the neighbors.
    vec potentials;         // Potentials for each configuration of the factor.

    int best_configuration; // Index of the best configuration at a subgradient step.
                            // TODO: Maybe replace best_configuration by beliefs for the high order variables?
                            // In which case, beliefs would be vector<vec> beliefs.
    vec beliefs;            // Posterior values for the configurations after averaging (projected DD, unary variables only).

    vertex_data(): nvars(0) {}

    void load(graphlab::iarchive& arc)
    {
      arc >> nvars >> degree >> cards >> neighbors >> potentials >> best_configuration >> beliefs;
    }
    void save(graphlab::oarchive& arc) const
    {
      arc << nvars << degree << cards << neighbors << potentials << best_configuration << beliefs;
    }
}; // end of vertex_data

/**
 * \brief There is an edge connecting each factor to each singleton
 * in its scope.
 */
struct edge_data
{
    int varid; // Do we need this? (afm)
    int card; // Do we need this? (afm)
    vec potentials; // TODO: Unary potentials distributed evenly through the edges (i.e. unary potentials divided by degree).

    vec multiplier_messages; // Dual variables, i.e. Lagrangian multipliers.
    vec local_messages;      // Local MAP variables (for projected DD).

    edge_data(): varid(0), card(0) {}

    void load(graphlab::iarchive& arc) {
      arc >> varid >> card >> potentials >> multiplier_messages >> local_messages;
    }
    void save(graphlab::oarchive& arc) const {
      arc << varid << card << potentials << multiplier_messages << local_messages;
    }
};

/**
 * The graph type
 */
typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

/**
 * \brief The Dual Decomposition Vertex Program.
 */
struct dd_vertex_program :
  public graphlab::ivertex_program< graph_type, factor_type,
                                    graphlab::messages::sum_priority >,
  public graphlab::IS_POD_TYPE {

  /////////////////////////////////////////////////////////////////////////
  // Find the configuration index of a factor given the array of states.
  /////////////////////////////////////////////////////////////////////////
  int get_configuration_index(const graph_type::vertex_type& vertex,
                              const std::vector<int>& states) const {
    const vertex_data& vdata = vertex.data();
    int index = states[0];
    for (size_t i = 1; i < states.size(); ++i) {
      index *= vdata.cards[i];
      index += states[i];
    }
    return index;
  }

  /////////////////////////////////////////////////////////////////////////
  // Find the array of states corresponding to a factor configuration index.
  /////////////////////////////////////////////////////////////////////////
  void get_configuration_states(const graph_type::vertex_type& vertex,
                                int index, std::vector<int>* states) const {
    const vertex_data& vdata = vertex.data();
    int tmp = 1;
    for (size_t i = 1; i < states->size(); ++i) {
      tmp *= vdata.cards[i];
    }
    (*states)[0] = index / tmp;
    for (size_t i = 1; i < states->size(); ++i) {
      index = index % tmp;
      tmp /= vdata.cards[i];
      (*states)[i] = index / tmp;
    }
  }

  /**
   * \brief Given an edge and a vertex return the other vertex along
   * that edge.
   */
  inline vertex_type get_other_vertex(edge_type& edge,
                                      const vertex_type& vertex) const {
    return vertex.id() == edge.source().id()? edge.target() : edge.source();
  }; // end of other_vertex

  virtual edge_dir_type gather_edges(icontext_type& context,
                                     const vertex_type& vertex) const = 0;
  virtual factor_type gather(icontext_type& context, const vertex_type& vertex,
                             edge_type& edge) const = 0;
  virtual void apply(icontext_type& context, vertex_type& vertex,
                     const factor_type& total) = 0;
  virtual edge_dir_type scatter_edges(icontext_type& context,
                                      const vertex_type& vertex) const = 0;
  virtual void scatter(icontext_type& context, const vertex_type& vertex,
                       edge_type& edge) const = 0;
}; // end of class bp_vertex_program

////////////////////////////////////////////////////////////////////////////////
// This class implements the "symmetric" version of dual decomposition described
// in:
// D. Sontag, A. Globerson, T. Jaakkola.
// Introduction to Dual Decomposition for Inference.
// Optimization for Machine Learning, editors S. Sra, S. Nowozin, and S. J.
// Wright: MIT Press, 2011
////////////////////////////////////////////////////////////////////////////////

struct dd_vertex_program_symmetric : public dd_vertex_program {
  /**
   * \brief Since the MRF is undirected we will use all edges for gather and
   * scatter
   */
  edge_dir_type gather_edges(icontext_type& context,
                             const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of gather_edges

  /**
   * \brief The gather function takes a vertex and an edge as inputs and outputs
   a vector of numeric values. Vectors of numeric values will later be summed
   over all edges incident in this vertex. So, if the vertex is a unary factor,
   we can just return the vector of Lagrange multipliers stored in "edge.messages".
   Otherwise (if vertex is a general factor), things are a little more tricky.
   Suppose the factor is linked to K variables, with cardinalities C_1, ..., C_K.
   Suppose this edge is with respect to the k-th variable. Then, we return a
   vector of size C_1 + ... + C_K which is zero everywhere except in the
   k-th slot, where the Lagrange multipliers in "edge.messages" will be copied
   to. This way, when the "gather sum" takes place, and since all these slots
   are disjoint, we will just get the Lagrange multipliers of all the variables.
   */
  factor_type gather(icontext_type& context, const vertex_type& vertex,
                     edge_type& edge) const {
    cout << "gather begin" << endl;
    const vertex_type other_vertex = get_other_vertex(edge, vertex);
    const vertex_data& vdata = vertex.data();
    edge_data& edata = edge.data();

    if (vdata.nvars == 1) {
      // Unary factor.
      cout << "This unary factor has " << vertex.num_in_edges() <<
        " in edges and " << vertex.num_out_edges() << " out edges" << endl;
      return edata.multiplier_messages;
    } else {
      // General factor.
      factor_type messages;
      messages.resize(vdata.potentials.size());
      int offset = 0;
      int index_neighbor = -1;
      for (int k = 0; k < vdata.nvars; ++k) {
        int vertex_id = vdata.neighbors[k];
        if (vertex_id == other_vertex.id()) {
          index_neighbor = k;
          break;
        }
        offset += vdata.cards[k];
      }
      CHECK_GE(index_neighbor, 0);
      for (int state = 0; state < vdata.cards[index_neighbor]; ++state) {
        messages[offset + state] = -edata.multiplier_messages[state];
      }
      return messages;
    }
    cout << "gather end" << endl;
  }; // end of gather function

  /**
   * \brief The apply function takes a vertex and a vector of numeric values
   (a total) as input. For unary vertices, this will be the sum of Lagrange
   multipliers, and we just need to sum that to the vertex potential and compute
   the argmax. For general factors, the vector of numeric values, as stated above,
   will contain all the Lagrange multipliers of the neighboring variables.
   So we need to loop through all possible factor configurations, get the
   sequence of states of each configuration, fetch the Lagrange multipliers for
   those states, and add them to the factor potential. Then we compute the argmax.
   */
  void apply(icontext_type& context, vertex_type& vertex,
             const factor_type& total) {
    vertex_data& vdata = vertex.data();
    cout << "begin apply" << endl;
    if (vdata.nvars == 1) {
      // Unary factor.
      ASSERT_EQ(vdata.potentials.size(), total.size());
      vec belief = vdata.potentials + total;
      // Save the best configuration for this vertex.
      belief.maxCoeff(&vdata.best_configuration);
      cout << "vdata.best_configuration = " << vdata.best_configuration << endl;
    } else {
      // General factor.
      vec belief = vdata.potentials;
      int num_configurations = vdata.potentials.size();
      for (int index_configuration = 0;
           index_configuration < num_configurations;
           ++index_configuration) {
        vector<int> states(vdata.nvars, -1);
        // This could be made more efficient by defining an iterator over factor
        // configurations.
        get_configuration_states(vertex, index_configuration, &states);
        int offset = 0;
        for (int k = 0; k < vdata.nvars; ++k) {
          belief[index_configuration] += total[offset + states[k]];
          offset += vdata.cards[k];
        }
      }
      // Save the best configuration for this factor.
      belief.maxCoeff(&vdata.best_configuration);
    }
    cout << "end apply" << endl;
  }; // end of apply

  /**
   * \brief Since the MRF is undirected we will use all edges for gather and
   * scatter
   */
  edge_dir_type scatter_edges(icontext_type& context,
                              const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of scatter edges

  /**
   * \brief The scatter function takes a vertex and an edge as input.
   We just need to update the messages (Lagrange multipliers) by looking at the
   saved argmaxes.
   */
  void scatter(icontext_type& context, const vertex_type& vertex,
               edge_type& edge) const {
    const vertex_type other_vertex = get_other_vertex(edge, vertex);
    const vertex_type *unary_vertex;
    const vertex_type *factor_vertex;
    cout << "begin scatter" << endl;
    if (vertex.data().nvars == 1) {
      // Unary factor.
      unary_vertex = &vertex;
      factor_vertex = &other_vertex;
    } else {
      // General factor.
      unary_vertex = &other_vertex;
      factor_vertex = &vertex;
    }
    const vertex_data& vdata = unary_vertex->data();
    const vertex_data& other_vdata = factor_vertex->data();
    edge_data& edata = edge.data();
    double stepsize = 1.0; // TODO: Make this decay over iteration number.

    CHECK_GE(vdata.best_configuration, 0);
    CHECK_LT(vdata.best_configuration, vdata.cards[0]);
    edata.multiplier_messages[vdata.best_configuration] += stepsize;
    vector<int> states(other_vdata.nvars, -1);
    get_configuration_states(*factor_vertex, other_vdata.best_configuration, &states);
    int offset = 0;
    int index_neighbor = -1;
    for (int k = 0; k < other_vdata.nvars; ++k) {
      int vertex_id = other_vdata.neighbors[k];
      if (vertex_id == unary_vertex->id()) {
        index_neighbor = k;
        break;
      }
      offset += other_vdata.cards[k];
    }
    CHECK_GE(index_neighbor, 0);
    CHECK_GE(states[index_neighbor], 0);
    CHECK_LT(states[index_neighbor], other_vdata.cards[index_neighbor]);
    CHECK_EQ(other_vdata.cards[index_neighbor], vdata.cards[0]);
    edata.multiplier_messages[states[index_neighbor]] -= stepsize;
    cout << "end scatter" << endl;
  }; // end of scatter
}; // end of class bp_vertex_program_symmetric

////////////////////////////////////////////////////////////////////////////////
// This class implements the "projected" version of dual decomposition described
// in:
// Komodakis, N., Paragios, N., and Tziritas, G. (2007).
// MRF optimization via dual decomposition: Message-passing revisited.
// In Proc. of International Conference on Computer Vision.
//
// The formulation used is the one in Algorithm 1 of:
//
// André F. T. Martins, Mário A. T. Figueiredo, Pedro M. Q. Aguiar,
// Noah A. Smith, and Eric P. Xing.
// "An Augmented Lagrangian Approach to Constrained MAP Inference."
// International Conference on Machine Learning (ICML), 2011.
////////////////////////////////////////////////////////////////////////////////

struct dd_vertex_program_projected : public dd_vertex_program {
  /**
   * \brief Since the MRF is undirected we will use all edges for gather and
   * scatter
   */
  edge_dir_type gather_edges(icontext_type& context,
                             const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of gather_edges

  /**
   * \brief The gather function takes a vertex and an edge as inputs and outputs
   a vector of numeric values. Vectors of numeric values will later be summed
   over all edges incident in this vertex.
   If the vertex is a unary factor, compute the sum of all the local MAP variables,
   which in the "apply" function will serve to compute the global MAP.
   Otherwise (if vertex is a general factor), things are a little more tricky.
   Suppose the factor is linked to K variables, with cardinalities C_1, ..., C_K.
   Suppose this edge is with respect to the k-th variable. Then, we return a
   vector of size C_1 + ... + C_K which is zero everywhere except in the
   k-th slot, where the Lagrange multipliers in "edge.messages" will be copied
   to. This way, when the "gather sum" takes place, and since all these slots
   are disjoint, we will just get the Lagrange multipliers of all the variables.
   */
  factor_type gather(icontext_type& context, const vertex_type& vertex,
                     edge_type& edge) const {
    cout << "gather begin" << endl;
    const vertex_type other_vertex = get_other_vertex(edge, vertex);
    const vertex_data& vdata = vertex.data();
    edge_data& edata = edge.data();

    if (vdata.nvars == 1) {
      // Unary factor.
      cout << "This unary factor has " << vertex.num_in_edges() <<
        " in edges and " << vertex.num_out_edges() << " out edges" << endl;
      factor_type messages = edata.local_messages;
      return messages;
    } else {
      // General factor.
      factor_type messages;
      messages.resize(vdata.potentials.size());
      int offset = 0;
      int index_neighbor = -1;
      for (int k = 0; k < vdata.nvars; ++k) {
        int vertex_id = vdata.neighbors[k];
        if (vertex_id == other_vertex.id()) {
          index_neighbor = k;
          break;
        }
        offset += vdata.cards[k];
      }
      CHECK_GE(index_neighbor, 0);
      //const vec &unary_potential = other_vertex.data().potential;
      //int degree = other_vertex.data().degree;

      for (int state = 0; state < vdata.cards[index_neighbor]; ++state) {
        messages[offset + state] = edata.multiplier_messages[state];
        // TODO: somehow set the "edge potential" to be the potential of the
        // unary variable divided by the number of factors in which that
        // variable appears.
        messages[offset + state] += edata.potentials[state];
        //message[offset + state] += unary_potential[state] / static_cast<double>(degree);
      }
      return messages;
    }
    cout << "gather end" << endl;
  }; // end of gather function

  /**
   * \brief The apply function takes a vertex and a vector of numeric values
   (a total) as input.
   For a unary vertex, "total" will be the sum of local MAP vectors, and we
   just need to divide by the vertex degree and save the result as global MAP.
   For higher-order factors, "total" will contain all the Lagrange multipliers
   of the neighboring variables. So we need to loop through all possible factor
   configurations, get the sequence of states of each configuration, fetch the
   Lagrange multipliers for those states, and add them to the factor potential.
   Then we compute the argmax and save result to local MAP for each variable
   connected to the factor.
   */
  void apply(icontext_type& context, vertex_type& vertex,
             const factor_type& total) {
    vertex_data& vdata = vertex.data();
    cout << "begin apply" << endl;
    if (vdata.nvars == 1) {
      // Unary factor. Divide by vertex degree.
      vdata.beliefs = total / static_cast<double>(vdata.degree);
      return;
    } else {
      // General factor.
      vec beliefs = vdata.potentials;
      int num_configurations = vdata.potentials.size();
      for (int index_configuration = 0;
           index_configuration < num_configurations;
           ++index_configuration) {
        vector<int> states(vdata.nvars, -1);
        // This could be made more efficient by defining an iterator over factor
        // configurations.
        get_configuration_states(vertex, index_configuration, &states);
        int offset = 0;
        for (int k = 0; k < vdata.nvars; ++k) {
          beliefs[index_configuration] += total[offset + states[k]];
          offset += vdata.cards[k];
        }
      }
      // Save the best configuration for this factor.
      beliefs.maxCoeff(&vdata.best_configuration);
    }
    cout << "end apply" << endl;
  }; // end of apply

  /**
   * \brief Since the MRF is undirected we will use all edges for gather and
   * scatter
   */
  edge_dir_type scatter_edges(icontext_type& context,
                              const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }; // end of scatter edges

  /**
   * \brief The scatter function takes a vertex and an edge as input.
   (1) If the vertex is a unary factor, we update the messages (Lagrange multipliers)
   by subtracting the global MAP by the local MAP.
   (2) If the vertex is a higher order factor, this function will take the best
   configuration (obtained at the apply function) and save the local MAP
   at the corresponding edge.
   */
  void scatter(icontext_type& context, const vertex_type& vertex,
               edge_type& edge) const {
    const vertex_data& vdata = vertex.data();
    edge_data& edata = edge.data();
    cout << "begin scatter" << endl;
    if (vdata.nvars == 1) {
      // Unary factor. Update the messages (Lagrange multipliers).
      double stepsize = 1.0; // TODO: Make this decay over iteration number.
      edata.multiplier_messages += (vdata.beliefs - edata.local_messages) * stepsize;
    } else {
      // General factor. Update the local MAPs.
      const vertex_type &unary_vertex = get_other_vertex(edge, vertex);
      vector<int> states(vdata.nvars, -1);
      get_configuration_states(vertex, vdata.best_configuration, &states);
      int offset = 0;
      int index_neighbor = -1;
      for (int k = 0; k < vdata.nvars; ++k) {
        int vertex_id = vdata.neighbors[k];
        if (vertex_id == unary_vertex.id()) {
          index_neighbor = k;
          break;
        }
        offset += vdata.cards[k];
      }
      CHECK_GE(index_neighbor, 0);
      CHECK_GE(states[index_neighbor], 0);
      CHECK_LT(states[index_neighbor], vdata.cards[index_neighbor]);
      CHECK_EQ(vdata.cards[index_neighbor], unary_vertex.data().cards[0]);
      edata.local_messages.setZero();
      edata.local_messages[states[index_neighbor]] += 1.0;
    }
    cout << "end scatter" << endl;
  }; // end of scatter
}; // end of class dd_vertex_program_projected

#endif


// Copyright (c) 2014-2018, The Monero Project
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification, are
// permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this list of
//    conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice, this list
//    of conditions and the following disclaimer in the documentation and/or other
//    materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its contributors may be
//    used to endorse or promote products derived from this software without specific
//    prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL
// THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
// THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Parts of this file are originally copyright (c) 2012-2013 The Cryptonote developers

#include <unistd.h>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <boost/thread/mutex.hpp>
#include <boost/thread/lock_guard.hpp>
#include <boost/shared_ptr.hpp>

#include "common/varint.h"
#include "warnings.h"
#include "crypto.h"
#include "hash.h"

// Add some log capabilities
#include "misc_log_ex.h"

#undef MONERO_DEFAULT_LOG_CATEGORY
#define MONERO_DEFAULT_LOG_CATEGORY "crypto"

namespace crypto {

  using std::abort;
  using std::int32_t;
  using std::int64_t;
  using std::size_t;
  using std::uint32_t;
  using std::uint64_t;

  extern "C" {
#include "crypto-ops.h"
#include "random.h"
  }

  const crypto::public_key null_pkey = crypto::public_key{};
  const crypto::secret_key null_skey = crypto::secret_key{};

  static inline unsigned char *operator &(ec_point &point) {
    return &reinterpret_cast<unsigned char &>(point);
  }

  static inline const unsigned char *operator &(const ec_point &point) {
    return &reinterpret_cast<const unsigned char &>(point);
  }

  static inline unsigned char *operator &(ec_scalar &scalar) {
    return &reinterpret_cast<unsigned char &>(scalar);
  }

  static inline const unsigned char *operator &(const ec_scalar &scalar) {
    return &reinterpret_cast<const unsigned char &>(scalar);
  }

  void generate_random_bytes_thread_safe(size_t _N, uint8_t *bytes)
  {
    static boost::mutex random_lock;
    boost::lock_guard<boost::mutex> lock(random_lock);
    generate_random_bytes_not_thread_safe(_N, bytes);
  }

  static inline bool less32(const unsigned char *k0, const unsigned char *k1)
  {
    for (int n = 31; n >= 0; --n)
    {
      if (k0[n] < k1[n])
        return true;
      if (k0[n] > k1[n])
        return false;
    }
    return false;
  }

  void random32_unbiased(unsigned char *bytes)
  {
    // l = 2^252 + 27742317777372353535851937790883648493.
    // it fits 15 in 32 bytes
    static const unsigned char limit[32] = { 0xe3, 0x6a, 0x67, 0x72, 0x8b, 0xce, 0x13, 0x29, 0x8f, 0x30, 0x82, 0x8c, 0x0b, 0xa4, 0x10, 0x39, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0 };
    do
    {
      generate_random_bytes_thread_safe(32, bytes);
    } while (!sc_isnonzero(bytes) && !less32(bytes, limit)); // should be good about 15/16 of the time
    sc_reduce32(bytes);
  }
  /* generate a random 32-byte (256-bit) integer and copy it to res */
  static inline void random_scalar(pq_seed &res) {
  	random32_unbiased((unsigned char*)res.data);
  }

  void hash_to_scalar(const void *data, size_t length, ec_scalar &res) {
    cn_fast_hash(data, length, reinterpret_cast<hash &>(res));
    sc_reduce32(&res);
  }

  rand_seed crypto_ops::generate_keys(public_key &pub, secret_key &sec, const rand_seed& recovery_key, bool recover) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
    rand_seed rng;

    unsigned char pk[CRYPTO_PUBLICKEYBYTES];
    unsigned char sk[CRYPTO_SECRETKEYBYTES];

    if (recover)
    {
        rng = recovery_key;
    }
    else
    {
        random_scalar(rng);
    }
    // Dilithium Key generation with Monero 256 bit generated bytes
    auto rc = crypto_sign_dilithium_keypair(pk, sk, (unsigned char *)&rng);

    assert(rc == 0);

    std::memcpy(&sec, sk, CRYPTO_SECRETKEYBYTES);
    std::memcpy(&pub, pk, CRYPTO_PUBLICKEYBYTES);

    // Returns the random 32 byte seed
    return rng;
  }

  bool crypto_ops::check_key(const public_key &key) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    ge_p3 point;
    return ge_frombytes_vartime(&point, &key) == 0;
  }

  bool crypto_ops::secret_key_to_public_key(const secret_key &sec, public_key &pub) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
    ge_p3 point;
    if (sc_check(&unwrap(sec)) != 0) {
      return false;
    }
    ge_scalarmult_base(&point, &unwrap(sec));
    ge_p3_tobytes(&pub, &point);
    return true;
  }

  bool crypto_ops::generate_key_derivation(const public_key &key1, const secret_key &key2, key_derivation &derivation) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    // TODO: No Dilithium implementation yet, so derivation = PVk
    std::memcpy(&derivation, &key1, CRYPTO_PUBLICKEYBYTES);

    return true;
  }

  void crypto_ops::derivation_to_scalar(const key_derivation &derivation, size_t output_index, ec_scalar &res) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
    struct {
      key_derivation derivation;
      char output_index[(sizeof(size_t) * 8 + 6) / 7];
    } buf;
    char *end = buf.output_index;
    buf.derivation = derivation;
    tools::write_varint(end, output_index);
    assert(end <= buf.output_index + sizeof buf.output_index);
    hash_to_scalar(&buf, end - reinterpret_cast<char *>(&buf), res);
  }

  bool crypto_ops::derive_public_key(const key_derivation &derivation, size_t output_index,
    const public_key &base, public_key &derived_key) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
    // TODO: No Dilithium implementation yet, so derived_key = PSk
    std::memcpy(&derived_key, &base, CRYPTO_PUBLICKEYBYTES);

    return true;
  }

  void crypto_ops::derive_secret_key(const key_derivation &derivation, size_t output_index,
    const secret_key &base, secret_key &derived_key) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    // TODO: No Dilithium implementation yet, so derived_key = sSk
    std::memcpy(&derived_key, &base, CRYPTO_SECRETKEYBYTES);
  }

  bool crypto_ops::derive_subaddress_public_key(const public_key &out_key, const key_derivation &derivation, std::size_t output_index, public_key &derived_key) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    // TODO: No Dilithium implementation yet, so derived_key = Poutk
    std::memcpy(&derived_key, &out_key, CRYPTO_PUBLICKEYBYTES);

    return true;
  }

// Dilithium Signature - crypto_sign
  void crypto_ops::generate_signature(const hash &prefix_hash, const public_key &pub, const secret_key &sec, signature &sig) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    unsigned long long signatureLen = 0L;
    auto result = crypto_sign_dilithium((unsigned char*)&sig, &signatureLen, (unsigned char *)&prefix_hash, sizeof(prefix_hash), &sec);
    assert(result == 0);
  }
// Dilithium Signature - crypto_open
  bool crypto_ops::check_signature(const hash &prefix_hash, const public_key &pub, const signature &sig) {
    LOG_PRINT_L1("crypto_ops " << __func__);

    unsigned long long mLen = 0L;
    unsigned char m[HASH_SIZE + CRYPTO_BYTES];
    auto result = crypto_sign_dilithium_open(m, &mLen, (unsigned char *)&sig, sizeof(sig), &pub);
    LOG_PRINT_L1("crypto_ops signature: " << result);

    return result == 0 ? true : false;
  }

  void crypto_ops::generate_tx_proof(const hash &prefix_hash, const public_key &R, const public_key &A, const boost::optional<public_key> &B, const public_key &_D, const secret_key &r, signature &sig) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
  }

  bool crypto_ops::check_tx_proof(const hash &prefix_hash, const public_key &R, const public_key &A, const boost::optional<public_key> &B, const public_key &_D, const signature &sig) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
      return true;
  }

  void crypto_ops::generate_key_image(const public_key &pub, const secret_key &sec, key_image &image) {
    LOG_PRINT_L1("crypto_ops " <<__func__);
    std::memcpy(&image, &pub, CRYPTO_PUBLICKEYBYTES);
  }

PUSH_WARNINGS
DISABLE_VS_WARNINGS(4200)
  struct ec_point_pair {
    ec_point a, b;
  };
  struct rs_comm {
    hash h;
    struct ec_point_pair ab[];
  };
POP_WARNINGS

  void crypto_ops::generate_ring_signature(const hash &prefix_hash, const key_image &image,
    const public_key *const *pubs, size_t pubs_count,
    const secret_key &sec, size_t sec_index,
    signature *sig) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

  }

  bool crypto_ops::check_ring_signature(const hash &prefix_hash, const key_image &image,
    const public_key *const *pubs, size_t pubs_count,
    const signature *sig) {
    LOG_PRINT_L1("crypto_ops " <<__func__);

    return true;
  }
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//         Copyright 2012 - 2014 MetaScale SAS
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_ARITHMETIC_FUNCTIONS_GENERIC_IDIVFIX_HPP_INCLUDED
#define BOOST_SIMD_ARITHMETIC_FUNCTIONS_GENERIC_IDIVFIX_HPP_INCLUDED

#include <boost/simd/arithmetic/functions/idivfix.hpp>
#include <boost/simd/include/functions/simd/divides.hpp>
#include <boost/simd/include/functions/simd/divfix.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/dispatch/attributes.hpp>

namespace boost { namespace simd { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT          ( idivfix_, tag::cpu_, (A0)
                                    , (generic_< arithmetic_<A0> >)
                                      (generic_< arithmetic_<A0> >)
                                    )
  {
    typedef A0 result_type;
    BOOST_FORCEINLINE BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return divfix(a0, a1);
    }
  };

#ifdef BOOST_MSVC
  #pragma warning(push)
  #pragma warning(disable: 4723) // potential divide by 0
#endif

  BOOST_DISPATCH_IMPLEMENT          ( idivfix_, tag::cpu_, (A0)
                                    , (generic_< floating_<A0> > )
                                      (generic_< floating_<A0> > )
                                    )

  {
    typedef typename boost::dispatch::meta::as_integer <A0>::type result_type;
    BOOST_FORCEINLINE BOOST_SIMD_FUNCTOR_CALL_REPEAT(2)
    {
      return itrunc(a0/a1);
    }
  };
} } }

#ifdef BOOST_MSVC
  #pragma warning(pop)
#endif

#endif


#ifndef STAN_MATH_PRIM_SCAL_FUN_INC_BETA_HPP
#define STAN_MATH_PRIM_SCAL_FUN_INC_BETA_HPP

#include <boost/math/special_functions/beta.hpp>

namespace stan {
namespace math {

inline double inc_beta(double a, double b, double x) {
  using boost::math::ibeta;
  return ibeta(a, b, x);
}

}  // namespace math
}  // namespace stan
#endif


/*
 * Copyright 2020 Ria Sonecha, Massachusetts Institute of Technology in Cambridge, MA, USA
 * Copyright 2020 Giuseppe Silano, University of Sannio in Benevento, Italy
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <thread>
#include <chrono>

#include "rotors_gazebo/Quaternion.h"
#include "rotors_gazebo/transform_datatypes.h"
#include "rotors_gazebo/parameters_ros.h"
#include <nav_msgs/Odometry.h>
#include "rotors_gazebo/Matrix3x3.h"
#include <ros/console.h>
#include <time.h>

#include <Eigen/Core>
#include <mav_msgs/conversions.h>
#include <mav_msgs/default_topics.h>
#include <ros/ros.h>
#include <std_srvs/Empty.h>
#include <trajectory_msgs/MultiDOFJointTrajectory.h>
#include <math.h>
#include "rotors_gazebo/spline_trajectory_generator.h"

#define DELTA_1   0.5
#define ALPHA_1   2
#define ALPHA_2   20
#define ALPHA_3   8
#define ALPHA_4   12
#define ALPHA_5   3
#define BETA_1    30
#define BETA_2    14
#define BETA_3    16
#define BETA_4    2
#define GAMMA_1   12
#define GAMMA_2   6
#define SAMPLING_TIME  10e-3       /* SAMPLING CONTROLLER TIME [s] - 100Hz */
#define START_SIMULATION_TIME 3   /* TIME GAZEBO NEEDS TO INITIALIZE THE ENVIRONMENT */

namespace rotors_gazebo {

  SplineTrajectoryGenerator::SplineTrajectoryGenerator(){

    // Parameters initializion
    a0_.setZero(); a1_.setZero(); a2_.setZero(); a3_.setZero(); a4_.setZero(); a5_.setZero();
    b0_.setZero(); b1_.setZero(); b2_.setZero(); b3_.setZero(); b4_.setZero();
    c0_.setZero(); c1_.setZero(); c2_.setZero(); c3_.setZero();

    g0_.setZero(); g1_.setZero(); g2_.setZero(); g3_.setZero(); g4_.setZero(); g5_.setZero();
    h0_.setZero(); h1_.setZero(); h2_.setZero(); h3_.setZero(); h4_.setZero();
    i0_.setZero(); i1_.setZero(); i2_.setZero(); i3_.setZero();

  }

  SplineTrajectoryGenerator::~SplineTrajectoryGenerator(){}

  void SplineTrajectoryGenerator::TrajectoryCallback(mav_msgs::EigenDroneState* odometry, double* time_final, double* time_init) {
    assert(odometry);
    assert(time_final);
    assert(time_init);

    double time_spline;
    time_spline = *time_final - *time_init;

    if (enable_parameter_computation_){
      ComputeSplineParameters(&time_final_);
      enable_parameter_computation_ = false;

      ROS_DEBUG("Publishing time_final_: %f", time_final_);
    }

    ROS_DEBUG("Publishing time_spline: %f", time_spline);

    // Desired Position
    double x_component, y_component, z_component;
    x_component = a5_.x() * pow(time_spline, 5) + a4_.x() * pow(time_spline, 4) + a3_.x() * pow(time_spline, 3) + a2_.x() * pow(time_spline, 2)
                  + a1_.x() * time_spline + a0_.x();
    y_component = a5_.y() * pow(time_spline, 5) + a4_.y() * pow(time_spline, 4) + a3_.y() * pow(time_spline, 3) + a2_.y() * pow(time_spline, 2)
                  + a1_.y() * time_spline + a0_.y();
    z_component = a5_.z() * pow(time_spline, 5) + a4_.z() * pow(time_spline, 4) + a3_.z() * pow(time_spline, 3) + a2_.z() * pow(time_spline, 2)
                  + a1_.z() * time_spline + a0_.z();

    ROS_DEBUG("Publishing position spline parameters along x-axis: [%f, %f, %f, %f, %f, %f].", a0_.x(), a1_.x(), a2_.x(), a3_.x(), a4_.x(), a5_.x());
    ROS_DEBUG("Publishing position spline parameters along y-axis: [%f, %f, %f, %f, %f, %f].", a0_.y(), a1_.y(), a2_.y(), a3_.y(), a4_.y(), a5_.y());
    ROS_DEBUG("Publishing position spline parameters along z-axis: [%f, %f, %f, %f, %f, %f].", a0_.z(), a1_.z(), a2_.z(), a3_.z(), a4_.z(), a5_.z());

    odometry->position_W = Eigen::Vector3f(x_component, y_component, z_component);

    ROS_DEBUG("Publishing position waypoint: [%f, %f, %f].", x_component, y_component, z_component);

    // Desired Linear Velocity
    x_component = b4_.x() * pow(time_spline, 4) + b3_.x() * pow(time_spline, 3) + b2_.x() * pow(time_spline, 2) + b1_.x() * time_spline
                  + b0_.x();
    y_component = b4_.y() * pow(time_spline, 4) + b3_.y() * pow(time_spline, 3) + b2_.y() * pow(time_spline, 2) + b1_.y() * time_spline
                  + b0_.y();
    z_component = b4_.z() * pow(time_spline, 4) + b3_.z() * pow(time_spline, 3) + b2_.z() * pow(time_spline, 2) + b1_.z() * time_spline
                  + b0_.z();

    ROS_DEBUG("Publishing velocity spline parameters along x-axis: [%f, %f, %f, %f, %f].", b0_.x(), b1_.x(), b2_.x(), b3_.x(), b4_.x());
    ROS_DEBUG("Publishing velocity spline parameters along y-axis: [%f, %f, %f, %f, %f].", b0_.y(), b1_.y(), b2_.y(), b3_.y(), b4_.y());
    ROS_DEBUG("Publishing velocity spline parameters along z-axis: [%f, %f, %f, %f, %f].", b0_.z(), b1_.z(), b2_.z(), b3_.z(), b4_.z());

    odometry->velocity = Eigen::Vector3f(x_component, y_component, z_component);

    ROS_DEBUG("Publishing velocity waypoint: [%f, %f, %f].", x_component, y_component, z_component);

    // Desired Acceleration
    x_component = c3_.x() * pow(time_spline, 3) + c2_.x() * pow(time_spline, 2) + c1_.x() * time_spline + c0_.x();
    y_component = c3_.y() * pow(time_spline, 3) + c2_.y() * pow(time_spline, 2) + c1_.y() * time_spline + c0_.y();
    z_component = c3_.z() * pow(time_spline, 3) + c2_.z() * pow(time_spline, 2) + c1_.z() * time_spline + c0_.z();

    ROS_DEBUG("Publishing acceleration spline parameters along y-axis: [%f, %f, %f, %f].", c0_.y(), c1_.y(), c2_.y(), c3_.y());
    ROS_DEBUG("Publishing acceleration spline parameters along x-axis: [%f, %f, %f, %f].", c0_.x(), c1_.x(), c2_.x(), c3_.x());
    ROS_DEBUG("Publishing acceleration spline parameters along z-axis: [%f, %f, %f, %f].", c0_.z(), c1_.z(), c2_.z(), c3_.z());

    odometry->acceleration = Eigen::Vector3f(x_component, y_component, z_component);

    ROS_DEBUG("Publishing acceleration waypoint: [%f, %f, %f].", x_component, y_component, z_component);

    // Desired Attitute
    rollDesRad_ = g5_.x() * pow(time_spline, 5) + g4_.x() * pow(time_spline, 4) + g3_.x() * pow(time_spline, 3) + g2_.x() * pow(time_spline, 2)
                + g1_.x() * time_spline + g0_.x();
    pitchDesRad_ = g5_.y() * pow(time_spline, 5) + g4_.y() * pow(time_spline, 4) + g3_.y() * pow(time_spline, 3) + g2_.y() * pow(time_spline, 2)
                + g1_.y() * time_spline + g0_.y();
    yawDesRad_ = g5_.z() * pow(time_spline, 5) + g4_.z() * pow(time_spline, 4) + g3_.z() * pow(time_spline, 3) + g2_.z() * pow(time_spline, 2)
                + g1_.z() * time_spline + g0_.z();

    ROS_DEBUG("Publishing orientation spline parameters along x-axis: [%f, %f, %f, %f, %f, %f].", g0_.x(), g1_.x(), g2_.x(), g3_.x(), g4_.x(), g5_.x());
    ROS_DEBUG("Publishing orientation spline parameters along y-axis: [%f, %f, %f, %f, %f, %f].", g0_.y(), g1_.y(), g2_.y(), g3_.y(), g4_.y(), g5_.y());
    ROS_DEBUG("Publishing orientation spline parameters along z-axis: [%f, %f, %f, %f, %f, %f].", g0_.z(), g1_.z(), g2_.z(), g3_.z(), g4_.z(), g5_.z());

    // Converts radians in quaternions
    double x, y, z, w;
    Euler2QuaternionCommandTrajectory(&x, &y, &z, &w);
    odometry->orientation_W_B = Eigen::Quaterniond(w, x, y, z);

    ROS_DEBUG("Publishing attitude waypoint: [%f, %f, %f].", rollDesRad_, pitchDesRad_, yawDesRad_);

    // Desired Angular Velocity
    double roll_component, pitch_component, yaw_component;
    roll_component = h4_.x() * pow(time_spline, 4) + h3_.x() * pow(time_spline, 3) + h2_.x() * pow(time_spline, 2) + h1_.x() * time_spline + h0_.x();
    pitch_component = h4_.y() * pow(time_spline, 4) + h3_.y() * pow(time_spline, 3) + h2_.y() * pow(time_spline, 2) + h1_.y() * time_spline + h0_.y();
    yaw_component = h4_.z() * pow(time_spline, 4) + h3_.z() * pow(time_spline, 3) + h2_.z() * pow(time_spline, 2) + h1_.z() * time_spline + h0_.z();

    ROS_DEBUG("Publishing angular velocity waypoint: [%f, %f, %f].", roll_component, pitch_component, yaw_component);

    // The angular rates are rotated from the inertial to the body frame
    // https://home.aero.polimi.it/trainelli/downloads/Bottasso_ThreeDimensionalRotations.pdf, eqs. 1.206 and 1.208
    double roll_component_B, pitch_component_B, yaw_component_B;
    roll_component_B = roll_component - sin(pitchDesRad_) * yaw_component;
    pitch_component_B = cos(rollDesRad_) * pitch_component + sin(rollDesRad_) * cos(pitchDesRad_) * yaw_component;
    yaw_component_B = -sin(rollDesRad_) * pitch_component + cos(rollDesRad_) * cos(pitchDesRad_) * yaw_component;

    odometry->angular_velocity_B = Eigen::Vector3f(roll_component_B, pitch_component_B, yaw_component_B);

    ROS_DEBUG("Publishing angular velocity waypoint in the body frame: [%f, %f, %f].", roll_component_B, pitch_component_B, yaw_component_B);

  }

  void SplineTrajectoryGenerator::ComputeSplineParameters(double* time_spline){
      assert(time_spline);

      /*        POSITION       */
      // Parameters used for the position and its derivatives. The coefficient
      // refers to the x, y and z-avis
      a0_ = position_initial_;

      ROS_DEBUG("Publishing position initial: [%f, %f, %f].", position_initial_[0], position_initial_[1], position_initial_[2]);
      ROS_DEBUG("Content of the a0 coefficient: [%f, %f, %f].", a0_[0], a0_[1], a0_[2]);

      a1_ = velocity_initial_;

      ROS_DEBUG("Publishing velocity initial: [%f, %f, %f].", velocity_initial_[0], velocity_initial_[1], velocity_initial_[2]);
      ROS_DEBUG("Content of the a1 coefficient: [%f, %f, %f].", a1_[0], a1_[1], a1_[2]);

      a2_.x() = acceleration_initial_.x() * DELTA_1;
      a2_.y() = acceleration_initial_.y() * DELTA_1;
      a2_.z() = acceleration_initial_.z() * DELTA_1;

      ROS_DEBUG("Publishing acceleration initial: [%f, %f, %f].", acceleration_initial_[0], acceleration_initial_[1], acceleration_initial_[2]);
      ROS_DEBUG("Content of the a2 coefficient: [%f, %f, %f].", a2_[0], a2_[1], a2_[2]);

      a3_.x() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (position_final_.x() - position_initial_.x()) -
                (ALPHA_3 * velocity_final_.x() + ALPHA_4 * velocity_initial_.x()) * *time_spline -
                (ALPHA_5 * acceleration_final_.x() - acceleration_initial_.x()) * pow(*time_spline, 2));

      a3_.y() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (position_final_.y() - position_initial_.y()) -
                (ALPHA_3 * velocity_final_.y() + ALPHA_4 * velocity_initial_.y()) * *time_spline -
                (ALPHA_5 * acceleration_final_.y() - acceleration_initial_.y()) * pow(*time_spline, 2));

      a3_.z() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (position_final_.z() - position_initial_.z()) -
                (ALPHA_3 * velocity_final_.z() + ALPHA_4 * velocity_initial_.z()) * *time_spline -
                (ALPHA_5 * acceleration_final_.z() - acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the a3 coefficient: [%f, %f, %f].", a3_[0], a3_[1], a3_[2]);

      a4_.x() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (position_initial_.x() - position_final_.x()) +
                (BETA_2 * velocity_final_.x() + BETA_3 * velocity_initial_.x()) * *time_spline +
                (BETA_3 * acceleration_final_.x() - BETA_4 * acceleration_initial_.x()) * pow(*time_spline, 2));

      a4_.y() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (position_initial_.y() - position_final_.y()) +
                (BETA_2 * velocity_final_.y() + BETA_3 * velocity_initial_.y()) * *time_spline +
                (BETA_3 * acceleration_final_.y() - BETA_4 * acceleration_initial_.y()) * pow(*time_spline, 2));

      a4_.z() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (position_initial_.z() - position_final_.z()) +
                (BETA_2 * velocity_final_.z() + BETA_3 * velocity_initial_.z()) * *time_spline +
                (BETA_3 * acceleration_final_.z() - BETA_4 * acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the a4 coefficient: [%f, %f, %f].", a4_[0], a4_[1], a4_[2]);

      a5_.x() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (position_final_.x() - position_initial_.x()) -
                GAMMA_2 * (velocity_final_.x() + velocity_initial_.x()) * *time_spline -
                (acceleration_final_.x() - acceleration_initial_.x()) * pow(*time_spline, 2));

      a5_.y() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (position_final_.y() - position_initial_.y()) -
                GAMMA_2 * (velocity_final_.y() + velocity_initial_.y()) * *time_spline -
                (acceleration_final_.y() - acceleration_initial_.y()) * pow(*time_spline, 2));

      a5_.z() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (position_final_.z() - position_initial_.z()) -
                GAMMA_2 * (velocity_final_.z() + velocity_initial_.z()) * *time_spline -
                (acceleration_final_.z() - acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the a5 coefficient: [%f, %f, %f].", a5_[0], a5_[1], a5_[2]);

      // The coefficients are used for having the velocity reference trajectory. The polynomial is computed
      // derivating the a one. In other words,
      // a5 * x^5 + a4 * x^4 + a3 * x^3 + a2 * x^2 + a1 * x + a0
      // 5 * a5 * x^4 + 4 * a4 * x^3 + 3 * a3 * x^2 + 2 * a2 * x + a1
      // b4 = 5 * a5 -- b3 = 4 * a4 -- b2 = 3 * a3 - b1 = 2 * a2 - b0 = a1
      // and so on
      b4_ = 5 * a5_;
      b3_ = 4 * a4_;
      b2_ = 3 * a3_;
      b1_ = 2 * a2_;
      b0_ = a1_;

      ROS_DEBUG("Content of the b0 coefficient: [%f, %f, %f].", b0_[0], b0_[1], b0_[2]);
      ROS_DEBUG("Content of the b1 coefficient: [%f, %f, %f].", b1_[0], b1_[1], b1_[2]);
      ROS_DEBUG("Content of the b2 coefficient: [%f, %f, %f].", b2_[0], b2_[1], b2_[2]);
      ROS_DEBUG("Content of the b3 coefficient: [%f, %f, %f].", b3_[0], b3_[1], b3_[2]);
      ROS_DEBUG("Content of the b4 coefficient: [%f, %f, %f].", b4_[0], b4_[1], b4_[2]);

      // The coefficients are used for having the acceleration reference trajectory. The polynomial is computed
      // derivating the b one. In other words,
      // 4 * b4 * x^4 + 3 * b3 * x^3 + 2 * b2 * x^2 + b1 * x
      // c3 = 4 * b4 -- c2 = 3 * b3 -- c1 = 2 * b2 - c0 = b1
      c3_ = 4 * b4_;
      c2_ = 3 * b3_;
      c1_ = 2 * b2_;
      c0_ = b1_;

      ROS_DEBUG("Content of the c0 coefficient: [%f, %f, %f].", c0_[0], c0_[1], c0_[2]);
      ROS_DEBUG("Content of the c1 coefficient: [%f, %f, %f].", c1_[0], c1_[1], c1_[2]);
      ROS_DEBUG("Content of the c2 coefficient: [%f, %f, %f].", c2_[0], c2_[1], c2_[2]);
      ROS_DEBUG("Content of the c3 coefficient: [%f, %f, %f].", c3_[0], c3_[1], c3_[2]);

      /*        ORIENTATION      */
      // Parameters used for the orientation and its derivatives. The coefficient
      // refers to the ROLL (X), PITCH (Y) and YAW (Z)
      g0_ = orientation_initial_;

      ROS_DEBUG("Publishing orientation initial: [%f, %f, %f].", orientation_initial_[0], orientation_initial_[1], orientation_initial_[2]);
      ROS_DEBUG("Content of the g0 coefficient: [%f, %f, %f].", g0_[0], g0_[1], g0_[2]);

      g1_ = angular_velocity_initial_;

      ROS_DEBUG("Publishing angular rate initial: [%f, %f, %f].", angular_velocity_initial_[0], angular_velocity_initial_[1], angular_velocity_initial_[2]);
      ROS_DEBUG("Content of the g1 coefficient: [%f, %f, %f].", g1_[0], g1_[1], g1_[2]);

      g2_.x() = angular_acceleration_initial_.x() * DELTA_1;
      g2_.y() = angular_acceleration_initial_.y() * DELTA_1;
      g2_.z() = angular_acceleration_initial_.z() * DELTA_1;

      ROS_DEBUG("Publishing angular acceleration initial: [%f, %f, %f].", angular_acceleration_initial_[0], angular_acceleration_initial_[1],
                angular_acceleration_initial_[2]);
      ROS_DEBUG("Content of the g2 coefficient: [%f, %f, %f].", g2_[0], g2_[1], g2_[2]);

      g3_.x() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (orientation_final_.x() - orientation_initial_.x()) -
                (ALPHA_3 * angular_velocity_final_.x() + ALPHA_4 * angular_velocity_initial_.x()) * *time_spline -
                (ALPHA_5 * angular_acceleration_final_.x() - angular_acceleration_initial_.x()) * pow(*time_spline, 2));

      g3_.y() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (orientation_final_.y() - orientation_initial_.y()) -
                (ALPHA_3 * angular_velocity_final_.y() + ALPHA_4 * angular_velocity_initial_.y()) * *time_spline -
                (ALPHA_5 * angular_acceleration_final_.y() - angular_acceleration_initial_.y()) * pow(*time_spline, 2));

      g3_.z() = (1/(ALPHA_1*pow(*time_spline,3))) * (ALPHA_2 * (orientation_final_.z() - orientation_initial_.z()) -
                (ALPHA_3 * angular_velocity_final_.z() + ALPHA_4 * angular_velocity_initial_.z()) * *time_spline -
                (ALPHA_5 * angular_acceleration_final_.z() - angular_acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the g3 coefficient: [%f, %f, %f].", g3_[0], g3_[1], g3_[2]);

      g4_.x() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (orientation_initial_.x() - orientation_final_.x()) +
                (BETA_2 * angular_velocity_final_.x() + BETA_3 * angular_velocity_initial_.x()) * *time_spline +
                (BETA_3 * angular_acceleration_final_.x() - BETA_4 * angular_acceleration_initial_.x()) * pow(*time_spline, 2));

      g4_.y() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (orientation_initial_.y() - orientation_final_.y()) +
                (BETA_2 * angular_velocity_final_.y() + BETA_3 * angular_velocity_initial_.y()) * *time_spline +
                (BETA_3 * angular_acceleration_final_.y() - BETA_4 * angular_acceleration_initial_.y()) * pow(*time_spline, 2));

      g4_.z() = (1/(ALPHA_1*pow(*time_spline,4))) * (BETA_1 * (orientation_initial_.z() - orientation_final_.z()) +
                (BETA_2 * angular_velocity_final_.z() + BETA_3 * angular_velocity_initial_.z()) * *time_spline +
                (BETA_3 * angular_acceleration_final_.z() - BETA_4 * angular_acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the g4 coefficient: [%f, %f, %f].", g4_[0], g4_[1], g4_[2]);

      g5_.x() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (orientation_final_.x() - orientation_initial_.x()) -
                GAMMA_2 * (angular_velocity_final_.x() + angular_velocity_initial_.x()) * *time_spline -
                (angular_acceleration_final_.x() - angular_acceleration_initial_.x()) * pow(*time_spline, 2));

      g5_.y() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (orientation_final_.y() - orientation_initial_.y()) -
                GAMMA_2 * (angular_velocity_final_.y() + angular_velocity_initial_.y()) * *time_spline -
                (angular_acceleration_final_.y() - angular_acceleration_initial_.y()) * pow(*time_spline, 2));

      g5_.z() = (1/(ALPHA_1*pow(*time_spline,5))) * (GAMMA_1 * (orientation_final_.z() - orientation_initial_.z()) -
                GAMMA_2 * (angular_velocity_final_.z() + angular_velocity_initial_.z()) * *time_spline -
                (angular_acceleration_final_.z() - angular_acceleration_initial_.z()) * pow(*time_spline, 2));

      ROS_DEBUG("Content of the g5 coefficient: [%f, %f, %f].", g5_[0], g5_[1], g5_[2]);

      // The coefficients are used for having the angular velocity reference trajectory. The polynomial is computed
      // derivating the g one. In other words,
      // g5 * x^5 + g4 * x^4 + g3 * x^3 + g2 * x^2 + g1 * x + g0
      // 5 * g5 * x^4 + 4 * g4 * x^3 + 3 * g3 * x^2 + 2 * g2 * x + g1
      // h4 = 5 * g5 -- h3 = 4 * g4 -- h2 = 3 * g3 - h1 = 2 * g2 - h0 = g1
      // and so on
      h4_ = 5 * g5_;
      h3_ = 4 * g4_;
      h2_ = 3 * g3_;
      h1_ = 2 * g2_;
      h0_ = g1_;

      ROS_DEBUG("Content of the h0 coefficient: [%f, %f, %f].", h0_[0], h0_[1], h0_[2]);
      ROS_DEBUG("Content of the h1 coefficient: [%f, %f, %f].", h1_[0], h1_[1], h1_[2]);
      ROS_DEBUG("Content of the h2 coefficient: [%f, %f, %f].", h2_[0], h2_[1], h2_[2]);
      ROS_DEBUG("Content of the h3 coefficient: [%f, %f, %f].", h3_[0], h3_[1], h3_[2]);
      ROS_DEBUG("Content of the h4 coefficient: [%f, %f, %f].", h4_[0], h4_[1], h4_[2]);

      // The coefficients are used for having the angular acceleration reference trajectory. The polynomial is computed
      // derivating the b one. In other words,
      // 4 * h4 * x^4 + 3 * h3 * x^3 + 2 * h2 * x^2 + h1 * x
      // i3 = 4 * h4 -- i2 = 3 * h3 -- i1 = 2 * h2 - i0 = h1
      i3_ = 4 * h4_;
      i2_ = 3 * h3_;
      i1_ = 2 * h2_;
      i0_ = h1_;

      ROS_DEBUG("Content of the i0 coefficient: [%f, %f, %f].", i0_[0], i0_[1], i0_[2]);
      ROS_DEBUG("Content of the i1 coefficient: [%f, %f, %f].", i1_[0], i1_[1], i1_[2]);
      ROS_DEBUG("Content of the i2 coefficient: [%f, %f, %f].", i2_[0], i2_[1], i2_[2]);
      ROS_DEBUG("Content of the i3 coefficient: [%f, %f, %f].", i3_[0], i3_[1], i3_[2]);

  }

  void SplineTrajectoryGenerator::Euler2QuaternionCommandTrajectory(double* x, double* y, double* z, double* w) const {
      assert(x);
      assert(y);
      assert(z);
      assert(w);

      // Abbreviations for the various angular functions
      double cy = cos(yawDesRad_ * 0.5);
      double sy = sin(yawDesRad_ * 0.5);
      double cp = cos(pitchDesRad_ * 0.5);
      double sp = sin(pitchDesRad_ * 0.5);
      double cr = cos(rollDesRad_ * 0.5);
      double sr = sin(rollDesRad_ * 0.5);

      *w = cy * cp * cr + sy * sp * sr;
      *x = cy * cp * sr - sy * sp * cr;
      *y = sy * cp * sr + cy * sp * cr;
      *z = sy * cp * cr - cy * sp * sr;

      ROS_DEBUG("x Trajectory: %f, y Trajectory: %f, z Trajectory: %f, w Trajectory: %f", *x, *y, *z, *w);
    }

   void SplineTrajectoryGenerator::InitializeParams(){

     ros::NodeHandle pnh("~");

     // Parameters reading from rosparam.
     GetRosParameter(pnh, "time_final/time", time_final_, &time_final_);
     GetRosParameter(pnh, "position_initial/x", position_initial_.x(), &position_initial_.x());
     GetRosParameter(pnh, "position_initial/y", position_initial_.y(), &position_initial_.y());
     GetRosParameter(pnh, "position_initial/z", position_initial_.z(), &position_initial_.z());

     GetRosParameter(pnh, "position_final/x", position_final_.x(), &position_final_.x());
     GetRosParameter(pnh, "position_final/y", position_final_.y(), &position_final_.y());
     GetRosParameter(pnh, "position_final/z", position_final_.z(), &position_final_.z());

     GetRosParameter(pnh, "velocity_initial/x", velocity_initial_.x(), &velocity_initial_.x());
     GetRosParameter(pnh, "velocity_initial/y", velocity_initial_.y(), &velocity_initial_.y());
     GetRosParameter(pnh, "velocity_initial/z", velocity_initial_.z(), &velocity_initial_.z());

     GetRosParameter(pnh, "velocity_final/x", velocity_final_.x(), &velocity_final_.x());
     GetRosParameter(pnh, "velocity_final/y", velocity_final_.y(), &velocity_final_.y());
     GetRosParameter(pnh, "velocity_final/z", velocity_final_.z(), &velocity_final_.z());

     GetRosParameter(pnh, "acceleration_initial/x", acceleration_initial_.x(), &acceleration_initial_.x());
     GetRosParameter(pnh, "acceleration_initial/y", acceleration_initial_.y(), &acceleration_initial_.y());
     GetRosParameter(pnh, "acceleration_initial/z", acceleration_initial_.z(), &acceleration_initial_.z());

     GetRosParameter(pnh, "acceleration_final/x", acceleration_final_.x(), &acceleration_final_.x());
     GetRosParameter(pnh, "acceleration_final/y", acceleration_final_.y(), &acceleration_final_.y());
     GetRosParameter(pnh, "acceleration_final/z", acceleration_final_.z(), &acceleration_final_.z());

     GetRosParameter(pnh, "orientation_initial/roll", orientation_initial_.x(), &orientation_initial_.x());
     GetRosParameter(pnh, "orientation_initial/pitch", orientation_initial_.y(), &orientation_initial_.y());
     GetRosParameter(pnh, "orientation_initial/yaw", orientation_initial_.z(), &orientation_initial_.z());

     GetRosParameter(pnh, "orientation_final/roll", orientation_final_.x(), &orientation_final_.x());
     GetRosParameter(pnh, "orientation_final/pitch", orientation_final_.y(), &orientation_final_.y());
     GetRosParameter(pnh, "orientation_final/yaw", orientation_final_.z(), &orientation_final_.z());

     GetRosParameter(pnh, "angularRate_initial/roll", angular_velocity_initial_.x(), &angular_velocity_initial_.x());
     GetRosParameter(pnh, "angularRate_initial/pitch", angular_velocity_initial_.y(), &angular_velocity_initial_.y());
     GetRosParameter(pnh, "angularRate_initial/yaw", angular_velocity_initial_.z(), &angular_velocity_initial_.z());

     GetRosParameter(pnh, "angularRate_final/roll", angular_velocity_final_.x(), &angular_velocity_final_.x());
     GetRosParameter(pnh, "angularRate_final/pitch", angular_velocity_final_.y(), &angular_velocity_final_.y());
     GetRosParameter(pnh, "angularRate_final/yaw", angular_velocity_final_.z(), &angular_velocity_final_.z());

     GetRosParameter(pnh, "angularDoubleRate_initial/roll", angular_acceleration_initial_.x(), &angular_acceleration_initial_.x());
     GetRosParameter(pnh, "angularDoubleRate_initial/pitch", angular_acceleration_initial_.y(), &angular_acceleration_initial_.y());
     GetRosParameter(pnh, "angularDoubleRate_initial/yaw", angular_acceleration_initial_.z(), &angular_acceleration_initial_.z());

     GetRosParameter(pnh, "angularDoubleRate_final/roll", angular_acceleration_final_.x(), &angular_acceleration_final_.x());
     GetRosParameter(pnh, "angularDoubleRate_final/pitch", angular_acceleration_final_.y(), &angular_acceleration_final_.y());
     GetRosParameter(pnh, "angularDoubleRate_final/yaw", angular_acceleration_final_.z(), &angular_acceleration_final_.z());

   }

   template<typename T> inline void SplineTrajectoryGenerator::GetRosParameterHovering(const ros::NodeHandle& nh,
                                                            const std::string& key,
                                                            const T& default_value,
                                                            T* value) {

     ROS_ASSERT(value != nullptr);
     bool have_parameter = nh.getParam(key, *value);
     if (!have_parameter) {
       ROS_WARN_STREAM("[rosparam]: could not find parameter " << nh.getNamespace()
                       << "/" << key << ", setting to default: " << default_value);
       *value = default_value;
     }
   }

}


#ifndef NYX_STATE_HPP
#define NYX_STATE_HPP

#include <vector>
#include <armadillo>

#include "nyx_link.hpp"
#include "nyx_pose.hpp"

// Class to keep track of state of arm during solve
class State {
    public:
        // Constructor
        State(std::vector<Link> chain);
        State(std::vector<Link> chain, Pose currentPose, Pose destinationPose);
        State(const State &state);

        // Create children in tree, return vector of States
        std::vector<State> getNext();
        std::vector<State> getNext(int stepSize);

        // Operator overloading to compare state scores
        bool operator>(const State &rhs) const;
        bool operator==(const State &rhs);

        void print() const;
        double getScore() const;

        std::string hash() const;
        Pose forwardKinematics() const;
        bool withinTolerance(const State &rhs, int tolerance) const;

        void setChain(std::vector<Link> chain);

        std::vector<Link> getChain();
    protected:
        std::vector<Link> chain;

    private:
        State();

        Pose currentPose;
        Pose destinationPose;

        int stepCount;
};

State::State(std::vector<Link> chain) :
    State(chain, Pose(), Pose())
{
}

State::State(std::vector<Link> chain, Pose currentPose, Pose destinationPose) :
    chain (chain),
    currentPose (currentPose),
    destinationPose (destinationPose)
{
    stepCount = 0;
}
State::State(const State &state) {
    this->chain = state.chain;
    this->currentPose = state.currentPose;
    this->destinationPose = state.destinationPose;
    this->stepCount = state.stepCount + 1;
}

std::vector<State> State::getNext() {
    return getNext(1);
}
std::vector<State> State::getNext(int stepSize) {
    // Create all next states that can result from current configuration
    std::vector<State> out = std::vector<State>();

    for (unsigned int i=0; i<chain.size(); i++) {
        // Get link from chain, and increase and decrease angle by 1 degree
        Link decrement = chain[i];
        decrement.setCurrentAngle(decrement.getAngle() - stepSize);
        // Create states that will be put into vector
        State first  = State(*this);
        // Check if valid link
        if (decrement.isValid()) {
            // Update link angles in each state
            first.chain[i] = decrement;
            // Push states to vector
            out.push_back(first);
        }

        Link increment = chain[i];
        increment.setCurrentAngle(increment.getAngle() + stepSize);
        State second = State(*this);
        if (increment.isValid()) {
            second.chain[i] = increment;
            out.push_back(second);
        }
    }

    return out;
}

void State::print() const {
    std::cout << "State information\n";
    currentPose.print();
    currentPose.print();
    for (unsigned int i=0; i<chain.size(); i++) {
        //output += "  " + chain[i].print() + "\n";
        chain[i].print();
    }
    std::cout << "End pose: ";
    Pose temp = forwardKinematics();
    temp.print();
}
double State::getScore() const {
    // Constants to multiply by
    double pastCoefficient   = 0.3;
    double futureCoefficient = 0.7;

    // Get distance between end effector and destinationPose
    double dist = destinationPose - forwardKinematics();

    return pastCoefficient * stepCount + futureCoefficient * dist;
}

bool State::operator>(const State &rhs) const {
    return (*this).getScore() > rhs.getScore();
}
bool State::operator==(const State &rhs) {
    return (*this).getScore() == rhs.getScore();
}

// Hash of this state
std::string State::hash() const {
    std::string out = "";
    for (unsigned int i=0; i<chain.size(); i++) {
        Link temp = chain[i];
        out += temp.getAngle();
    }
    return out;
}

// Get position of end manipulator with @chain
Pose State::forwardKinematics() const {
    // Initialize lastValue to be the final vector V
    double x, y, z;
    double rho, theta, phi;
    double c_phi, c_theta,
           s_phi, s_theta;

    Link current = chain[chain.size()-1];
    // Get X/Y/Z of link
    x = current.getX();
    y = current.getY();
    z = current.getZ();
    // Get angles
    cartesianToSpherical(x,y,z, rho,theta,phi);
    c_phi = cos(phi);
    s_phi = sin(phi);
    c_theta = cos(theta);
    s_theta = sin(theta);

    /*
    arma::dmat lastValue { rho * c_phi * c_theta ,
                           rho * c_phi * s_theta ,
                           rho * s_phi           };
                           */
    //std::cout << "Creating first lastValue\n";
    arma::dmat lastValue (3,1);
    lastValue[0] = rho * c_phi * c_theta;
    lastValue[1] = rho * c_phi * s_theta;
    lastValue[2] = rho * s_phi;
    //std::cout << "Finish first lastValue\n";

    // Loop through the links backwards, and apply the coordinate transforms
    for (int i=chain.size()-2; i>=0; i--) {
        // Y
        // [ cos(phi), 0, sin(phi),
        //   0       , 1, 0       ,
        //  -sin(phi), 0, cos(phi) ]
        //
        //  Z
        //  [ cos(theta), -sin(theta), 0,
        //    sin(theta),  cos(theta), 0,
        //    0         , 0          , 1 ]

        //std::cout << "A + " <<i<< "\n";
        x = chain[i].getX();
        y = chain[i].getY();
        z = chain[i].getZ();
        //std::cout << "B\n";
        // Get angles
        cartesianToSpherical(x,y,z, rho,theta,phi);
        c_phi = cos(phi);
        s_phi = sin(phi);
        c_theta = cos(theta);
        s_theta = sin(theta);
        //std::cout << "C\n";

        // Coordinate transform matrix
        arma::dmat y { { c_phi, 0 , s_phi},
                       {  0 , 1 ,  0 },
                       {-s_phi, 0 ,  c_phi} };
        //std::cout << "D\n";
        arma::dmat z { { c_theta, -s_theta, 0 },
                       { s_theta,  c_theta, 0 },
                       {  0 ,   0 , 1 } };

        //std::cout << "E\n";
        arma::dmat v (3,1);
        v[0] = rho * c_phi * c_theta;
        v[1] = rho * c_phi * s_theta;
        v[2] = rho * s_phi;;
        //std::cout << "F\n";

        lastValue = v + y * z * lastValue;
        //std::cout << "G\n";
    }

    //std::cout << "H\n";
    Pose end = Pose(lastValue[0], lastValue[1], lastValue[2]);
    //std::cout << "I\n";
    return end;
    /*
    // Initial pose
    Pose end = Pose(0,0,0);

    // End effector orientation
    double yaw = 0.0;
    double pitch = 0.0;

    // Loop through links and adjust pose and orientation
    for (unsigned int i=0; i<chain.size(); i++) {
        Link temp = chain[i];
        // If yaw joint, apply yaw
        if (temp.getAxisOfRotation() == 1) {
            yaw += temp.getAngle();
        }
        // If pitch joint, apply pitch
        if (temp.getAxisOfRotation() == 2) {
            pitch += temp.getAngle();
        }

        // Adjust pose appropriately

        end.setX( end.getX() + temp.getLength()*cos(radians(pitch))*cos(radians(yaw)) );
        end.setY( end.getY() + temp.getLength()*cos(radians(pitch))*sin(radians(yaw)) );
        end.setZ( end.getZ() + temp.getLength()*sin(radians(pitch)) );
    }

    return end;
    */
}

// Check that state is within bounds to given state
bool State::withinTolerance(const State &rhs, int tolerance) const {
    for (unsigned int i=0; i<chain.size(); i++) {
        if ( abs(chain[i].getAngle() - rhs.chain[i].getAngle()) > tolerance ) {
            return false;
        }
    }
    return true;
}

// Set the chain, this may need to be removed after testing
void State::setChain(std::vector<Link> chain) {
    this->chain = chain;
}

// Get chain
std::vector<Link> State::getChain() {
    return chain;
}

#endif


//    Copyright 2017 Rainer Gemulla
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
#include <iostream>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <boost/random/uniform_real.hpp>

#include <util/evaluation.h>

#include <mpi2/mpi2.h>

// added by me ******************
#include "detail/mfdsgd-args.h"
#include "parse.h"
#include <mf/matrix/io/generateDistributedMatrix.h>

//*******************************
log4cxx::LoggerPtr logger(log4cxx::Logger::getLogger("main"));

using namespace mf;
using namespace rg;

typedef UpdateLock<UpdateNzsl> UpdateLockNzsl;
typedef UpdateLock<UpdateNzslL2> UpdateLockNzslL2;
typedef UpdateLock<UpdateNzslNzl2> UpdateLockNzslNzl2;
typedef UpdateLock<UpdateNzslTruncate> UpdateLockNzslTruncate;
typedef UpdateLock<UpdateNzslL2Truncate> UpdateLockNzslL2Truncate;
typedef UpdateLock<UpdateNzslNzl2Truncate> UpdateLockNzslNzl2Truncate;
MPI2_TYPE_TRAITS(UpdateLockNzsl);
MPI2_TYPE_TRAITS(UpdateLockNzslL2);
MPI2_TYPE_TRAITS(UpdateLockNzslNzl2);
MPI2_TYPE_TRAITS(UpdateLockNzslTruncate);
MPI2_TYPE_TRAITS(UpdateLockNzslL2Truncate);
MPI2_TYPE_TRAITS(UpdateLockNzslNzl2Truncate);

template<typename U,typename R,typename L, typename D>
//void runAsgd2(Args& args, U update, R regularize, L loss, D decay,
//		AsgdJob<U,R>& asgdJob, DistributedDenseMatrix& dw, DistributedDenseMatrixCM& dh, Trace& trace) {
void runAsgd2(Args& args, U update, R regularize, L loss, D decay,
			AsgdJob<U,R>& asgdJob, std::pair<DistributedDenseMatrix, DistributedDenseMatrixCM>& factorsPair,
			std::vector<DistributedSparseMatrix>& dataVector, Trace& trace) {

	mf_size_type blocks1 = args.worldSize;
	mf_size_type blocks2 = args.worldSize;
	Timer t;
	AsgdRunner asgdRunner(args.random);
	if (args.inputTestMatrixFile.length() == 0) {
		// run DSGD
		t.start();
		asgdRunner.run(asgdJob, loss, args.epochs, decay, trace, args.balanceType, args.balanceMethod);
		t.stop();
		LOG4CXX_INFO(logger, "Total time: " << t);
	} else {
//		DistributedSparseMatrix dvTest=loadMatrix<SparseMatrix>(args.inputTestMatrixFile,
//									"Vtest", true, 1, args.worldSize, blocks1, 1);
//		LOG4CXX_INFO(logger, "Test matrix: "
//			<< dvTest.size1() << " x " << dvTest.size2() << ", " << nnz(dvTest) << " nonzeros, "
//			<< dvTest.blocks1() << " x " << dvTest.blocks2() << " blocks");
//
//		AsgdFactorizationData<> testData(dvTest,dw,dh,args.tasksPerRank);
		AsgdFactorizationData<> testData(dataVector[1],factorsPair.first,factorsPair.second,args.tasksPerRank);
		LOG4CXX_INFO(logger, "Using NzslLoss for test data");
		NzslLoss testLoss;
		// run ASGD
		t.start();
		asgdRunner.run(asgdJob, loss, args.epochs, decay, trace, args.balanceType, args.balanceMethod, &testData, &testLoss);
		t.stop();
		LOG4CXX_INFO(logger, "Total time: " << t);
	}
}

// run ASGD
template<typename U,typename R,typename L>
void runAsgd(Args& args, U update, R regularize, L loss) {

	mf_size_type blocks1 = args.worldSize;
	mf_size_type blocks2 = args.worldSize;

	std::vector<DistributedSparseMatrix> dataVector;
	if (args.inputTestMatrixFile.length() == 0) {
		dataVector=getDataMatrices<SparseMatrix>(args.inputMatrixFile,"V",true,args.tasksPerRank, args.worldSize, blocks1, 1,true,false);
	}else{
		dataVector=getDataMatrices<SparseMatrix>(args.inputMatrixFile,"V",true,args.tasksPerRank, args.worldSize, blocks1, 1,true,false, &args.inputTestMatrixFile);
	}

	std::pair<DistributedDenseMatrix, DistributedDenseMatrixCM> factorsPair= getFactors(args.inputRowFacFile,
			args.inputColFacFile,  args.tasksPerRank, args.worldSize,blocks1,blocks2,true);

//	// distribute the input matrices
//	DistributedSparseMatrix dv=loadMatrix<SparseMatrix>(args.inputMatrixFile,
//							"V", true, 1, args.worldSize,blocks1,1);
//	LOG4CXX_INFO(logger, "Data matrix: "
//		<< dv.size1() << " x " << dv.size2() << ", " << nnz(dv) << " nonzeros, "
//		<< dv.blocks1() << " x " << dv.blocks2() << " blocks");
//
//	DistributedDenseMatrix dw=loadMatrix<DenseMatrix>(args.inputRowFacFile,
//											"W", true, 1, args.worldSize,blocks1,1);
//	LOG4CXX_INFO(logger, "Row factor matrix: "
//		<< dw.size1() << " x " << dw.size2() << ", " << dw.blocks1() << " x " << dw.blocks2() << " blocks");
//	DistributedDenseMatrixCM dh=loadMatrix<DenseMatrixCM>(args.inputColFacFile,
//											"H", false, 1, args.worldSize,1,blocks2);
//	LOG4CXX_INFO(logger, "Column factor matrix: "
//		<< dh.size1() << " x " << dh.size2() << ", " << dh.blocks1() << " x " << dh.blocks2() << " blocks");
//
//	AsgdJob<U,R> asgdJob(dv, dw, dh, update, regularize, args.sgdOrder, args.stratumOrder, args.tasksPerRank, args.averageDeltas);

	AsgdJob<U,R> asgdJob(dataVector[0], factorsPair.first, factorsPair.second,
			update, regularize, args.sgdOrder, args.stratumOrder, args.tasksPerRank, args.averageDeltas);

	Trace trace;
	// add trace fields
	args.createTraceFields(trace);

	if (args.decayName.compare("BoldDriver") == 0){

		if (std::isnan(args.epsDecrease) && (std::isnan(args.epsIncrease))) {
			BoldDriver decay(args.epsilon);
//			runAsgd2<>(args, update, regularize, loss, decay, asgdJob, dw, dh, trace);
			runAsgd2<>(args, update, regularize, loss, decay, asgdJob, factorsPair, dataVector, trace);
		} else {
			BoldDriver decay(args.epsilon, args.epsDecrease, args.epsIncrease);
//			runAsgd2<>(args, update, regularize, loss, decay, asgdJob, dw, dh, trace);
			runAsgd2<>(args, update, regularize, loss, decay, asgdJob, factorsPair, dataVector, trace);
		}

	} else if (args.decayName.compare("Const") == 0) {

		DecayConstant decay(args.epsilon);
//		runAsgd2<>(args, update, regularize, loss, decay, asgdJob, dw, dh, trace);
		runAsgd2<>(args, update, regularize, loss, decay, asgdJob, factorsPair, dataVector, trace);
	} else {
		RG_THROW(rg::NotImplementedException, "decay");
	}

	// write trace to an R file
	if (args.traceFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing trace to " << args.traceFile);
		trace.toRfile(args.traceFile, args.traceVar);
	}

	// write computed factors to file
	if (args.outputRowFacFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing row factors to " << args.outputRowFacFile);
		DenseMatrix w0;
//		unblock(dw, w0);
		unblock(factorsPair.first, w0);
		writeMatrix(args.outputRowFacFile, w0);
	}
	if (args.outputColFacFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing column factors to " << args.outputColFacFile);
		DenseMatrixCM h0;
//		unblock(dh, h0);
		unblock(factorsPair.second, h0);
		writeMatrix(args.outputColFacFile, h0);
	}
}

bool runArgs(Args& args) {
	using namespace mf;
	using namespace mpi2;
	if (!(args.abs)) {
		if (!(!std::isnan(args.truncateArgs[0]) && !std::isnan(args.truncateArgs[1]))) {
			if ((args.updateName.compare("Nzsl") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl") == 0)) {
				if (args.updateArgs.size()<0 || args.updateArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzsl update(UNINITIALIZED);
				if (args.updateArgs.size()==0) {
					update = UpdateNzsl();
				}
				UpdateNzsl updateAbs = update;
				UpdateNzsl updateTruncate = updateAbs;
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeNone regularizeTruncate = regularizeAbs;
				if (args.lossArgs.size()<0 || args.lossArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				NzslLoss loss(UNINITIALIZED);
				if (args.lossArgs.size()==0) {
					loss = NzslLoss();
				}
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
			if ((args.updateName.compare("Nzsl_L2") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl_L2") == 0)) {
				if (args.updateArgs.size()<1 || args.updateArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzslL2 update(UNINITIALIZED);
				if (args.updateArgs.size()==1) {
					update = UpdateNzslL2(args.updateArgs[0]);
				}
				UpdateNzslL2 updateAbs = update;
				UpdateNzslL2 updateTruncate = updateAbs;
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeNone regularizeTruncate = regularizeAbs;
				NzslLoss loss1;
				if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				L2Loss loss2(UNINITIALIZED);
				if (args.lossArgs.size()==1) {
					loss2 = L2Loss(args.lossArgs[0]);
				}
				SumLoss<NzslLoss, L2Loss> loss = SumLoss<NzslLoss, L2Loss>(loss1, loss2);
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
			if ((args.updateName.compare("Nzsl_Nzl2") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl_Nzl2") == 0)) {
				if (args.updateArgs.size()<1 || args.updateArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzslNzl2 update(UNINITIALIZED);
				if (args.updateArgs.size()==1) {
					update = UpdateNzslNzl2(args.updateArgs[0]);
				}
				UpdateNzslNzl2 updateAbs = update;
				UpdateNzslNzl2 updateTruncate = updateAbs;
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeNone regularizeTruncate = regularizeAbs;
				NzslLoss loss1;
				if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				Nzl2Loss loss2(UNINITIALIZED);
				if (args.lossArgs.size()==1) {
					loss2 = Nzl2Loss(args.lossArgs[0]);
				}
				SumLoss<NzslLoss, Nzl2Loss> loss = SumLoss<NzslLoss, Nzl2Loss>(loss1, loss2);
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
		}
		if ((!std::isnan(args.truncateArgs[0]) && !std::isnan(args.truncateArgs[1]))) {
			if ((args.updateName.compare("Nzsl") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl") == 0)) {
				if (args.updateArgs.size()<0 || args.updateArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzsl update(UNINITIALIZED);
				if (args.updateArgs.size()==0) {
					update = UpdateNzsl();
				}
				UpdateNzsl updateAbs = update;
				UpdateTruncate<UpdateNzsl > updateTruncate(updateAbs, args.truncateArgs[0], args.truncateArgs[1]);
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeTruncate<RegularizeNone > regularizeTruncate(regularizeAbs, args.truncateArgs[0], args.truncateArgs[1]);
				if (args.lossArgs.size()<0 || args.lossArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				NzslLoss loss(UNINITIALIZED);
				if (args.lossArgs.size()==0) {
					loss = NzslLoss();
				}
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
			if ((args.updateName.compare("Nzsl_L2") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl_L2") == 0)) {
				if (args.updateArgs.size()<1 || args.updateArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzslL2 update(UNINITIALIZED);
				if (args.updateArgs.size()==1) {
					update = UpdateNzslL2(args.updateArgs[0]);
				}
				UpdateNzslL2 updateAbs = update;
				UpdateTruncate<UpdateNzslL2 > updateTruncate(updateAbs, args.truncateArgs[0], args.truncateArgs[1]);
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeTruncate<RegularizeNone > regularizeTruncate(regularizeAbs, args.truncateArgs[0], args.truncateArgs[1]);
				NzslLoss loss1;
				if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				L2Loss loss2(UNINITIALIZED);
				if (args.lossArgs.size()==1) {
					loss2 = L2Loss(args.lossArgs[0]);
				}
				SumLoss<NzslLoss, L2Loss> loss = SumLoss<NzslLoss, L2Loss>(loss1, loss2);
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
			if ((args.updateName.compare("Nzsl_Nzl2") == 0) && (args.regularizeName.compare("None") == 0) && (args.lossName.compare("Nzsl_Nzl2") == 0)) {
				if (args.updateArgs.size()<1 || args.updateArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.updateString << std::endl;					return false;
				}
				UpdateNzslNzl2 update(UNINITIALIZED);
				if (args.updateArgs.size()==1) {
					update = UpdateNzslNzl2(args.updateArgs[0]);
				}
				UpdateNzslNzl2 updateAbs = update;
				UpdateTruncate<UpdateNzslNzl2 > updateTruncate(updateAbs, args.truncateArgs[0], args.truncateArgs[1]);
				if (args.regularizeArgs.size()<0 || args.regularizeArgs.size()>0) {
					std::cout << "Invalid number of arguments in " << args.regularizeString << std::endl;					return false;
				}
				RegularizeNone regularize(UNINITIALIZED);
				if (args.regularizeArgs.size()==0) {
					regularize = RegularizeNone();
				}
				RegularizeNone regularizeAbs = regularize;
				RegularizeTruncate<RegularizeNone > regularizeTruncate(regularizeAbs, args.truncateArgs[0], args.truncateArgs[1]);
				NzslLoss loss1;
				if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
					std::cout << "Invalid number of arguments in " << args.lossString << std::endl;					return false;
				}
				Nzl2Loss loss2(UNINITIALIZED);
				if (args.lossArgs.size()==1) {
					loss2 = Nzl2Loss(args.lossArgs[0]);
				}
				SumLoss<NzslLoss, Nzl2Loss> loss = SumLoss<NzslLoss, Nzl2Loss>(loss1, loss2);
				runAsgd(args, updateTruncate, regularizeTruncate, loss);
				return true;
			}
		}
	}
	std::cerr << "Invalid combination of update, regularize, and loss arguments" << std::endl;
	std::cout << "Valid combinations are:" << std::endl;
	std::cout << "	Nzsl / None / Nzsl" << std::endl;
	std::cout << "	Nzsl / L2 / Nzsl_L2" << std::endl;
	std::cout << "	Nzsl_Nzl2 / None / Nzsl_Nzl2" << std::endl;
	return false;
}

int main(int argc, char *argv[]) {
	using namespace std;
	using namespace mf;
	using namespace mpi2;
	using namespace rg;
	using namespace boost::program_options;

	// initialize mf library and mpi2
	boost::mpi::communicator& world = mfInit(argc, argv);

	mfStart();
	bool result = true;

	if (world.rank() == 0) {
		boost::this_thread::sleep(boost::posix_time::milliseconds(100)); // so that arguments are logged nicely
		Args args;

		// parse command line
		options_description desc("Options");
		desc.add_options()
			("help", "produce help message")
			("input-file", value<string>(&args.inputMatrixFile), "filename of data matrix")
			("input-test-file", value<string>(&args.inputTestMatrixFile), "filename of test matrix")
			("input-row-file", value<string>(&args.inputRowFacFile), "filename of initial row factors")
			("input-col-file", value<string>(&args.inputColFacFile), "filename of initial column factors")
			("output-row-file", value<string>(&args.outputRowFacFile), "filename of final row factors")
			("output-col-file", value<string>(&args.outputColFacFile), "filename of final column factors")
			("trace", value<string>(&args.traceFile), "filename of trace [trace.R]")
			("trace-var", value<string>(&args.traceVar), "variable name for trace [traceVar]")
			("epochs", value<mf_size_type>(&args.epochs), "number of epochs to run [10]")
			("tasks-per-rank", value<int>(&args.tasksPerRank), "number of concurrent tasks per rank [1]")
			("sgd-order", value<string>(&args.sgdOrderString), "order of SGD steps [WOR] (e.g., \"SEQ\", \"WR\", \"WOR\")")
			("seed", value<unsigned>(&args.seed), "seed for random number generator (system time if not set)")
			("rank", value<mf_size_type>(&args.rank), "rank of factorization")
			("update", value<string>(&args.updateString), "SGD update function (e.g., \"Sl\", \"Nzsl\", \"GklData\")")
			("regularize", value<string>(&args.regularizeString), "SGD regularization function (e.g., \"None\", \"L2(0.05)\", \"Nzl2(0.05)\")")
			("loss", value<string>(&args.lossString), "loss function (e.g., \"Nzsl\", \"Nzsl_L2(0.5)\"))")
			("abs", "if present, absolute values are taken after every SGD step")
			("truncate", value<string>(&args.truncateString), "if present, truncatation is enabled (e.g., --truncate \"(-1000, 1000)\"")
			("decay", value<string>(&args.decayString), "decay function (constant, bold driver, or auto)")
			("balance", value<string>(&args.balanceString), "Type of balancing (None, L2, Nzl2)")
			("average-deltas", value<bool>(&args.averageDeltas), "Whether to average deltas when synchronizing [false]")
		;

		positional_options_description pdesc;
		pdesc.add("input-file", 1);
		pdesc.add("input-test-file", 2);
		pdesc.add("input-row-file", 3);
		pdesc.add("input-col-file", 4);
		pdesc.add("output-row-file", 5);
		pdesc.add("output-col-file", 6);

		variables_map vm;
		store(command_line_parser(argc, argv).options(desc).positional(pdesc).run(), vm);
		notify(vm);

		// check required arguments
		if (vm.count("help") || vm.count("input-file") == 0 || vm.count("input-row-file") == 0 ||
			vm.count("input-col-file") == 0 || vm.count("update") == 0 || vm.count("regularize") == 0 ||
			vm.count("loss") == 0 || vm.count("rank") == 0 || vm.count("decay") == 0)
		{
			cerr << "Error: Options input-file, input-row-file, input-col-file, update, regularize, loss, rank, and decay have to be present" << endl;
			cout << "mfasgd [options]" << endl << endl;
			cout << desc << endl;
			exit(1);
		}

		// set default arguments
		if (vm.count("seed") == 0) args.seed = time(NULL);
		if (vm.count("trace") == 0) { args.traceFile = ""; }
		if (vm.count("trace-var") == 0) args.traceVar = "trace";
		if (vm.count("input-test-file") == 0) args.inputTestMatrixFile = "";
		if (vm.count("seed") == 0) args.seed = time(NULL);
		if (vm.count("epochs") == 0) args.epochs = 10;
		if (vm.count("tasks-per-rank") == 0) args.tasksPerRank = 1;
		if (vm.count("sgd-order") == 0) { args.sgdOrderString = "WOR"; args.sgdOrder = SGD_ORDER_WOR; }
		if (vm.count("stratum-order") == 0) { args.stratumOrderString = "WOR"; args.stratumOrder = STRATUM_ORDER_WOR; }
		if (vm.count("output-row-file") == 0) { args.outputRowFacFile = ""; }
		if (vm.count("output-col-file") == 0) { args.outputColFacFile = ""; }
		if (vm.count("balance") == 0) { args.balanceString = "None"; }
		if (vm.count("average-deltas") == 0) { args.averageDeltas = false; }

		// print some information
		LOG4CXX_INFO(logger, "Input");
		LOG4CXX_INFO(logger, "    Input file: " << args.inputMatrixFile);
		LOG4CXX_INFO(logger, "    Input test file: " << (args.inputTestMatrixFile.length() == 0 ? "Disabled" :args.inputTestMatrixFile));
		LOG4CXX_INFO(logger, "    Input row factors: " << args.inputRowFacFile);
		LOG4CXX_INFO(logger, "    Input column factors: " << args.inputColFacFile);
		LOG4CXX_INFO(logger, "Output");
		LOG4CXX_INFO(logger, "    Output row factors: " << (args.outputRowFacFile.length() == 0 ? "Disabled" :args.outputRowFacFile));
		LOG4CXX_INFO(logger, "    Output column factors: " << (args.outputColFacFile.length() == 0 ? "Disabled" :args.outputColFacFile));
		if (args.traceFile.length() == 0) {
			LOG4CXX_INFO(logger, "    Trace: Disabled");
		} else {
			LOG4CXX_INFO(logger, "    Trace: " << args.traceFile << " (" << args.traceVar << ")");
		}
		LOG4CXX_INFO(logger, "Parallelization");
		LOG4CXX_INFO(logger, "    MPI ranks: " << world.size());
		LOG4CXX_INFO(logger, "    Tasks per rank: " << args.tasksPerRank);
		LOG4CXX_INFO(logger, "ASGD options");
		LOG4CXX_INFO(logger, "    Seed: " << args.seed);
		LOG4CXX_INFO(logger, "    Epochs: " << args.epochs);
		LOG4CXX_INFO(logger, "    Factorization rank: " << args.rank);
		LOG4CXX_INFO(logger, "    Update function: " << args.updateString);
		LOG4CXX_INFO(logger, "    Regularize function: " << args.regularizeString);
		LOG4CXX_INFO(logger, "    Loss function: " << args.lossString);
		LOG4CXX_INFO(logger, "    Decay: " << args.decayString);

		// parse balancing
		args.balanceMethod = BALANCE_SIMPLE;
		if (args.balanceString.compare("None") == 0) {
			LOG4CXX_INFO(logger, "    Balancing: Disabled");
			args.balanceType = BALANCE_NONE;
		} else if (args.balanceString.compare("L2") == 0) {
			LOG4CXX_INFO(logger, "    Balancing: L2");
			args.balanceType = BALANCE_L2;
		} else if (args.balanceString.compare("NZL2") == 0) {
			LOG4CXX_INFO(logger, "    Balancing: Nzl2");
			args.balanceType = BALANCE_NZL2;
		} else {
			cerr << "Invalid arguments for balance; expected \"None\", \"L2\" or \"Nzl2\"" << endl;
			exit(1);
		}

		// parse abs
		if (vm.count("abs") == 0) {
			LOG4CXX_INFO(logger, "    Absolute function: Disabled");
			args.abs = false;
		} else {
			LOG4CXX_INFO(logger, "    Absolute function: Enabled");
			args.abs = true;
		}

		// parse truncate
		if (vm.count("truncate") == 0) {
			args.truncateArgs.resize(2);
			args.truncateArgs[0] = NAN;
			args.truncateArgs[1] = NAN;
			LOG4CXX_INFO(logger, "    Truncation: Disabled");
		} else {
			parse::parseTruncate("truncate", args.truncateString, args);
			LOG4CXX_INFO(logger, "    Truncation: (" << args.truncateArgs[0] << "," << args.truncateArgs[1] << ")");
		}

		// parse sgdOrder
		if (args.sgdOrderString.compare("SEQUENTIAL") == 0) {
			LOG4CXX_INFO(logger, "    SGD step sequence: SEQ");
			args.sgdOrder = SGD_ORDER_SEQ;
		} else if (args.sgdOrderString.compare("WOR") == 0) {
			LOG4CXX_INFO(logger, "    SGD step sequence: WOR");
			args.sgdOrder = SGD_ORDER_WOR;
		} else if (args.sgdOrderString.compare("WR") == 0) {
			LOG4CXX_INFO(logger, "    SGD step sequence: WR");
			args.sgdOrder = SGD_ORDER_WR;
		} else {
			cerr << "Invalid arguments for sgdOrder; expected \"SEQUENTIAL\", \"WOR\" or \"WR\"" << endl;
			exit(1);
		}

		// print averaging
		LOG4CXX_INFO(logger, "    Average deltas: " << args.averageDeltas);

		// fill fields
		args.random = Random32(args.seed);
		args.world = world;
		args.worldSize = world.size();
		args.worldRank = world.rank();
		args.tries = args.worldSize * args.tasksPerRank;
		while (args.tries < 7) args.tries += args.worldSize;
		args.blocks1 = args.worldSize;
		args.blocks2 = args.worldSize;

		// parsing
		parse::parseArg("update", args.updateString, args.updateName, args.updateArgs);
		parse::parseArg("regularize", args.regularizeString, args.regularizeName, args.regularizeArgs);
		parse::parseArg("loss", args.lossString, args.lossName, args.lossArgs);
		parse::parseDecay("decay", args.decayString, args);

		// let's go
		result = runArgs(args);
	}

	mfStop();
	mfFinalize();

	// everything OK
	return result ? 0 : 1;

}