/**
 * 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
 *
 */

/**
 * \file
 * This code iplements the NMF algorithm described in the paper:
 * Lee, D..D., and Seung, H.S., (2001), 'Algorithms for Non-negative Matrix
 * Factorization', Adv. Neural Info. Proc. Syst. 13, 556-562.
 */

#include <graphlab/util/stl_util.hpp>
#include <graphlab.hpp>

#include <Eigen/Dense>
#include "eigen_serialization.hpp"
#include <graphlab/macros_def.hpp>
#include <graphlab/util/timer.hpp>
#include "stats.hpp"

typedef Eigen::VectorXd vec;
typedef Eigen::MatrixXd mat_type;

//when using negative node id range, we are not allowed to use
//0 and 1 so we add 2.
const static int SAFE_NEG_OFFSET=2;
const double epsilon = 1e-16;
static bool debug;
int iter = 0;

bool isuser(uint node){
  return ((int)node) >= 0;
}
/**
 * \ingroup toolkit_matrix_pvecization
 *
 * \brief the vertex data type which contains the latent pvec.
 *
 * Each row and each column in the matrix corresponds to a different
 * vertex in the SGD graph.  Associated with each vertex is a pvec
 * (vector) of latent parameters that represent that vertex.  The goal
 * of the SGD algorithm is to find the values for these latent
 * parameters such that the non-zero entries in the matrix can be
 * predicted by taking the dot product of the row and column pvecs.
 */
struct vertex_data {
  /**
   * \brief A shared "constant" that specifies the number of latent
   * values to use.
   */
  static size_t NLATENT;
  /** \brief The latent pvec for this vertex */
  vec pvec;

  double train_rmse;
  double validation_rmse;
  /**
   * \brief Simple default constructor which randomizes the vertex
   *  data
   */
  vertex_data() { if (debug) pvec = vec::Ones(NLATENT); else randomize(); train_rmse = validation_rmse = 0; }
  /** \brief Randomizes the latent pvec */
  void randomize() { pvec.resize(NLATENT); pvec.setRandom(); }
  /** \brief Save the vertex data to a binary archive */
  void save(graphlab::oarchive& arc) const {
    arc << pvec << train_rmse << validation_rmse;
  }
  /** \brief Load the vertex data from a binary archive */
  void load(graphlab::iarchive& arc) {
    arc >> pvec >>train_rmse >> validation_rmse;
  }
}; // end of vertex data

/**
 * \brief The edge data stores the entry in the matrix.
 *
 * In addition the edge data nmfo stores the most recent error estimate.
 */
struct edge_data : public graphlab::IS_POD_TYPE {
  /**
   * \brief The type of data on the edge;
   *
   * \li *Train:* the observed value is correct and used in training
   * \li *Validate:* the observed value is correct but not used in training
   * \li *Predict:* The observed value is not correct and should not be
   *        used in training.
   */
  enum data_role_type { TRAIN, VALIDATE, PREDICT  };

  /** \brief the observed value for the edge */
  float weight;

  /** \brief The train/validation/test designation of the edge */
  data_role_type role;

  /** \brief basic initialization */
  edge_data(float weight = 0, data_role_type role = PREDICT) :
    weight(weight), role(role) { }

}; // end of edge data

/**
 * \brief The graph type is defined in terms of the vertex and edge
 * data.
 */
typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

double extract_l2_error(const graph_type::edge_type & edge);

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

class gather_type {
  public:
    vec pvec;
    double training_rmse;
    double validation_rmse;
    gather_type() { training_rmse = validation_rmse = 0; }
    gather_type(const vec & _pvec, double _train_rmse, double _validation_rmse){ pvec = _pvec; training_rmse = _train_rmse; validation_rmse = _validation_rmse; }
    void reset(){ pvec = vec::Zero(vertex_data::NLATENT); training_rmse = 0; validation_rmse = 0; }
    void save(graphlab::oarchive& arc) const { arc << pvec << training_rmse << validation_rmse; }
    void load(graphlab::iarchive& arc) { arc >> pvec >> training_rmse >> validation_rmse; }
    gather_type& operator+=(const gather_type& other) {
      pvec += other.pvec;
      training_rmse += other.training_rmse;
      validation_rmse += other.validation_rmse;
      return *this;
    }

};

gather_type x1;
gather_type x2;
gather_type * px;

bool isuser_node(const graph_type::vertex_type& vertex){
  return isuser(vertex.id());
}
/**
 * SGD vertex program type
 */
class nmf_vertex_program :
  public graphlab::ivertex_program<graph_type, gather_type, gather_type>,
  public graphlab::IS_POD_TYPE{
    public:
      /** The convergence tolerance */
      static double TOLERANCE;
      static double MAXVAL;
      static double MINVAL;
      static bool debug;
      static size_t MAX_UPDATES;

      /** compute a missing value based on NMF algorithm */
      static float nmf_predict(const vertex_data& user,
          const vertex_data& movie,
          const float rating,
          double & prediction){

        prediction = user.pvec.dot(movie.pvec);
        //truncate prediction to allowed values
        prediction = std::min((double)prediction, nmf_vertex_program::MAXVAL);
        prediction = std::max((double)prediction, nmf_vertex_program::MINVAL);
        //return the squared error
        float err = rating - prediction;
        assert(!std::isnan(err));
        return err*err;

      }

      /** The set of edges to gather along */
      edge_dir_type gather_edges(icontext_type& context,
          const vertex_type& vertex) const {
        //UNUSED
        return graphlab::ALL_EDGES;
      }; // end of gather_edges

      /** The gather function computes XtX and Xy */
      gather_type gather(icontext_type& context, const vertex_type& vertex,
          edge_type& edge) const {

        if (edge.data().role == edge_data::TRAIN || edge.data().role == edge_data::VALIDATE){
          const vertex_type other_vertex = get_other_vertex(edge, vertex);
          double prediction = 0;
          double rmse = nmf_predict(vertex.data(), other_vertex.data(), edge.data().weight, prediction);
          if (prediction == 0)
            logstream(LOG_FATAL)<<"Got into numerical error!" << std::endl;
          if (edge.data().role == edge_data::TRAIN)
            return gather_type(other_vertex.data().pvec * (edge.data().weight / prediction), rmse, 0);
          else //validation
            return gather_type(vec::Zero(vertex_data::NLATENT), 0, rmse);

        }
        return gather_type(vec::Zero(vertex_data::NLATENT), 0, 0);
      } // end of gather function

      void apply(icontext_type& context, vertex_type& vertex,
          const gather_type& sum) {
        vertex_data& vdata = vertex.data();
        if (vdata.pvec.sum() != 0){
          for (uint i=0; i< vertex_data::NLATENT; i++){
            vdata.pvec[i] *= sum.pvec[i] / px->pvec[i];
            ASSERT_NE(px->pvec[i] , 0);
            if (vdata.pvec[i] < epsilon)
              vdata.pvec[i] = epsilon;
          }
        }
        vdata.train_rmse = sum.training_rmse;
        vdata.validation_rmse = sum.validation_rmse;
      }

      edge_dir_type scatter_edges(icontext_type& context,
          const vertex_type& vertex) const {
        //UNUSED
        return graphlab::ALL_EDGES;
      }; // end of scatter edges

      void scatter(icontext_type& context, const vertex_type& vertex,
          edge_type& edge) const {
        //we do not schedule any more neighbors to run
      }
      static void verify_rows(graph_type::vertex_type& vertex){
        if (isuser(vertex.id()) && vertex.num_out_edges() == 0)
          logstream(LOG_FATAL)<<"NMF algorithm can not work when the row " << vertex.id() << " of the matrix contains all zeros" << std::endl;
      }

      static gather_type pre_iter(const graph_type::vertex_type & vertex){
        gather_type ret;
        ret.pvec = vertex.data().pvec;
        ret.training_rmse = vertex.data().train_rmse;
        ret.validation_rmse = vertex.data().validation_rmse;
        return ret;
      }

      static graphlab::empty signal_left(icontext_type& context,
          const vertex_type& vertex) {
        if(vertex.num_out_edges() > 0) context.signal(vertex);
        return graphlab::empty();
      } // end of signal_left

      static graphlab::empty signal_right(icontext_type& context,
          const vertex_type& vertex) {
        if(vertex.num_in_edges() > 0) context.signal(vertex);
        return graphlab::empty();
      } // end of signal_left

  }; // end of nmf vertex program

gather_type count_edges(nmf_vertex_program::icontext_type & context, const graph_type::edge_type& edge) {
  gather_type ret;
  if (edge.data().role == edge_data::TRAIN){
    ret.training_rmse = 1;
  }
  else if (edge.data().role == edge_data::VALIDATE){
    ret.validation_rmse = 1;
  }
  if (edge.data().weight < 0)
    logstream(LOG_FATAL)<<"Found a negative entry in matirx row " << edge.source().id() << " with value: " << edge.data().weight << std::endl;
  return ret;
}

struct prediction_saver {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  /* save the linear model, using the format:
     nodeid) factor1 factor2 ... factorNLATENT \n
     */
  std::string save_vertex(const vertex_type& vertex) const {
    return "";
  }
  std::string save_edge(const edge_type& edge) const {
    if (edge.data().role != edge_data::PREDICT)
      return "";

    std::stringstream strm;
    const double prediction =
      edge.source().data().pvec.dot(edge.target().data().pvec);
    strm << edge.source().id() << '\t'
      << -edge.target().id()-SAFE_NEG_OFFSET << '\t'
      << prediction << '\n';
    return strm.str();
  }
}; // end of prediction_saver

struct linear_model_saver_U {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  /* save the linear model, using the format:
     nodeid) factor1 factor2 ... factorNLATENT \n
     */
  std::string save_vertex(const vertex_type& vertex) const {
    if (vertex.num_out_edges() > 0){
      std::string ret = boost::lexical_cast<std::string>(vertex.id()) + ") ";
      for (uint i=0; i< vertex_data::NLATENT; i++)
        ret += boost::lexical_cast<std::string>(vertex.data().pvec[i]) + " ";
      ret += "\n";
      return ret;
    }
    else return "";
  }
  std::string save_edge(const edge_type& edge) const {
    return "";
  }
};

struct linear_model_saver_V {
  typedef graph_type::vertex_type vertex_type;
  typedef graph_type::edge_type   edge_type;
  /* save the linear model, using the format:
     nodeid) factor1 factor2 ... factorNLATENT \n
     */
  std::string save_vertex(const vertex_type& vertex) const {
    if (vertex.num_out_edges() == 0){
      std::string ret = boost::lexical_cast<std::string>(-vertex.id()-SAFE_NEG_OFFSET) + ") ";
      for (uint i=0; i< vertex_data::NLATENT; i++)
        ret += boost::lexical_cast<std::string>(vertex.data().pvec[i]) + " ";
      ret += "\n";
      return ret;
    }
    else return "";
  }
  std::string save_edge(const edge_type& edge) const {
    return "";
  }
};

/**
 * \brief The graph loader function is a line parser used for
 * distributed graph construction.
 */
inline bool graph_loader(graph_type& graph,
    const std::string& filename,
    const std::string& line) {
  ASSERT_FALSE(line.empty());

 // Parse the line
  std::stringstream strm(line);
  graph_type::vertex_id_type source_id(-1), target_id(-1);
  float weight(0);
  strm >> source_id >> target_id;

  if (source_id == graph_type::vertex_id_type(-1) || target_id == graph_type::vertex_id_type(-1)){
    logstream(LOG_WARNING)<<"Failed to read input line: "<< line << " in file: "  << filename << " (or node id is -1). " << std::endl;
    return true;
  }

  // Determine the role of the data
  edge_data::data_role_type role = edge_data::TRAIN;
  if(boost::ends_with(filename,".validate")) role = edge_data::VALIDATE;
  else if(boost::ends_with(filename, ".predict")) role = edge_data::PREDICT;

  // for test files (.predict) no need to read the actual rating value.
  if(role == edge_data::TRAIN || role == edge_data::VALIDATE){
    strm >> weight;
    if (weight < nmf_vertex_program::MINVAL || weight > nmf_vertex_program::MAXVAL)
      logstream(LOG_FATAL)<<"Rating values should be between " << nmf_vertex_program::MINVAL << " and " << nmf_vertex_program::MAXVAL << ". Got value: " << weight << " [ user: " << source_id << " to item: " <<target_id << " ] " << std::endl;
  }
  target_id = -(graphlab::vertex_id_type(target_id + SAFE_NEG_OFFSET));

  // Create an edge and add it to the graph
  graph.add_edge(source_id, target_id, edge_data(weight, role));
  return true; // successful load
} // end of graph_loader

size_t vertex_data::NLATENT = 20;
double nmf_vertex_program::TOLERANCE = 1e-3;
size_t nmf_vertex_program::MAX_UPDATES = -1;
double nmf_vertex_program::MAXVAL = 1e+100;
double nmf_vertex_program::MINVAL = -1e+100;
bool nmf_vertex_program::debug = false;

/**
 * \brief The engine type used by the ALS matrix factorization
 * algorithm.
 *
 * The ALS matrix factorization algorithm currently uses the
 * synchronous engine.  However we plan to add support for alternative
 * engines in the future.
 */
typedef graphlab::omni_engine<nmf_vertex_program> engine_type;

int main(int argc, char** argv) {
  global_logger().set_log_level(LOG_INFO);
  global_logger().set_log_to_console(true);

  // Parse command line options -----------------------------------------------
  const std::string description =
    "Compute the ALS factorization of a matrix.";
  graphlab::command_line_options clopts(description);
  std::string input_dir;
  std::string predictions;
  std::string exec_type = "synchronous";
  clopts.attach_option("matrix", input_dir,
      "The directory containing the matrix file");
  clopts.add_positional("matrix");
  clopts.attach_option("D", vertex_data::NLATENT,
      "Number of latent parameters to use.");
  clopts.attach_option("engine", exec_type,
      "The engine type synchronous or asynchronous");
  clopts.attach_option("max_iter", nmf_vertex_program::MAX_UPDATES,
      "The maxumum number of udpates allowed for a vertex");
  clopts.attach_option("debug", nmf_vertex_program::debug,
      "debug - additional verbose info");
  clopts.attach_option("maxval", nmf_vertex_program::MAXVAL, "max allowed value");
  clopts.attach_option("minval", nmf_vertex_program::MINVAL, "min allowed value");
  clopts.attach_option("predictions", predictions,
      "The prefix (folder and filename) to save predictions.");

  if(!clopts.parse(argc, argv) || input_dir == "") {
    std::cout << "Error in parsing command line arguments." << std::endl;
    clopts.print_description();
    return EXIT_FAILURE;
  }
  debug = nmf_vertex_program::debug;

  graphlab::mpi_tools::init(argc, argv);
  graphlab::distributed_control dc;

  dc.cout() << "Loading graph." << std::endl;
  graphlab::timer timer;
  graph_type graph(dc, clopts);
  graph.load(input_dir, graph_loader);
  dc.cout() << "Loading graph. Finished in "
    << timer.current_time() << std::endl;
  dc.cout() << "Finalizing graph." << std::endl;
  timer.start();
  graph.finalize();
  dc.cout() << "Finalizing graph. Finished in "
    << timer.current_time() << std::endl;

  if (!graph.num_edges() || !graph.num_vertices())
     logstream(LOG_FATAL)<< "Failed to load graph. Check your input path: " << input_dir << std::endl;

  dc.cout()
    << "========== Graph statistics on proc " << dc.procid()
    << " ==============="
    << "\n Num vertices: " << graph.num_vertices()
    << "\n Num edges: " << graph.num_edges()
    << "\n Num replica: " << graph.num_replicas()
    << "\n Replica to vertex ratio: "
    << float(graph.num_replicas())/graph.num_vertices()
    << "\n --------------------------------------------"
    << "\n Num local own vertices: " << graph.num_local_own_vertices()
    << "\n Num local vertices: " << graph.num_local_vertices()
    << "\n Replica to own ratio: "
    << (float)graph.num_local_vertices()/graph.num_local_own_vertices()
    << "\n Num local edges: " << graph.num_local_edges()
    //<< "\n Begin edge id: " << graph.global_eid(0)
    << "\n Edge balance ratio: "
    << float(graph.num_local_edges())/graph.num_edges()
    << std::endl;

  dc.cout() << "Creating engine" << std::endl;
  engine_type engine(dc, graph, exec_type, clopts);

  // Run the NMF ---------------------------------------------------------
  dc.cout() << "Running NMF" << std::endl;
  dc.cout() << "(C) Code by Danny Bickson, CMU " << std::endl;
  dc.cout() << "Please send bug reports to danny.bickson@gmail.com" << std::endl;
  dc.cout() << "Time   Training    Validation" <<std::endl;
  dc.cout() << "       RMSE        RMSE " <<std::endl;
  timer.start();

  gather_type edge_count = engine.map_reduce_edges<gather_type>(count_edges);
  dc.cout()<<"Training edges: " << edge_count.training_rmse << " validation edges: " << edge_count.validation_rmse << std::endl;

  graphlab::vertex_set left = graph.select(isuser_node);
  graphlab::vertex_set right = ~left;
  graph.transform_vertices(nmf_vertex_program::verify_rows, left);

  graphlab::timer mytimer; mytimer.start();

  for (uint j=0; j< nmf_vertex_program::MAX_UPDATES; j++){
    x1 = graph.map_reduce_vertices<gather_type>(nmf_vertex_program::pre_iter,right);
    px = &x1;
    for (int i=0; i< (int)vertex_data::NLATENT; i++)
      ASSERT_NE(px->pvec[i], 0);

    dc.cout()<< std::setw(8) << mytimer.current_time() << " " << sqrt(x1.training_rmse/edge_count.training_rmse);
    if (edge_count.validation_rmse > 0)
      dc.cout() << " " << std::setw(8) << sqrt(x1.validation_rmse/edge_count.validation_rmse) << std::endl;
    else dc.cout() << std::endl;
    engine.map_reduce_vertices<graphlab::empty>(nmf_vertex_program::signal_left);
    engine.start();
    x1.reset();

    x2 = graph.map_reduce_vertices<gather_type>(nmf_vertex_program::pre_iter,left);
    px = &x2;

    engine.map_reduce_vertices<graphlab::empty>(nmf_vertex_program::signal_right);
    engine.start();
    x2.reset();
  }

  const double runtime = timer.current_time();
  dc.cout() << "----------------------------------------------------------"
    << std::endl
    << "Final Runtime (seconds):   " << runtime
                                        << std::endl
                                        << "Updates executed: " << engine.num_updates() << std::endl
                                        << "Update Rate (updates/second): "
                                          << engine.num_updates() / runtime << std::endl;

  // Make predictions ---------------------------------------------------------
  if(!predictions.empty()) {
    std::cout << "Saving predictions" << std::endl;
    const bool gzip_output = false;
    const bool save_vertices = false;
    const bool save_edges = true;
    const size_t threads_per_machine = 1;
    //save the predictions
    graph.save(predictions, prediction_saver(),
        gzip_output, save_vertices,
        save_edges, threads_per_machine);
    //save the linear model
    graph.save(predictions + ".U", linear_model_saver_U(),
        gzip_output, save_edges, save_vertices, threads_per_machine);
    graph.save(predictions + ".V", linear_model_saver_V(),
        gzip_output, save_edges, save_vertices, threads_per_machine);

  }

  graphlab::mpi_tools::finalize();
  return EXIT_SUCCESS;
} // end of main


#ifndef BRESENHAM_LINE_ALGORITHMS_HPP
#define BRESENHAM_LINE_ALGORITHMS_HPP

#include <boost/gil/gil_all.hpp>
#include "line-struct.hpp"

// uses the bresenham algorithm to draw a 1 pixel grayscale line on the image view img
void drawGrayscaleLine(boost::gil::gray8_view_t &img, LineData line, int opacity, bool additive);
void drawGrayscaleLineNormalAngleAdditive(boost::gil::gray8_view_t &img, LineData line, int opacity);
void drawGrayscaleLineSteepAngleAdditive(boost::gil::gray8_view_t &img, LineData line, int opacity);
void drawGrayscaleLineNormalAngleSubtractive(boost::gil::gray8_view_t &img, LineData line, int opacity);
void drawGrayscaleLineSteepAngleSubtractive(boost::gil::gray8_view_t &img, LineData line, int opacity);

// uses the bresenham algorithm to calculate a score of how much a drawn 1 pixel grayscale line would change the image
unsigned long grayscaleLineScore(boost::gil::gray8_view_t &img, LineData line, int opacity, bool additive);
unsigned long grayscaleLineScoreNormalAngleAdditive(boost::gil::gray8_view_t &img, LineData line, int opacity);
unsigned long grayscaleLineScoreSteepAngleAdditive(boost::gil::gray8_view_t &img, LineData line, int opacity);
unsigned long grayscaleLineScoreNormalAngleSubtractive(boost::gil::gray8_view_t &img, LineData line, int opacity);
unsigned long grayscaleLineScoreSteepAngleSubtractive(boost::gil::gray8_view_t &img, LineData line, int opacity);

// uses the bresenham algorithm to draw a 1 pixel rgb line on the image view img
void drawRGBLine(boost::gil::rgb8_view_t &img, LineData line, int opacity, bool additive, int channel);
void drawRGBLineNormalAngleAdditive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
void drawRGBLineSteepAngleAdditive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
void drawRGBLineNormalAngleSubtractive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
void drawRGBLineSteepAngleSubtractive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);

// uses the bresenham algorithm to calculate a score of how much a drawn 1 pixel rgb line would change the image
unsigned long rgbLineScore(boost::gil::rgb8_view_t &img, LineData line, int opacity, bool additive, int channel);
unsigned long rgbLineScoreNormalAngleAdditive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
unsigned long rgbLineScoreSteepAngleAdditive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
unsigned long rgbLineScoreNormalAngleSubtractive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);
unsigned long rgbLineScoreSteepAngleSubtractive(boost::gil::rgb8_view_t &img, LineData line, int opacity, int c);

#endif // BRESENHAM_LINE_ALGORITHMS_HPP


//
// Copyright 2019-2020 Mateusz Loskot <mateusz at loskot dot net>
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//
#include <boost/gil.hpp>
#include <boost/gil/extension/numeric/affine.hpp>
#include <boost/gil/extension/numeric/resample.hpp>
#include <boost/gil/extension/numeric/sampler.hpp>

#include <boost/core/lightweight_test.hpp>

#include <cmath>

namespace gil = boost::gil;

// FIXME: Remove when https://github.com/boostorg/core/issues/38 happens
#define BOOST_GIL_TEST_IS_CLOSE(a, b, epsilon) BOOST_TEST_LT(std::fabs((a) - (b)), (epsilon))

namespace {
constexpr double HALF_PI = 1.57079632679489661923;
}

void test_matrix3x2_default_constructor()
{
    gil::matrix3x2<int> m1;
    BOOST_TEST_EQ(m1.a, 1);
    BOOST_TEST_EQ(m1.b, 0);
    BOOST_TEST_EQ(m1.c, 0);
    BOOST_TEST_EQ(m1.d, 1);
    BOOST_TEST_EQ(m1.e, 0);
    BOOST_TEST_EQ(m1.f, 0);
}

void test_matrix3x2_parameterized_constructor()
{
    gil::matrix3x2<int> m1(1, 2, 3, 4, 5, 6);
    BOOST_TEST_EQ(m1.a, 1);
    BOOST_TEST_EQ(m1.b, 2);
    BOOST_TEST_EQ(m1.c, 3);
    BOOST_TEST_EQ(m1.d, 4);
    BOOST_TEST_EQ(m1.e, 5);
    BOOST_TEST_EQ(m1.f, 6);
}

void test_matrix3x2_copy_constructor()
{
    gil::matrix3x2<int> m1(1, 2, 3, 4, 5, 6);
    gil::matrix3x2<int> m2(m1);
    BOOST_TEST_EQ(m2.a, 1);
    BOOST_TEST_EQ(m2.b, 2);
    BOOST_TEST_EQ(m2.c, 3);
    BOOST_TEST_EQ(m2.d, 4);
    BOOST_TEST_EQ(m2.e, 5);
    BOOST_TEST_EQ(m2.f, 6);
}

void test_matrix3x2_assignment_operator()
{
    gil::matrix3x2<int> m1(1, 2, 3, 4, 5, 6);
    gil::matrix3x2<int> m2;
    m2 = m1;
    BOOST_TEST_EQ(m2.a, 1);
    BOOST_TEST_EQ(m2.b, 2);
    BOOST_TEST_EQ(m2.c, 3);
    BOOST_TEST_EQ(m2.d, 4);
    BOOST_TEST_EQ(m2.e, 5);
    BOOST_TEST_EQ(m2.f, 6);
}

void test_matrix3x2_multiplication_assignment()
{
    gil::matrix3x2<int> m1;
    gil::matrix3x2<int> m2;
    m2 *= m1;
    BOOST_TEST_EQ(m2.a, 1);
    BOOST_TEST_EQ(m2.b, 0);
    BOOST_TEST_EQ(m2.c, 0);
    BOOST_TEST_EQ(m2.d, 1);
    BOOST_TEST_EQ(m2.e, 0);
    BOOST_TEST_EQ(m2.f, 0);

    gil::matrix3x2<int> m3(0, 0, 0, 0, 0, 0);
    m2 *= m3;
    BOOST_TEST_EQ(m2.a, 0);
    BOOST_TEST_EQ(m2.b, 0);
    BOOST_TEST_EQ(m2.c, 0);
    BOOST_TEST_EQ(m2.d, 0);
    BOOST_TEST_EQ(m2.e, 0);
    BOOST_TEST_EQ(m2.f, 0);
}

void test_matrix3x2_matrix3x2_multiplication()
{
    gil::matrix3x2<int> m1;
    gil::matrix3x2<int> m2(0, 0, 0, 0, 0, 0);
    gil::matrix3x2<int> m3;
    m3 = m1 * m2;
    BOOST_TEST_EQ(m3.a, 0);
    BOOST_TEST_EQ(m3.b, 0);
    BOOST_TEST_EQ(m3.c, 0);
    BOOST_TEST_EQ(m3.d, 0);
    BOOST_TEST_EQ(m3.e, 0);
    BOOST_TEST_EQ(m3.f, 0);
}

void test_matrix3x2_vector_multiplication()
{
    gil::matrix3x2<int> m1;
    gil::point<int> v1{2, 4};

    gil::point<int> v2 = v1 * m1;
    BOOST_TEST_EQ(v2.x, 2);
    BOOST_TEST_EQ(v2.y, 4);

    gil::point<int> v3 = gil::transform(m1, v1);
    BOOST_TEST_EQ(v3.x, 2);
    BOOST_TEST_EQ(v3.y, 4);
}

void test_matrix3x2_get_rotate()
{
    auto m1 = gil::matrix3x2<double>::get_rotate(HALF_PI);
    BOOST_GIL_TEST_IS_CLOSE(m1.a, std::cos(HALF_PI), 0.03);
    BOOST_TEST_EQ(m1.b, 1);
    BOOST_TEST_EQ(m1.c, -1);
    BOOST_GIL_TEST_IS_CLOSE(m1.d, std::cos(HALF_PI), 0.03);
    BOOST_TEST_EQ(m1.e, 0);
    BOOST_TEST_EQ(m1.f, 0);
}

void test_matrix3x2_get_scale()
{
    gil::matrix3x2<int> m1;
    m1 = gil::matrix3x2<int>::get_scale(2);
    BOOST_TEST_EQ(m1.a, 2);
    BOOST_TEST_EQ(m1.b, 0);
    BOOST_TEST_EQ(m1.c, 0);
    BOOST_TEST_EQ(m1.d, 2);
    BOOST_TEST_EQ(m1.e, 0);
    BOOST_TEST_EQ(m1.f, 0);
    m1 = gil::matrix3x2<int>::get_scale(2, 4);
    BOOST_TEST_EQ(m1.a, 2);
    BOOST_TEST_EQ(m1.d, 4);
    m1 = gil::matrix3x2<int>::get_scale(gil::point<int>{4, 8});
    BOOST_TEST_EQ(m1.a, 4);
    BOOST_TEST_EQ(m1.d, 8);
}

void test_matrix3x2_get_translate()
{
    gil::matrix3x2<int> m1;
    m1 = gil::matrix3x2<int>::get_translate(2, 4);
    BOOST_TEST_EQ(m1.a, 1);
    BOOST_TEST_EQ(m1.b, 0);
    BOOST_TEST_EQ(m1.c, 0);
    BOOST_TEST_EQ(m1.d, 1);
    BOOST_TEST_EQ(m1.e, 2);
    BOOST_TEST_EQ(m1.f, 4);
    m1 = gil::matrix3x2<int>::get_translate(gil::point<int>{4, 8});
    BOOST_TEST_EQ(m1.e, 4);
    BOOST_TEST_EQ(m1.f, 8);
}

void test_matrix3x2_transform()
{
    gil::matrix3x2<int> m1;
    gil::point<int> v1{2, 4};
    gil::point<int> v2 = gil::transform(m1, v1);
    BOOST_TEST_EQ(v2.x, 2);
    BOOST_TEST_EQ(v2.y, 4);
}

int main()
{
    test_matrix3x2_default_constructor();
    test_matrix3x2_parameterized_constructor();
    test_matrix3x2_copy_constructor();
    test_matrix3x2_assignment_operator();
    test_matrix3x2_multiplication_assignment();
    test_matrix3x2_matrix3x2_multiplication();
    test_matrix3x2_vector_multiplication();
    test_matrix3x2_get_rotate();
    test_matrix3x2_get_scale();
    test_matrix3x2_get_translate();
    test_matrix3x2_transform();

    return ::boost::report_errors();
}


//   Copyright (c) 2014-2016 SSPA Sweden AB

#pragma once

#include <Eigen/Dense>
namespace pysim {
    typedef Eigen::RowVectorXd vector;
    typedef Eigen::MatrixXd matrix;
};


#pragma once
#include <boost/serialization/vector.hpp>
#include <cmath>
#include <iostream>

#include <tbb/concurrent_vector.h>
#include <tbb/concurrent_unordered_map.h>
#include <tbb/parallel_sort.h>
#include <tbb/parallel_for_each.h>

#include "BasisJz.hpp"
#include "AbstractBasis1D.hpp"

namespace edlib
{
template<typename UINT>
class Basis1D final
	: public AbstractBasis1D<UINT>
{
private:
	tbb::concurrent_vector<std::pair<UINT, uint32_t>> rpts_; //Representatives

	int checkState(UINT s)
	{
		UINT sr = s;
		const auto N = this->getN();
		const auto k = this->getK();

		for(uint32_t r = 1; r <= N; r++)
		{
			sr = this->rotl(s, r);
			if(sr < s)
			{
				return -1; //s is not a representative
			}
			else if(sr == s)
			{
				/* s is smller than rotl(s,1), rot(s, 2), ..., rot(s, r-1)
				 * when we fall in this else if clause. As rot(s,r) == s,
				 * s is the smallest among rotl(s, 1), ..., rotl(s, N-1).
				 */
				if((k % (N/r)) != 0)
					return -1; //this representative is not allowed for k
				return r;
			}
		}
		return -1;
	}

	void constructBasisFull()
	{
		// insert 0
		{
			UINT s = 0;
			int r = checkState(s);
			if(r > 0)
			{
				rpts_.emplace_back(s,r);
			}
		}

		// iterate over all odd numbers
		const uint32_t N = this->getN();
		tbb::parallel_for(UINT(1), (UINT(1)<<UINT(N)), UINT(2), [&](UINT s)
		{
			int r = checkState(s);
			if(r > 0)
			{
				rpts_.emplace_back(s,r);
			}
		});
		tbb::parallel_sort(rpts_.begin(), rpts_.end());
	}

	void constructBasisJz()
	{
		const uint32_t n = this->getN();
		const uint32_t nup = n/2;

		BasisJz<UINT> basis(n,nup);

		tbb::parallel_for_each(basis.begin(), basis.end(), [&](UINT s)
		{
			int r = checkState(s);
			if(r > 0)
			{
				rpts_.emplace_back(s,r);
			}
		});
		tbb::parallel_sort(rpts_.begin(), rpts_.end());
	}

public:
	Basis1D(uint32_t N, uint32_t k, bool useU1)
		: AbstractBasis1D<UINT>(N, k)
	{
		assert( (!useU1) || (N % 2 == 0));
		if(useU1)
		{
			constructBasisJz();
		}
		else
		{
			constructBasisFull();
		}
	}

	Basis1D(const Basis1D& ) = default;
	Basis1D(Basis1D&& ) = default;

	Basis1D& operator=(const Basis1D& ) = default;
	Basis1D& operator=(Basis1D&& ) = default;

	~Basis1D() = default;

	uint32_t stateIdx(UINT rep) const
	{
		auto comp = [](const std::pair<UINT, uint32_t>& v1, UINT v2)
		{
			return v1.first < v2;
		};
		auto iter = lower_bound(rpts_.begin(), rpts_.end(), rep, comp);
		if((iter == rpts_.end()) || (iter->first != rep))
		{
			return getDim();
		}
		else
		{
			return distance(rpts_.begin(), iter);
		}
	}

	tbb::concurrent_vector<std::pair<UINT,uint32_t>> getRepresentatives() const
	{
		return rpts_;
	}

	std::size_t getDim() const override
	{
		return rpts_.size();
	}

	UINT getNthRep(uint32_t n) const override
	{
		return rpts_[n].first;
	}

	inline uint32_t rotRpt(uint32_t n) const
	{
		return rpts_[n].second;
	}

	std::pair<int, double> hamiltonianCoeff(UINT bSigma, int aidx) const override
	{
		using std::sqrt;
		using std::pow;

		const auto k = this->getK();
		double expk = (k == 0)?1.0:-1.0;

		UINT bRep;
		int bRot;
		std::tie(bRep, bRot) = this->findMinRots(bSigma);

		auto bidx = stateIdx(bRep);

		if(bidx >= getDim())
		{
			return std::make_pair(-1, 0.0);
		}

		double Na = 1.0/rpts_[aidx].second;
		double Nb = 1.0/rpts_[bidx].second;

		return std::make_pair(bidx, sqrt(Nb/Na)*pow(expk, bRot));
	}

	std::vector<std::pair<UINT, double>> basisVec(uint32_t n) const override
	{
		const auto k = this->getK();
		const double expk = (k == 0)?1.0:-1.0;
		std::vector<std::pair<UINT,double>> res;

		auto rep = rpts_[n].first;
		double norm = 1.0/sqrt(rpts_[n].second);
		for(uint32_t r = 0; r < rpts_[n].second; r++)
		{
			res.emplace_back( this->rotl(rep, r), pow(expk, r)*norm);
		}
		return res;
	}
};
} // namespace edlib


// A dummy GLUP application

#include <polyfem/Navigation.hpp>
#include <polyfem/MeshUtils.hpp>
#include <polyfem/PolygonUtils.hpp>
#include <polyfem/Singularities.hpp>
#include <polyfem/Refinement.hpp>
#include <geogram/basic/file_system.h>
#include <geogram/basic/command_line.h>
#include <geogram/basic/command_line_args.h>
#include <geogram/mesh/mesh_geometry.h>
#include <geogram/mesh/mesh_io.h>
#include <geogram/mesh/mesh_preprocessing.h>
#include <geogram_gfx/glup_viewer/glup_viewer.h>
#include <geogram_gfx/glup_viewer/glup_viewer_gui.h>
#include <geogram_gfx/mesh/mesh_gfx.h>
#include <Eigen/Dense>

namespace {
	using namespace std;
	using namespace GEO;
	using namespace polyfem;

	Navigation::Index idx_;

	Eigen::VectorXi singular_vertices_;
	Eigen::MatrixX2i singular_edges_;

	float t_;

	bool show_index_;

	class DemoGlupApplication : public SimpleMeshApplication {
	public:
		DemoGlupApplication(int argc, char** argv) :
			SimpleMeshApplication(argc, argv, "<filename>")
		{
			name_ = "[Float] Navigation";
			t_ = 0.5;
			show_index_ = true;
		}

		virtual void init_graphics() override {
			SimpleMeshApplication::init_graphics();
			glup_viewer_disable(GLUP_VIEWER_BACKGROUND);
			glup_viewer_disable(GLUP_VIEWER_3D);
			ImGui::GetStyle().WindowRounding = 7.0f;
			ImGui::GetStyle().FrameRounding = 0.0f;
			ImGui::GetStyle().GrabRounding = 0.0f;
			retina_mode_ = false;
			scaling_ = 1.0f;
		}

		virtual bool load(const std::string &filename) override {
			if(!GEO::FileSystem::is_file(filename)) {
				GEO::Logger::out("I/O") << "is not a file" << std::endl;
				return false;
			}
			SimpleMeshApplication::load(filename);
			mesh_.vertices.set_double_precision();
			polyfem::orient_normals_2d(mesh_);
			mesh_.vertices.set_single_precision();

			// Compute mesh connectivity
			Navigation::prepare_mesh(mesh_);
			GEO::Attribute<GEO::index_t> c2e(mesh_.facet_corners.attributes(), "edge_id");

			// Initialize the key
			idx_ = Navigation::get_index_from_face(mesh_, c2e, 0, 0);

			// Compute singularities
			polyfem::singularity_graph(mesh_, singular_vertices_, singular_edges_);

			// Compute element types
			compute_types();

			return true;
		}

		void compute_types() {
			std::vector<ElementType> tags;
			polyfem::compute_element_tags(mesh_, tags);
			GEO::Attribute<int> attrs(mesh_.facets.attributes(), "tags");
			for (index_t f = 0; f < mesh_.facets.nb(); ++f) {
				switch (tags[f]) {
				case ElementType::RegularInteriorCube:        attrs[f] = 0; break;
				case ElementType::SimpleSingularInteriorCube: attrs[f] = 1; break;
				case ElementType::MultiSingularInteriorCube:  attrs[f] = 2; break;
				case ElementType::RegularBoundaryCube:        attrs[f] = 3; break;
				case ElementType::SimpleSingularBoundaryCube: attrs[f] = 4; break;
				case ElementType::MultiSingularBoundaryCube:  attrs[f] = 5; break;
				case ElementType::InteriorPolytope:           attrs[f] = 6; break;
				case ElementType::BoundaryPolytope:           attrs[f] = 7; break;
				case ElementType::Undefined:                  attrs[f] = 8; break;
				default: attrs[f] = -1; break;
				}
				// std::cout << attrs[f] << std::endl;
			}
		}

		void compute_visibility_kernel() {
			index_t f = idx_.face;
			if (f < mesh_.facets.nb()) {
				Eigen::MatrixXd IV(mesh_.facets.nb_vertices(f), 2);
				for (index_t lv = 0; lv < mesh_.facets.nb_vertices(f); ++lv) {
					index_t v = mesh_.facets.vertex(f, lv);
					vec3 p = mesh_vertex(mesh_, v);
					// std::cout << p << std::endl;
					IV.row(lv) << p[0], p[1];
				}
				Eigen::MatrixXd OV;
				polyfem::compute_visibility_kernel(IV, OV);
				mesh_.clear();
				GEO::vector<index_t> poly;
				// std::cout << OV << std::endl;
				for (index_t v = 0; v < OV.rows(); ++v) {
					index_t vv = mesh_.vertices.create_vertex();
					assert(v == vv);
					float *p = mesh_.vertices.single_precision_point_ptr(v);
					p[0] = (float) OV(v, 0);
					p[1] = (float) OV(v, 1);
					poly.push_back(v);
					// std::cout << v << ' ' << p[0] << ' ' << p[1] << std::endl;
				}
				mesh_.facets.create_polygon(poly);
			}
		}

		virtual void draw_viewer_properties() override {
			ImGui::Text("Vertex: %d", idx_.vertex);
			ImGui::Text("Edge:   %d", idx_.edge);
			ImGui::Text("Facet:  %d", idx_.face);
			idx_.vertex = std::max(0, std::min((int) mesh_.vertices.nb(), idx_.vertex));
			idx_.edge = std::max(0, std::min((int) mesh_.edges.nb(), idx_.edge));
			idx_.face = std::max(0, std::min((int) mesh_.facets.nb(), idx_.face));
			GEO::Attribute<GEO::index_t> c2e(mesh_.facet_corners.attributes(), "edge_id");

			ImGui::Checkbox("Show Index", &show_index_);
			if (ImGui::Button("Switch Vertex", ImVec2(-1, 0))) {
				idx_ = Navigation::switch_vertex(mesh_, idx_);
			}
			if (ImGui::Button("Switch Edge", ImVec2(-1, 0))) {
				idx_ = Navigation::switch_edge(mesh_, c2e, idx_);
			}
			if (ImGui::Button("Switch Face", ImVec2(-1, 0))) {
				auto tmp = Navigation::switch_face(mesh_, c2e, idx_);
				if (tmp.face != -1) {
					idx_ = tmp;
				}
			}

			ImGui::Separator();

			if (ImGui::Button("Kill Singularities", ImVec2(-1, 0))) {
				polyfem::create_patch_around_singularities(mesh_, singular_vertices_, singular_edges_);
				Navigation::prepare_mesh(mesh_);
				polyfem::singularity_graph(mesh_, singular_vertices_, singular_edges_);
				compute_types();
				GEO::mesh_save(mesh_, "foo.obj");
			}

			ImGui::DragFloat("t", &t_, 0.01f, 0.0f, 1.0f);
			static bool split_polygons = false;
			ImGui::Checkbox("Split Polygons", &split_polygons);
			static bool force_quad_ring = false;
			ImGui::Checkbox("Force quad ring", &force_quad_ring);
			if (ImGui::Button("Refine", ImVec2(-1, 0))) {
				GEO::Mesh tmp;
				if (split_polygons == false) {
					polyfem::refine_polygonal_mesh(mesh_, tmp, polyfem::Polygons::no_split_func());
				} else if (force_quad_ring) {
					polyfem::refine_polygonal_mesh(mesh_, tmp, polyfem::Polygons::catmul_clark_split_func());
				} else {
					polyfem::refine_polygonal_mesh(mesh_, tmp, polyfem::Polygons::polar_split_func(t_));
				}
				mesh_.copy(tmp);
				// std::cout << mesh_.vertices.nb() << std::endl;
				Navigation::prepare_mesh(mesh_);
				polyfem::singularity_graph(mesh_, singular_vertices_, singular_edges_);
				compute_types();
				idx_ = Navigation::get_index_from_face(mesh_, c2e, 0, 0);
				// GEO::mesh_save(mesh_, "foo.obj");
				mesh_gfx_.set_mesh(&mesh_);
			}

			if (ImGui::Button("Visibility Kernel", ImVec2(-1, 0))) {
				compute_visibility_kernel();
				Navigation::prepare_mesh(mesh_);
				polyfem::singularity_graph(mesh_, singular_vertices_, singular_edges_);
				compute_types();
				idx_ = Navigation::get_index_from_face(mesh_, c2e, 0, 0);
			}
		}

		virtual void draw_scene() override {
			if (mesh()) {
				if (show_index_) {
					draw_selected();
				}
				//draw_singular();
			}
			SimpleMeshApplication::draw_scene();
		}

		virtual void draw_selected() {
			glupSetPointSize(GLfloat(10));
			glupEnable(GLUP_VERTEX_COLORS);
			glupColor3f(1.0f, 0.0f, 0.0f);

			// Selected vertex
			glupBegin(GLUP_POINTS);
			glupVertex(mesh_vertex(mesh_, idx_.vertex));
			glupEnd();

			// Selected edge
			glupSetMeshWidth(5);
			glupColor3f(0.0f, 0.8f, 0.0f);
			glupBegin(GLUP_LINES);
			{
				int v0 = mesh_.edges.vertex(idx_.edge, 0);
				int v1 = mesh_.edges.vertex(idx_.edge, 1);
				glupVertex(mesh_vertex(mesh_, v0));
				glupVertex(mesh_vertex(mesh_, v1));
			}
			// Boundary edges
			GEO::Attribute<bool> boundary(mesh_.edges.attributes(), "boundary_edge");
			glupColor3f(0.7f, 0.7f, 0.0f);
			for (int e = 0; e < (int) mesh_.edges.nb(); ++e) {
				if (boundary[e]) {
					int v0 = mesh_.edges.vertex(e, 0);
					int v1 = mesh_.edges.vertex(e, 1);
					// glupVertex(mesh_vertex(mesh_, v0));
					// glupVertex(mesh_vertex(mesh_, v1));
				}
			}
			glupEnd();

			// Selected facet
			glupSetMeshWidth(0);
			glupColor3f(1.0f, 0.0f, 0.0f);
			glupBegin(GLUP_TRIANGLES);
			for (int lv = 1; lv + 1 < (int) mesh_.facets.nb_vertices(idx_.face); ++lv) {
				int v0 = mesh_.facets.vertex(idx_.face, 0);
				int v1 = mesh_.facets.vertex(idx_.face, lv);
				int v2 = mesh_.facets.vertex(idx_.face, lv+1);
				glupVertex(mesh_vertex(mesh_, v0));
				glupVertex(mesh_vertex(mesh_, v1));
				glupVertex(mesh_vertex(mesh_, v2));
			}
			glupEnd();

			glupDisable(GLUP_VERTEX_COLORS);
		}

		void draw_singular() {
			glupSetPointSize(GLfloat(10));
			glupEnable(GLUP_VERTEX_COLORS);
			glupColor3f(0.0f, 0.0f, 0.8f);

			// Selected vertex
			glupBegin(GLUP_POINTS);
			for (int i = 0; i < singular_vertices_.size(); ++i) {
				glupVertex(mesh_vertex(mesh_, singular_vertices_[i]));
			}
			glupEnd();

			// Selected edge
			glupSetMeshWidth(5);
			glupBegin(GLUP_LINES);
			for (int e = 0; e < singular_edges_.rows(); ++e) {
				int v0 = singular_edges_(e, 0);
				int v1 = singular_edges_(e, 1);
				glupVertex(mesh_vertex(mesh_, v0));
				glupVertex(mesh_vertex(mesh_, v1));
			}
			glupEnd();
			glupDisable(GLUP_VERTEX_COLORS);
		}

	};
}

int main(int argc, char** argv) {
	#ifndef WIN32
	setenv("GEO_NO_SIGNAL_HANDLER", "1", 1);
	#endif
	GEO::initialize();
	GEO::CmdLine::import_arg_group("standard");
	GEO::CmdLine::import_arg_group("algo");
	GEO::CmdLine::import_arg_group("gfx");
	DemoGlupApplication app(argc, argv);
	GEO::CmdLine::set_arg("gfx:geometry", "1024x1024");
	app.start();
	return 0;
}


//
// Created by Alex Beccaro on 23/01/2019.
//

#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#include "../../src/problems/101-150/126/problem126.hpp"

BOOST_AUTO_TEST_SUITE( Problem126 )

    BOOST_AUTO_TEST_CASE( Example ) {
        auto res = problems::problem126::solve(10);
        BOOST_CHECK_EQUAL(res, 154);
    }

    BOOST_AUTO_TEST_CASE( Solution ) {
        auto res = problems::problem126::solve();
        BOOST_CHECK_EQUAL(res, 18522);
    }

BOOST_AUTO_TEST_SUITE_END()


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#include <iostream>
#include <stdint.h>
#include <vector>
#include <Eigen/Dense>

#include <boost/shared_ptr.hpp>

#include <dpMM/dpMM.hpp>
#include <dpMM/cat.hpp>
#include <dpMM/dir.hpp>
#include <dpMM/niw.hpp>
#include <dpMM/sampler.hpp>
#include <dpMM/clGMMData.hpp>
#include <dpMM/niwSphere.hpp>

using namespace Eigen;
using std::endl; using std::cout;
using boost::shared_ptr;

/* Dirichlet Mixture using ClGMMData for data storage - should be FAST */
template <class H, class T>
class DirMMcld : public DpMM<T>
{
  uint32_t K_;
  Dir<Cat<T>,T> dir_;
  Cat<T> pi_;
#ifdef CUDA
  SamplerGpu<T>* sampler_;
#else
  Sampler<T>* sampler_;
#endif
  Matrix<T,Dynamic,Dynamic> pdfs_;
//  Cat cat_;
  shared_ptr<BaseMeasure<T> > theta0_;
  vector<shared_ptr<BaseMeasure<T> > > thetas_;

  shared_ptr<ClGMMData<T> > cld_;

public:
  DirMMcld(const Dir<Cat<T>,T>& alpha, const shared_ptr<BaseMeasure<T> >& theta);
  DirMMcld(const Dir<Cat<T>,T>& alpha, const vector<shared_ptr<BaseMeasure<T> > >& thetas);
  ~DirMMcld();

  void initialize(const shared_ptr<ClGMMData<T> >& cld);
  void initialize(const Matrix<T,Dynamic,Dynamic>& x)
    {cout<<"not supported"<<endl; assert(false);};

  void sampleLabels();
  void sampleParameters();

  double logJoint();
  const VectorXu& z(){return (cld_->z());};
  const VectorXu& labels(){return (cld_->z());};
  const VectorXu& getLabels(){return (cld_->z());};
  const Matrix<T,Dynamic,1>& counts(){return cld_->counts();};
  const Matrix<T,Dynamic,Dynamic>& means(){return cld_->means();};

  virtual uint32_t getK() const { return K_;};

  Matrix<T,Dynamic,1> getCounts();

private:
};

// --------------------------------------- impl -------------------------------

template <class H, class T>
DirMMcld<H,T>::DirMMcld(const Dir<Cat<T>,T>& alpha,
    const shared_ptr<BaseMeasure<T> >& theta)
  : K_(alpha.K_), dir_(alpha), pi_(dir_.sample()), theta0_(theta)
{};

template <class H, class T>
DirMMcld<H,T>::DirMMcld(const Dir<Cat<T>,T>& alpha,
    const vector<shared_ptr<BaseMeasure<T> > >& thetas) :
  K_(alpha.K_), dir_(alpha), pi_(dir_.sample()), //cat_(dir_.sample()),
  thetas_(thetas)
{};

template <class H, class T>
DirMMcld<H,T>::~DirMMcld()
{
  if (sampler_ != NULL) delete sampler_;
};

template <class H, class T>
Matrix<T,Dynamic,1> DirMMcld<H,T>::getCounts()
{
  return counts();
};

//template <class H, class T>
//void DirMMcld<H,T>::initialize(const Matrix<T,Dynamic,Dynamic>& x)
//{
//
//};

template <class H, class T>
void DirMMcld<H,T>::initialize(const shared_ptr<ClGMMData<T> >& cld)
{
  cld_ = cld;
  assert(cld_->K() == K_);

  cout<<"init"<<endl;
  // randomly init labels from prior
  cout<<"sample pi"<<endl;
  pi_ = dir_.sample();
  cout<<"init pi="<<pi_.pdf().transpose()<<endl;
#ifdef CUDA
  sampler_ = new SamplerGpu<T>(cld_->N(),K_,dir_.pRndGen_);
#else
  sampler_ = new Sampler<T>(dir_.pRndGen_);
#endif
  //TODO: use sampler for this
  //pi_.sample(*(cld_->z()));
  pdfs_.setZero(cld_->N(),K_);
  for(uint32_t i=0; i<cld_->N(); ++i)
    pdfs_.row(i) = pi_.pdf();
  sampler_->sampleDiscPdf(pdfs_,(cld_->z()));
//  cout<<z->transpose()<<endl;
  assert((cld->z().array() < K_).all());
  // init the parameters
  if(thetas_.size() == 0)
  {
    cout<<"creating thetas"<<endl;
    for (uint32_t k=0; k<K_; ++k)
      thetas_.push_back(shared_ptr<BaseMeasure<T> >(theta0_->copy()));
  }
//  cld_->update(K_);
//  for(uint32_t k=0; k<K_; ++k)
//    thetas_[k]->posterior(cld_,k);
//  for (uint32_t k=0; k<K_; ++k)
//    thetas_[k].initialize(x_,z_);
};

template <class H, class T>
void DirMMcld<H,T>::sampleLabels()
{
  // obtain posterior categorical under labels
  pi_ = dir_.posterior(*(cld_->z())).sample();
//  cout<<pi_.pdf().transpose()<<endl;

  for(uint32_t i=0; i<cld_->N(); ++i)
  {
    //TODO: could buffer this better
    // compute categorical distribution over label z_i
    VectorXd logPdf_z = pi_.pdf().array().log();
    for(uint32_t k=0; k<K_; ++k)
    {
//      cout<<thetas_[k].logLikelihood(cld_->x()->col(i))<<" ";
      // TODO this does waste time since we are doing Log_p for all x on CPU
      logPdf_z[k] += thetas_[k]->logLikelihood(*(cld_->x()),i);
    }
//    cout<<endl;
    // make pdf sum to 1. and exponentiate
    pdfs_.row(i) = (logPdf_z.array()-logSumExp(logPdf_z)).exp().matrix().transpose();
//    cout<<pi_.pdf().transpose()<<endl;
//    cout<<pdf.transpose()<<" |.|="<<pdf.sum();
//    cout<<" z_i="<<z_[i]<<endl;
  }
  // sample z_i
  sampler_->sampleDiscPdf(pdfs_,cld_->z());
//    cout<<pdfs_<<endl;
//  cout<<" z="<<(*(cld_->z())).transpose()<<endl;
};

// specialization of the sampling function to run on GPU
template<>
void DirMMcld<NiwSphere<float>,float>::sampleLabels()
{
//  cout<<"sampling specialized to DirMMcld<NiwSphere,float>"<<endl;
  // obtain posterior categorical under labels
//  pi_ = dir_.posterior(*(cld_->z())).sample();
  pi_ = dir_.posteriorFromCounts(cld_->counts()).sample();
//  cout<<pi_.pdf().transpose()<<endl;
  vector<Matrix<float,Dynamic,Dynamic> > Sigmas(K_,
      Matrix<float,Dynamic,Dynamic>::Zero(cld_->D(),cld_->D()));
  Matrix<float,Dynamic,1> logNormalizers(K_);
  for(uint32_t k=0; k<K_; ++k)
  {
    Sigmas[k] = dynamic_cast<NiwSphere<float>* >(
        thetas_[k].get())->normalS_.Sigma();
    logNormalizers(k) = -0.5* dynamic_cast<NiwSphere<float>* >(
        thetas_[k].get())->normalS_.logDetSigma();
  }
//  cout<<logNormalizers.transpose()<<endl;
  cld_->sampleGMMpdf(pi_.pdf(), Sigmas, logNormalizers, sampler_);
};

template<>
void DirMMcld<NiwSphere<double>,double>::sampleLabels()
{
//  cout<<"sampling specialized to DirMMcld<NiwSphere,double>"<<endl;
  // obtain posterior categorical under labels
//  pi_ = dir_.posterior(*(cld_->z())).sample();
  pi_ = dir_.posteriorFromCounts(cld_->counts()).sample();
//  cout<<pi_.pdf().transpose()<<endl;
  vector<Matrix<double,Dynamic,Dynamic> > Sigmas(K_,
      Matrix<double,Dynamic,Dynamic>::Zero(cld_->D(),cld_->D()));
  Matrix<double,Dynamic,1> logNormalizers(K_);
  for(uint32_t k=0; k<K_; ++k)
  {
    Sigmas[k] = dynamic_cast<NiwSphere<double>* >(
        thetas_[k].get())->normalS_.Sigma();
    logNormalizers(k) = -0.5* dynamic_cast<NiwSphere<double>* >(
        thetas_[k].get())->normalS_.logDetSigma();
  }
//  cout<<logNormalizers.transpose()<<endl;

  cld_->sampleGMMpdf(pi_.pdf(), Sigmas, logNormalizers, sampler_);
};

template <class H, class T>
void DirMMcld<H,T>::sampleParameters()
{
  cld_->update(K_);
  for(uint32_t k=0; k<K_; ++k)
    thetas_[k]->posterior(cld_,k);
};

template <class H, class T>
double DirMMcld<H,T>::logJoint()
{
  double logJoint = dir_.logPdf(pi_);
  for (uint32_t k=0; k<K_; ++k)
    logJoint += thetas_[k]->logPdfUnderPrior();
  for (uint32_t i=0; i<cld_->N(); ++i)
    logJoint += thetas_[(cld_->z())(i)]->logLikelihood(*(cld_->x()),i);
  return logJoint;
};


/*
- Etapas por geração:
1. Inicializar população
2. Avaliação
3. Seleção
4. Crossover
5. Mutação
6. Rearranjo da população

// g++ -o main main.cpp -std=c++11  -O2 -larmadillo
[-lpthread -ldl -lglut -lGLU -lGL (-lm)] -> not used for now
// Warning flags: -Wall -Wextra -Werror

-> Atualmente existem as finais e as semifinais entre AGs (AEs). A final é composta de TOURNMENT_RATE semi-finais e cada semi-final é composta de TOURNMENT_RATE evolutionaryAlgorithms.
-> apos a ocorrencia de TOURNMENT_RATE EAS em uma semi-final, uma outra EA é feita, onde TOURNMENT_RATE dos individuos são os melhores individuos dessas EAs anteriores.
-> Na final é feito algo análogo à semi-final, só que ao invés dos melhores indivíduos virem das EAs, eles vem das semi-finais.

-> O que fazer agora?
1. Fazer a função de cálculo do fitness (em evaluatePop) (fitness[i] = abs(terminoReal - terminoEsperado)/terminoEsperado)
[]
2. Fazer o AGs gerar o arquivo de configuração com os parâmetros do jogo (em initializePop)
[OK]
3. Checar se a variância é pequena para aplicar alteração mutação. Além disso, trocar em "if(nbNoImprovementGens > 200)", o 200 por um parâmetro (usar #define)
[NÂO NECESSARIO/A DISCUTIR]
4. Salvar o melhor de cada AG em um vetor de "melhores globais" e trocar o while(true) da main por outra condição (o AG deve finalizar em algum momento - salvando seu melhor no vetor de "melhores globais")
[PARCIAL]
5. Fazer nas semi-finais e finais initializePop com valores não puramente randomicos e sim com valores limitados entre o min e o max dos valores já descobertos para cada parametro
[OK]

*/
#include <stdlib.h>  /* srand, rand */
#include <iostream>
#include <fstream> // creating (.csv) files
#include <vector>
#include <cstring>
#include <string>
#include <functional>
#include <armadillo> // http://arma.sourceforge.net/docs.html
// #include <unistd.h>
// #include <thread>   /* std::this_thread::sleep_for */
// #include <chrono>   /* std::chrono::seconds */
// #include <GL/glut.h> /* simple GUI + keyboard usage */
// #include <loguru.cpp> // logging lib -> https://emilk.github.io/loguru/index.html#loguru

const std::string parametersList[] = {
    "ON_ATTACK_MULTIPLIER",
    "POP_INCREASE_TAX",
    "POP_PER_CONSTRUCTION",
    "INIT_POP_VALUE",
    "INIT_STAT_VALUE",
    "ON_ATTACK_DECAY_TAX",
    "NORMAL_DECAY_TAX",
    "ATTACK_FREQUENCY",
    "INIT_RESOURCES_VALUE",
    "TAX_REDUCT",
    "DECAY_DELAY_MICRO",
    "MIN_LOSS_0",
    "MIN_LOSS_1",
    "MIN_LOSS_2",
    "MIN_LOSS_3",
    "MAX_LOSS_0",
    "MAX_LOSS_1",
    "MAX_LOSS_2",
    "MAX_LOSS_3",
    "TOT_ROBOTS_INI",
    "FREE_ROBOTS_INI",
    "PROD_COST_INI",
    "TIME_STEP",
    "INIT_TIME_STEP",
    "MAX_TIME_STEPS",
    "MIN_REWARD",
    "REWARD_RANGE",
    "FAILURE_TAX"
};

#define NB_PARAMETERS 28   // number of parameters that will be adjusted
#define POPULATION_SIZE 20

#define TOURNMENT_RATE 10

#define MAX_MUTATION_RATE 10
#define MUTATION_INCREASE_RATIO 1.1 // mutation increase 10% at a time
#define INITIAL_MUTATION 0.05

#define BEGIN_DECAY_THRESHOLD 50 // after how many generations without improvement the mutation starts to decay
#define APPLY_DECAY_MODULE 15 // decay will be applied when the generation index % APPLY_DECAY_MODULE = 0

#define APPLY_PREDATION_INTERVAL 10 // interval (in generations) to apply predation

#define MAX_PARAM_VALUE 1 // max value a parameter can reach

// NOT USED??
// enum events{
//     BASE_RATE = 0,
//     MUTA_RATE,
//     PRED_RATE,
//     PART_RATE,
//     GENO_RATE,
//     EV_A_RATE,
// };

enum tournments{ INITIALS,SEMI_FINALS, FINALS};

arma::mat::fixed<TOURNMENT_RATE,NB_PARAMETERS> indvSemiFinals; // arma::mat : armadillo matrix (Mat<double>)
arma::mat::fixed<TOURNMENT_RATE,NB_PARAMETERS> indvFinals; // arma::mat : armadillo matrix (Mat<double>)

arma::mat::fixed<TOURNMENT_RATE, NB_PARAMETERS> tournments[2] = {indvSemiFinals,indvFinals};

//arma::rowvec indvSemiFinals[TOURNMENT_RATE];
//arma::rowvec indvFinals[TOURNMENT_RATE];

const int eventTriggerModuleBase[] = {1, 10, 5, 10, 2, 10};
int eventTriggerModule[6];

enum selectionMethods { TOURNAMENT, ELITISM, ROULETTE };
selectionMethods partsSelectionMethods[] = {
    TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,TOURNAMENT,
    ELITISM,ELITISM,ELITISM,ELITISM,ELITISM,ELITISM,ELITISM,ELITISM,ELITISM,ELITISM
};
int partPos = 0;
selectionMethods selectionMethod = partsSelectionMethods[partPos];

arma::mat::fixed<POPULATION_SIZE, NB_PARAMETERS> population; // arma::mat : armadillo matrix (Mat<double>)
// Each LINE (ROW) is the individual's parameters
double fitness[POPULATION_SIZE];
double maxFitness = 0, minFitness = 2000000000;
int maxFitIndex = -1, minFitIndex = -1;
double mutationRate = INITIAL_MUTATION;
int nbNoImprovementGens = 0; // number of generations in a row without any improvement (best indiviNdual remains the same)
int generationIndex = 0;
bool continueEA;

// Choose if population evaluation will be according to max fitness value or min fitness
enum evalTypes{ BY_MIN, BY_MAX };
evalTypes evalType = BY_MIN;
bool greater(double par1,double par2) { return par1 > par2; }
bool less(double par1,double par2) { return par1 < par2; }
bool (*evaluateType[])(double,double) = {less,greater};

// 1st step: initialize population
void initializePop(int tournmentType) {
    double minVal,maxVal;
    std::cout << "INITIALIZING POPULATION\n";
    // for (int i = 0; i < POPULATION_SIZE; i++) {
    //     for (int j = 0; j < NB_PARAMETERS; j++) {
    //         population[i][j] = (double) (rand() % MAX_PARAM_VALUE); // number range = [0, MAX_PARAM_VALUE[
    //     }
    // }
    population.randu(); // initialize with values between 0 and 1
    if(!tournmentType) population *= MAX_PARAM_VALUE; // if a normal EA is taking place, just multiply the population by a givern value
    else{ // if not
        int middle = (POPULATION_SIZE - TOURNMENT_RATE)/2;
        for(int i = TOURNMENT_RATE + 1;i< middle;i++){ // Min and max values are defined for each col, and converting the current col values to a number between this constraints.
            // this starts after the first TOURNMENT_RATE individuls, and the remaining is divided into two pars, one with biased rand values, and other without
            minVal = 0.5 * tournments[tournmentType].col(i).min();
            maxVal = 1.5 * tournments[tournmentType].col(i).max();
            population.col(i) = minVal + (maxVal - minVal)*(population.col(i));
        }
        for(int i = middle + 1;i < POPULATION_SIZE;i++)
            population.col(i) *= MAX_PARAM_VALUE;
    }

    std::fstream myFile;
    myFile.open("constValues.cfg", std::ios_base::app);

    myFile << parametersList[0] << " = " << population(0,0);
    for (size_t i = 1; i < NB_PARAMETERS; i++) myFile << std::endl << parametersList[i] << " = " << population(0,i);
    //printf("k%sk\n", parametersList[i]);
    myFile.close();

    population.print("Population matrix initialized:");

    return;
}

// Function to normalize fitness (values between 0 and 1)
void normalizeFitness(double * normalizedFitness) {
    if (normalizedFitness == NULL)
        return;

    for (int i = 0; i < POPULATION_SIZE; i++)
        normalizedFitness[i] = (fitness[i]-minFitness)/(maxFitness-minFitness);

    return;
}

// 2nd step: evaluate population (calculate fitness)
void evaluatePop() {
    std::cout << "EVALUATING POPULATION\n";

    nbNoImprovementGens++; // we begin considering there was no improvement in the generation

    for (int i = 0; i < POPULATION_SIZE; i++) {
        // example of fitness calculation -> CHANGE!!!!
        // fitness[i] = abs(terminoReal - terminoEsperado)/terminoEsperado
        fitness[i] = arma::sum(population.row(i));

        if ((*evaluateType[evalType])(maxFitness,fitness[i])){ // searching for the  max fitnnes from new generation
            maxFitness = fitness[i];
            maxFitIndex = i;
            nbNoImprovementGens = 0; // this generation shows improvement -> reset counter
        }
        else if ((*evaluateType[evalType])(fitness[i],minFitness)){
            minFitness = fitness[i];
            minFitIndex = i;
        }

        // printf("fitness[%d] %lf\n", i, fitness[i]);
    }

    printf("MAX FITNESS: %lf - INDEX: %d\n", maxFitness, maxFitIndex);
    printf("MIN FITNESS: %lf - INDEX: %d\n", minFitness, minFitIndex);

    return;
}

void crossover(int i, arma::rowvec parent1, arma::rowvec parent2) {
    // for(int j = 0; j < NB_PARAMETERS; j++){
    //     population[i][j] = (parent1[j] + parent2[j])/2.0;
    // }
    population.row(i) = (parent1 + parent2)/2.0;

    return;
}

void elitism() {
    std::cout << "ELITISM\n";
    arma::rowvec bestIndv = population.row(maxFitIndex);

    for (int i = 0; i < POPULATION_SIZE; i++) {
        if (i == maxFitIndex) // keeping the best individual (no changes are made to it)
            continue;

        // crossover
        crossover(i, population.row(i), bestIndv);
    }

    return;
}

void tournament() {
    std::cout << "TOURNAMENT\n";
    arma::mat::fixed<POPULATION_SIZE, NB_PARAMETERS> oldPopulation;
    int parentIndex[2];

    // copying last population (new one will be different)
    // for (int i = 0; i < POPULATION_SIZE; i++) {
    //     for (int j = 0; j < NB_PARAMETERS; j++) {
    //         oldPopulation[i][j] = population[i][j];
    //     }
    // }
    oldPopulation = population;

    for (int i = 0; i < POPULATION_SIZE; i++) {
        if (i == maxFitIndex)
            continue;

        // chossing parents for new individual
        for (int j = 0; j < 2; j++) {
            int indexIndA = rand() % POPULATION_SIZE; // indv 1 that will "fight" to be parent
            int indexIndB = rand() % POPULATION_SIZE; // indv 2 that will "fight" to be parent

            parentIndex[j] = (fitness[indexIndA] > fitness[indexIndB]) ? indexIndA : indexIndB;
        }

        // crossover
        crossover(i, oldPopulation.row(parentIndex[0]), oldPopulation.row(parentIndex[1]));
    }

    return;
}

void roulette() {
    std::cout << "ROULETTE\n";
    arma::mat::fixed<POPULATION_SIZE, NB_PARAMETERS> oldPopulation;
    double standardizedFitness[POPULATION_SIZE];
    int parentIndex[2], rNb;
    double probSum = 0.0, partialSum = 0.0;

    // copying last population (new one will be different)
    // for (int i = 0; i < POPULATION_SIZE; i++) {
    //     for (int j = 0; j < NB_PARAMETERS; j++) {
    //         oldPopulation[i][j] = population[i][j];
    //     }
    // }
    oldPopulation = population;

    // Standardize fitness (set probabilites that add up to 100%)
    for (int i = 0; i < POPULATION_SIZE; i++)
        probSum += fitness[i];
    for (int i = 0; i < POPULATION_SIZE; i++)
        standardizedFitness[i] = fitness[i]/probSum;

    // Chosing new parents for each individual
    for (int i = 0; i < POPULATION_SIZE; i++) {
        if (i == maxFitIndex) // preserves best individual
            continue;

        for (int k = 0; k < 2; k++) { // chosing 2 parents
            rNb = ((double) rand() / RAND_MAX); // rand between 0 and 1
            partialSum = 0.0;

            for (int j = 0; j < POPULATION_SIZE; j++) { // randomly chosing and individual according to its fitness (+fitness = +probabity)
                partialSum += fitness[j];
                if (partialSum >= rNb) {
                    parentIndex[k] = j; // new parent at index j
                    break;
                }
            }
        }

        // crossover
        crossover(i, oldPopulation.row(parentIndex[0]), oldPopulation.row(parentIndex[1]));
    }
}

void (*selectionType[])() = { tournament, elitism, roulette };

void mutate(int indIndex) {
    int plusMinusSign = (rand() % 2 ? 1 : -1); // mutation will increase the param value or decrease it
    int mutateParamIndex = rand() % NB_PARAMETERS; // index of parameter that will be mutated

    population(indIndex, mutateParamIndex) += population(indIndex, mutateParamIndex) * mutationRate * plusMinusSign;

    return;
}

// 3rd step: selection + mutation and crossover
void selectionAndMutation() { // tournament, elitism, roulette...
    std::cout << "SELECTION\n";

    // Selection method (+ crossover)
    (*selectionType[selectionMethod])(); // void (*selectionType[])() => tournament(), elitism() or roulette()

    // Mutation

    // Mutating all params from individuals
    // arma::rowvec bestIndv = population.row(maxFitIndex);

    // arma::mat mutateMatrix(POPULATION_SIZE, NB_PARAMETERS, arma::fill::randu);
    // mutateMatrix = ((mutateMatrix * MAX_PARAM_VALUE) - MAX_PARAM_VALUE/2.0) * mutationRate;

    // population = population + mutateMatrix;

    // population.row(maxFitIndex) = bestIndv;
    // population.transform( [](double x) { return ((x < 0 || x > MAX_PARAM_VALUE) ? abs(MAX_PARAM_VALUE-abs(x)) : x); } );

    // Mutating only one parameter from each individual
    for(int i = 0; i < maxFitIndex; i++) mutate(i);
    // (don't mutate best one)
    for(int i = maxFitIndex+1; i < POPULATION_SIZE; i++) mutate(i);

    return;
}

// Initialize .csv file to sabe data from the EA
// Creates the column's headers
// generation, paramsIndv1, fitnessIndv1 ..., paramsIndvN, fitnessIndvN, bestParams, bestFitness
void createCSV(std::string pathPos) {
    std::ofstream csvFileWriter;

    csvFileWriter.open("historyEA-" + pathPos + ".csv");
    if (!csvFileWriter.good()) {
        std::cout << "[!] Error occurred while trying to create historyEA" << pathPos << ".csv!\n";
        return;
    }

    csvFileWriter << "generation,";
    for (int i = 0; i < POPULATION_SIZE; i++)
        csvFileWriter << "paramsIndv" << i << ",fitnessIndv" << i << ",";
    csvFileWriter << "paramsBestIndv,fitnessBestIndv\n";

    return;
}

std::string formatParamsString(int indvIndex) {
    std::string paramsFormated = "[ ";
    for (int i = 0; i < NB_PARAMETERS; i++) {
        paramsFormated += std::to_string(population(indvIndex, i));
        paramsFormated += " ";
    }
    paramsFormated += "]";

    return paramsFormated;
}

// Saves information about a generation in a .csv file
// Information: generation, paramsIndv1, fitnessIndv1 ..., paramsIndvN, fitnessIndvN, bestParams, bestFitness
void saveGenerationData(int generation) {
    std::ofstream csvFileWriter;

    csvFileWriter.open("historyEA.csv", std::ios_base::app); // append instead of overwrite
    if (!csvFileWriter.good()) {
        std::cout << "[!] Error occurred while trying to open historyEA.csv!\n";
        return;
    }

    csvFileWriter << generation << ",";
    for (int i = 0; i < POPULATION_SIZE; i++) {
        csvFileWriter << formatParamsString(i) << "," << fitness[i] << ",";
    }
    csvFileWriter << formatParamsString(maxFitIndex) << "," << maxFitness << "\n";
}

void increaseMutation() {
    if(!mutationRate) {
        mutationRate = INITIAL_MUTATION;
    }
    else {
        mutationRate *= MUTATION_INCREASE_RATIO;

        if (mutationRate > MAX_MUTATION_RATE) // mutation reached its max value
            mutationRate = MAX_MUTATION_RATE;
    }

    return;
}

// Function that will kill the worst individual each "APPLY_PREDATION_INTERVAL" number of generations
void predationOfOne() {
    arma::rowvec newInd(NB_PARAMETERS, arma::fill::randu); // creating totally new individual
    newInd = newInd * MAX_PARAM_VALUE;

    population.row(minFitIndex) = newInd;

    return;
}

// A sequence of selections will run, this method will indicate the selection method that will be used by the next batch of generations
void partIncrease(){
    selectionMethod = partsSelectionMethods[++partPos];

    return;
}

// Function that will kill all individuals but the best to reset the population and generate new individuals (without biases)
void oneRemainingPopReset() {
    std::cout << "APPLYING POPULATION RESET\n";

    arma::rowvec best = population.row(maxFitIndex);

    initializePop(INITIALS);
    population.row(0) = best;
    partIncrease();

    return;
}

// Function that kills all individuals and creates a new population
// Ends the batch of generations from current EA
void fullPopReset() {
    std::cout << "HARD POPULATION RESET";

    continueEA = false;

    return;
}

// Generic function that checks if certain event should happen (predation, population reset, mutation increase...)
bool eventHappens(int eventType) {
    if(nbNoImprovementGens % eventTriggerModule[eventType] == 0)
        return true; // the event being verified should happen!
    return false;
}

void (*eventTypes[])() = { increaseMutation, predationOfOne, partIncrease, oneRemainingPopReset, fullPopReset };

void checkEvents() {
    if(nbNoImprovementGens == 0)
        return;

    for(int i = 1; i < 6; i++){
        if(eventHappens(i)) {
            (*eventTypes[i-1])(); // calls one of these functions: increaseMutation, predationOfOne, partIncrease, oneRemainingPopReset, fullPopReset
            return;
        }
    }
}

void startEventTriggerList(){
    eventTriggerModule[0] = eventTriggerModuleBase[0];
    for(int i = 1; i < 6; i++)
        eventTriggerModule[i] = eventTriggerModule[i-1] * eventTriggerModuleBase[i];
}

void evolutionaryAlgorithm(std::string csvStr){
    nbNoImprovementGens = 0;
    createCSV(csvStr);

    partPos = 0; // defines selection method of current batch of generations

    continueEA = true;
    generationIndex = 0;
    while (continueEA) {
        std::cout << "\n==== Generation " << generationIndex << " ====\n";
        // population.print("Current population:");
        evaluatePop();
        selectionAndMutation();

        saveGenerationData(generationIndex);

        checkEvents(); // checks if mutation should increase, predation or population reset should occur etc.
        // if fullPopReset() is called, continueEA = false

        generationIndex++;
        scanf("%*c");
    }

    return;
}

void replaceByBests(int tournmentType){
    for(int i = 0; i < TOURNMENT_RATE; i++)
        population.row(i) = tournments[tournmentType - 1].row(i);
}

void semiFinalsTournment(int semiFinalPos){
    for(int i = 0; i < TOURNMENT_RATE; i++){
        initializePop(INITIALS);
        evolutionaryAlgorithm("SF-" + std::to_string(semiFinalPos) + "_EA-" + std::to_string(i));
        indvSemiFinals.row(i) = population.row(maxFitIndex); // saves the best individual from current EA
    }

    initializePop(SEMI_FINALS);
    replaceByBests(SEMI_FINALS);
    evolutionaryAlgorithm("SF-" + std::to_string(semiFinalPos)); // this EA will use the best individuals from each previous EA
    indvFinals.row(semiFinalPos) = population.row(maxFitIndex); // saves the best individual from current EA
    return;
}

arma::rowvec finalTournment(){
    for(int i = 0; i < TOURNMENT_RATE; i++)
        semiFinalsTournment(i);

    initializePop(FINALS);
    replaceByBests(FINALS);
    evolutionaryAlgorithm("Main"); // this EA will use the best individuals from each semifinal

    return population.row(maxFitIndex);
}

int main(int argc, char * argv[]) {
    srand(time(NULL));
    startEventTriggerList();

    arma::rowvec bestIndividual = finalTournment();

    system("clear");
}


/* -*-c++-*--------------------------------------------------------------------
 * 2018 Bernd Pfrommer bernd.pfrommer@gmail.com
 */

#include "tagslam/pose_noise.h"
#include <Eigen/Cholesky>

#include <iostream>

namespace tagslam {
  // static function
  PoseNoise PoseNoise::make(const Point3d &a, const Point3d &p) {
    Matrix6d m = Matrix6d::Zero();
    m.diagonal() << a(0)*a(0),a(1)*a(1),a(2)*a(2),p(0)*p(0),p(1)*p(1),
      p(2)*p(2);
    return (PoseNoise(m, true));
  }

  // static function
  PoseNoise PoseNoise::make(double a, double p) {
    return (make(Point3d(a, a, a), Point3d(p, p, p)));
  }

  static const PoseNoise::Matrix6d sqrt_info_to_sigma(
    const PoseNoise::Matrix6d &R) {
    // TODO: test this! Is it working at all???
    const PoseNoise::Matrix6d rsqi = (R.transpose() * R).inverse();
    return (rsqi);
  }

  // static method
  PoseNoise PoseNoise::makeFromR(const Matrix6d &R) {
    const Matrix6d sigma = sqrt_info_to_sigma(R);
    return (PoseNoise(sigma, false /*isdiag*/));
  }

  Eigen::Matrix<double, 6, 1> PoseNoise::getDiagonal() const {
    return (noise_.diagonal());
  }

  const PoseNoise::Matrix6d  PoseNoise::getCovarianceMatrix() const {
    return (noise_);
  }

  PoseNoise::Matrix6d PoseNoise::convertToR() const {
    //Sigma = R * R^T
    // Cholesky decomposition: sigma = L * L.transpose();
    const Matrix6d ni = noise_.inverse();
    Eigen::LLT<Matrix6d> llt(ni);
    const Matrix6d U = llt.matrixU();
    return (U);
  }

  std::ostream &operator<<(std::ostream &os, const PoseNoise &pn) {
    os << "is_diagonal: " << pn.isDiagonal_ << std::endl << pn.noise_;
    return (os);
  }

}  // namespace


#include "ros/ros.h"
#include "geometry_msgs/Vector3.h"
#include "sensor_msgs/Joy.h"

#include "create_controller/ControlMsgs.h"

#include <string>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <algorithm>
#include <cmath>

#define PI 3.141592
using namespace Eigen;
Vector3d position(-5,-5,-5);
double etheta[4]={0.0,0.0,0.0,0.0};

std::string robot1_name;
std::string robot2_name;
std::string robot3_name;
std::string robot4_name;

void Callbackzero(const create_controller::ControlMsgs::ConstPtr& msg)
{
	etheta[0] = fmod((msg->etheta), 2*PI);
	if (etheta[0] > PI) {
		etheta[0] = 2*PI-etheta[0];
	}
}
void Callbackone(const create_controller::ControlMsgs::ConstPtr& msg)
{
	etheta[1] = fmod((msg->etheta), 2*PI);
	if (etheta[1] > PI) {
		etheta[1] = 2*PI-etheta[1];
	}
}
void Callbacktwo(const create_controller::ControlMsgs::ConstPtr& msg)
{
	etheta[2] = fmod((msg->etheta), 2*PI);
	if (etheta[2] > PI) {
		etheta[2] = 2*PI-etheta[2];
	}
}
void Callbackthree(const create_controller::ControlMsgs::ConstPtr& msg)
{
	etheta[3] = fmod((msg->etheta), 2*PI);
	if (etheta[3] > PI) {
		etheta[3] = 2*PI-etheta[3];
	}
}

void position_callback(const geometry_msgs::Vector3::ConstPtr& msg)
{
	position(0) = msg->x;
	position(1) = msg->y;
	position(2) = msg->z;
}

int main(int argc, char **argv)
{
	// ROS Initalization
	ros::init(argc, argv, "falcon_to_joy");

	ros::NodeHandle n;
	ros::NodeHandle private_n("~");

	private_n.param<std::string>("robot1name", robot1_name, "gandhi");
	private_n.param<std::string>("robot2name", robot2_name, "alexander");
	private_n.param<std::string>("robot3name", robot3_name, "suleyman");
	private_n.param<std::string>("robot4name", robot4_name, "hiawatha");

	ros::Subscriber position_sub = n.subscribe("position", 10, position_callback);
	ros::Subscriber subzero = n.subscribe(  robot1_name + "/dia", 10, Callbackzero);
	ros::Subscriber subone =  n.subscribe(  robot2_name + "/dia", 10, Callbackone);
	ros::Subscriber subtwo =  n.subscribe(  robot3_name + "/dia", 10, Callbacktwo);
	ros::Subscriber subthree =n.subscribe(  robot4_name + "/dia", 10, Callbackthree);
	ros::Publisher joy_pub = n.advertise<sensor_msgs::Joy>("joy", 10);
	ros::Publisher force_pub = n.advertise<geometry_msgs::Vector3>("force", 10);

	// ROS loop
	ros::Rate loop_rate(1000); // 1 kHz
	bool uninitialized = true;
	while (ros::ok())
	{
		ros::spinOnce();

		Vector3d joy(position(0)/0.06,  position(1)/0.06, (2.0*(position(2)-0.073))/0.1035 - 1.0);

		Vector3d force;
		double resistance;

		if (joy.norm()<0.25) {
			resistance =1;
		} else {
			double radius = joy.norm() - 0.25;
			double sum_etheta = fabs(etheta[0])+ fabs(etheta[1])+fabs(etheta[2])+fabs(etheta[3]);
			resistance = 1 ;//+6*(radius*sum_etheta);
			printf("%f\n", sum_etheta);
		}

		for (int i = 0; i < 3; i++)
		{
			force(i) = -5*joy(i)*resistance;
		}

		geometry_msgs::Vector3 fmsg;
		fmsg.x = force(0);
		fmsg.y = force(1);
		fmsg.z = force(2);
		force_pub.publish(fmsg);

		if (joy.norm() < 0.25) {
			for (int i = 0; i < 2; ++i)
			{
				joy(i) =0;
			}
		}

		sensor_msgs::Joy jmsg;
		jmsg.axes.resize(3);
		jmsg.axes[0] = -joy(0);//position(0)/0.06; // -1.0 to 1.0
		jmsg.axes[1] = joy(1);//position(1)/0.06;
		jmsg.axes[2] = joy(2);//(2.0*(position(2)-0.073))/0.1035 - 1.0;//32767
		joy_pub.publish(jmsg);

		loop_rate.sleep();
	}

	return 0;
}


#include "lietorch_cpu.h"
#include <Eigen/Dense>

#include <iostream>
#include "common.h"
#include "dispatch.h"

#include "so3.h"
#include "rxso3.h"
#include "se3.h"
#include "sim3.h"

template <typename Group, typename scalar_t>
void exp_forward_kernel(const scalar_t* a_ptr, scalar_t* X_ptr, int batch_size) {
    // exponential map forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Data>(X_ptr + i*Group::N) = Group::Exp(a).data();
        }
    });
}

template <typename Group, typename scalar_t>
void exp_backward_kernel(const scalar_t* grad, const scalar_t* a_ptr, scalar_t* da, int batch_size) {
    // exponential map backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a(a_ptr + i*Group::K);
            Grad dX(grad + i*Group::N);
            Eigen::Map<Grad>(da + i*Group::K) = dX * Group::left_jacobian(a);
        }
    });
}

template <typename Group, typename scalar_t>
void log_forward_kernel(const scalar_t* X_ptr, scalar_t* a_ptr, int batch_size) {
    // logarithm map forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a = Group(X_ptr + i*Group::N).Log();
            Eigen::Map<Tangent>(a_ptr + i*Group::K) = a;
        }
    });
}

template <typename Group, typename scalar_t>
void log_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
    // logarithm map backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Tangent a = Group(X_ptr + i*Group::N).Log();
            Grad da(grad + i*Group::K);
            Eigen::Map<Grad>(dX + i*Group::N) = da * Group::left_jacobian_inverse(a);
        }
    });
}

template <typename Group, typename scalar_t>
void inv_forward_kernel(const scalar_t* X_ptr, scalar_t* Y_ptr, int batch_size) {
    // group inverse forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Data>(Y_ptr + i*Group::N) = X.inv().data();
        }
    });
}

template <typename Group, typename scalar_t>
void inv_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t *dX, int batch_size) {
    // group inverse backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group Y = Group(X_ptr + i*Group::N).inv();
            Grad dY(grad + i*Group::N);
            Eigen::Map<Grad>(dX + i*Group::N) = -dY * Y.Adj();
        }
    });
}

template <typename Group, typename scalar_t>
void mul_forward_kernel(const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* Z_ptr, int batch_size) {
    // group multiplication forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group Z = Group(X_ptr + i*Group::N) * Group(Y_ptr + i*Group::N);
            Eigen::Map<Data>(Z_ptr + i*Group::N) = Z.data();
        }
    });
}

template <class Group, typename scalar_t>
void mul_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* Y_ptr, scalar_t* dX, scalar_t* dY, int batch_size) {
    // group multiplication backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Grad dZ(grad + i*Group::N);
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Grad>(dX + i*Group::N) = dZ;
            Eigen::Map<Grad>(dY + i*Group::N) = dZ * X.Adj();
        }
    });
}

template <typename Group, typename scalar_t>
void adj_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // adjoint forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = X.Adj(a);
        }
    });
}

template <typename Group, typename scalar_t>
void adj_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Grad db(grad + i*Group::K);

            Tangent a(a_ptr + i*Group::K);
            Tangent b = X.Adj() * a;

            Eigen::Map<Grad>(da + i*Group::K) = db * X.Adj();
            Eigen::Map<Grad>(dX + i*Group::N) = -db * Group::adj(b);
        }
    });
}

template <typename Group, typename scalar_t>
void adjT_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // adjoint forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = X.AdjT(a);
        }
    });
}

template <typename Group, typename scalar_t>
void adjT_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* dX, scalar_t* da, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent db(grad + i*Group::K);
            Grad a(a_ptr + i*Group::K);

            Eigen::Map<Tangent>(da + i*Group::K) = X.Adj(db);
            Eigen::Map<Grad>(dX + i*Group::N) = -a * Group::adj(X.Adj(db));
        }
    });
}

template <typename Group, typename scalar_t>
void act_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
    // action on point forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Point = Eigen::Matrix<scalar_t,3,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*3);
            Eigen::Map<Point>(q_ptr + i*3) = X * p;
        }
    });
}

template <typename Group, typename scalar_t>
void act_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp, int batch_size) {
    // adjoint backward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Point = Eigen::Matrix<scalar_t,3,1>;
    using PointGrad = Eigen::Matrix<scalar_t,1,3>;
    using Transformation = Eigen::Matrix<scalar_t,4,4>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*3);
            PointGrad dq(grad + i*3);

            Eigen::Map<PointGrad>(dp + i*3) = dq * X.Matrix().template block<3,3>(0,0);
            Eigen::Map<Grad>(dX + i*Group::N) = dq * Group::act_jacobian(X*p);
        }
    });
}

// template <typename Group, typename scalar_t>
// void tovec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
//     // group inverse forward kernel
//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;

//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
//         for (int64_t i=start; i<end; i++) {
//             Group X(X_ptr + i*Group::N);
//             Grad g(grad + i*Group::N);
//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
//         }
//     });
// }

// template <typename Group, typename scalar_t>
// void fromvec_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, scalar_t* dX, int batch_size) {
//     // group inverse forward kernel
//     using Data = Eigen::Matrix<scalar_t,Group::N,1>;
//     using Grad = Eigen::Matrix<scalar_t,1,Group::N>;

//     at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
//         for (int64_t i=start; i<end; i++) {
//             Group X(X_ptr + i*Group::N);
//             Grad g(grad + i*Group::N);
//             Eigen::Map<Grad>(dX + i*Group::N) = g * X.vec_jacobian();
//         }
//     });
// }

template <typename Group, typename scalar_t>
void act4_forward_kernel(const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* q_ptr, int batch_size) {
    // action on homogeneous point forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Point = Eigen::Matrix<scalar_t,4,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*4);
            Eigen::Map<Point>(q_ptr + i*4) = X.act4(p);
        }
    });
}

template <typename Group, typename scalar_t>
void act4_backward_kernel(const scalar_t* grad, const scalar_t* X_ptr, const scalar_t* p_ptr, scalar_t* dX, scalar_t* dp, int batch_size) {
    // action on homogeneous point backward kernel

    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Grad = Eigen::Matrix<scalar_t,1,Group::K>;
    using Point = Eigen::Matrix<scalar_t,4,1>;
    using PointGrad = Eigen::Matrix<scalar_t,1,4>;
    using Transformation = Eigen::Matrix<scalar_t,4,4>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Point p(p_ptr + i*4);
            PointGrad dq(grad + i*4);

            Eigen::Map<PointGrad>(dp + i*4) = dq * X.Matrix4x4();
            const Point q = X.act4(p);
            Eigen::Map<Grad>(dX + i*Group::N) = dq * Group::act4_jacobian(q);
        }
    });
}

template <typename Group, typename scalar_t>
void as_matrix_forward_kernel(const scalar_t* X_ptr, scalar_t* T_ptr, int batch_size) {
    // group inverse forward kernel
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;
    using Matrix4 = Eigen::Matrix<scalar_t,4,4,Eigen::RowMajor>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Matrix4>(T_ptr + i*16) = X.Matrix4x4();
        }
    });
}

template <typename Group, typename scalar_t>
void orthogonal_projector_kernel(const scalar_t* X_ptr, scalar_t* P_ptr, int batch_size) {
    // group inverse forward kernel
    using Proj = Eigen::Matrix<scalar_t,Group::N,Group::N,Eigen::RowMajor>;
    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Eigen::Map<Proj>(P_ptr + i*Group::N*Group::N) = X.orthogonal_projector();
        }
    });
}

template <typename Group, typename scalar_t>
void jleft_forward_kernel(const scalar_t* X_ptr, const scalar_t* a_ptr, scalar_t* b_ptr, int batch_size) {
    // left-jacobian inverse action
    using Tangent = Eigen::Matrix<scalar_t,Group::K,1>;
    using Data = Eigen::Matrix<scalar_t,Group::N,1>;

    at::parallel_for(0, batch_size, 1, [&](int64_t start, int64_t end) {
        for (int64_t i=start; i<end; i++) {
            Group X(X_ptr + i*Group::N);
            Tangent a(a_ptr + i*Group::K);
            Tangent b = Group::left_jacobian_inverse(X.Log()) * a;
            Eigen::Map<Tangent>(b_ptr + i*Group::K) = b;
        }
    });
}

// unary operations

torch::Tensor exp_forward_cpu(int group_id, torch::Tensor a) {
    int batch_size = a.size(0);
    torch::Tensor X;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_forward_kernel", ([&] {
        X = torch::zeros({batch_size, group_t::N}, a.options());
        exp_forward_kernel<group_t, scalar_t>(
            a.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            batch_size);
    }));

    return X;
}

std::vector<torch::Tensor> exp_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor a) {
    int batch_size = a.size(0);
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, a.type(), "exp_backward_kernel", ([&] {
        exp_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {da};
}

torch::Tensor log_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor a;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_forward_kernel", ([&] {
        a = torch::zeros({batch_size, group_t::K}, X.options());
        log_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            batch_size);
    }));

    return a;
}

std::vector<torch::Tensor> log_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "log_backward_kernel", ([&] {
        log_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX};
}

torch::Tensor inv_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor Y = torch::zeros_like(X);

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_forward_kernel", ([&] {
        inv_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            batch_size);
    }));

    return Y;
}

std::vector<torch::Tensor> inv_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "inv_backward_kernel", ([&] {
        inv_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX};
}

// binary operations
torch::Tensor mul_forward_cpu(int group_id, torch::Tensor X, torch::Tensor Y) {
    int batch_size = X.size(0);
    torch::Tensor Z = torch::zeros_like(X);

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_forward_kernel", ([&] {
        mul_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            Z.data_ptr<scalar_t>(),
            batch_size);
    }));

    return Z;
}

std::vector<torch::Tensor> mul_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor Y) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dY = torch::zeros(Y.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "mul_backward_kernel", ([&] {
        mul_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            Y.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dY.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dY};
}

torch::Tensor adj_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_forward_kernel", ([&] {
        adj_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}

std::vector<torch::Tensor> adj_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adj_backward_kernel", ([&] {
        adj_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, da};
}

torch::Tensor adjT_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_forward_kernel", ([&] {
        adjT_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}

std::vector<torch::Tensor> adjT_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor da = torch::zeros(a.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "adjT_backward_kernel", ([&] {
        adjT_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            da.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, da};
}

torch::Tensor act_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor q = torch::zeros(p.sizes(), p.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_forward_kernel", ([&] {
        act_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            q.data_ptr<scalar_t>(),
            batch_size);
    }));

    return q;
}

std::vector<torch::Tensor> act_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act_backward_kernel", ([&] {
        act_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dp.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dp};
}

torch::Tensor act4_forward_cpu(int group_id, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor q = torch::zeros(p.sizes(), p.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_forward_kernel", ([&] {
        act4_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            q.data_ptr<scalar_t>(),
            batch_size);
    }));

    return q;
}

std::vector<torch::Tensor> act4_backward_cpu(int group_id, torch::Tensor grad, torch::Tensor X, torch::Tensor p) {
    int batch_size = X.size(0);
    torch::Tensor dX = torch::zeros(X.sizes(), grad.options());
    torch::Tensor dp = torch::zeros(p.sizes(), grad.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "act4_backward_kernel", ([&] {
        act4_backward_kernel<group_t, scalar_t>(
            grad.data_ptr<scalar_t>(),
            X.data_ptr<scalar_t>(),
            p.data_ptr<scalar_t>(),
            dX.data_ptr<scalar_t>(),
            dp.data_ptr<scalar_t>(),
            batch_size);
    }));

    return {dX, dp};
}

torch::Tensor as_matrix_forward_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor T4x4 = torch::zeros({X.size(0), 4, 4}, X.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "as_matrix_forward_kernel", ([&] {
        as_matrix_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            T4x4.data_ptr<scalar_t>(),
            batch_size);
    }));

    return T4x4;
}

torch::Tensor orthogonal_projector_cpu(int group_id, torch::Tensor X) {
    int batch_size = X.size(0);
    torch::Tensor P;

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "orthogonal_projector_kernel", ([&] {
        P = torch::zeros({X.size(0), group_t::N, group_t::N}, X.options());
        orthogonal_projector_kernel<group_t, scalar_t>(X.data_ptr<scalar_t>(), P.data_ptr<scalar_t>(), batch_size);
    }));

    return P;
}

torch::Tensor jleft_forward_cpu(int group_id, torch::Tensor X, torch::Tensor a) {
    int batch_size = X.size(0);
    torch::Tensor b = torch::zeros(a.sizes(), a.options());

    DISPATCH_GROUP_AND_FLOATING_TYPES(group_id, X.type(), "jleft_forward_kernel", ([&] {
        jleft_forward_kernel<group_t, scalar_t>(
            X.data_ptr<scalar_t>(),
            a.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            batch_size);
    }));

    return b;
}