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

/**
 * This implements the classical "k-means" clustering algorithm.
 *
 * It takes as input file a series of lines where each line is a comma separated
 * or space separated list of values representing a vector. For instance:
 *
 * \verbatim
 * 1.1, 1.5, 0.9
 * 0.3, 0.4, -1.1
 * ...
 * \endverbatim
 *
 * It constructs a graph with a single vertex for each data point and simply
 * uses the "Map-Reduce" scheme to perform a k-means clustering of all
 * the datapoints.
 */

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

#include <limits>
#include <vector>
#include <map>
#include <iostream>
#include <stdlib.h>

#include <graphlab.hpp>

size_t NUM_CLUSTERS = 0;
bool IS_SPARSE = false;

struct cluster {
  cluster(): count(0), changed(false) { }
  std::vector<double> center;
  std::map<size_t, double> center_sparse;
  size_t count;
  bool changed;

  void save(graphlab::oarchive& oarc) const {
    oarc << center << count << changed << center_sparse;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> center >> count >> changed >> center_sparse;
  }
};

std::vector<cluster> CLUSTERS;

// the current cluster to initialize
size_t KMEANS_INITIALIZATION;

struct vertex_data{
  std::vector<double> point;
  std::map<size_t, double> point_sparse;
  size_t best_cluster;
  double best_distance;
  bool changed;

  void save(graphlab::oarchive& oarc) const {
    oarc << point << best_cluster << best_distance << changed << point_sparse;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> point >> best_cluster >> best_distance >> changed >> point_sparse;
  }
};

//use edges when edge weight file is given
struct edge_data {
  double weight;

  edge_data() :
      weight(0.0) {
  }
  explicit edge_data(double w) :
      weight(w) {
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << weight;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> weight;
  }
};

// helper function to compute distance between points
double sqr_distance(const std::vector<double>& a,
                    const std::vector<double>& b) {
  ASSERT_EQ(a.size(), b.size());
  double total = 0;
  for (size_t i = 0;i < a.size(); ++i) {
    double d = a[i] - b[i];
    total += d * d;
  }
  return total;
}

double sqr_distance(const std::map<size_t, double>& a,
                    const std::map<size_t, double>& b) {
  double total = 0.0;
  for(std::map<size_t, double>::const_iterator iter = a.begin();
      iter != a.end(); ++iter){
    size_t id = (*iter).first;
    double val = (*iter).second;
    if(b.find(id) != b.end()){
      double d = val - b.at(id);
      total += d*d;
    }else{
      total += val * val;
    }
  }
  for(std::map<size_t, double>::const_iterator iter = b.begin();
      iter != b.end(); ++iter){
    double val = (*iter).second;
    if(a.find((*iter).first) == a.end()){
      total += val * val;
    }
  }

  return total;

////   cosine distance is better for sparse datapoints?
//    double ip = 0.0;
//    double lenA = 0.0;
//    double lenB = 0.0;
//    for(std::map<size_t, double>::const_iterator iter = a.begin();
//        iter != a.end(); ++iter){
//      size_t id = (*iter).first;
//      double val = (*iter).second;
//      if(b.find(id) != b.end()){
//        ip += val * b.at(id);
//      }
//      lenA += val*val;
//    }
//
//    if(ip == 0.0 || lenA == 0.0)
//      return 1.0;
//
//    for(std::map<size_t, double>::const_iterator iter = b.begin();
//        iter != b.end(); ++iter){
//      double val = (*iter).second;
//      lenB += val * val;
//    }
//
//    if(lenB == 1.0)
//      return 1.0;
//
//    return 1.0 - ip/(sqrt(lenA)*sqrt(lenB));

}

// helper function to add two vectors
std::vector<double>& plus_equal_vector(std::vector<double>& a,
                                       const std::vector<double>& b) {
  ASSERT_EQ(a.size(), b.size());
  for (size_t i = 0;i < a.size(); ++i) {
    a[i] += b[i];
  }
  return a;
}

// helper function to add two vectors
std::map<size_t, double>& plus_equal_vector(std::map<size_t, double>& a,
                                       const std::map<size_t, double>& b) {
  for(std::map<size_t, double>::const_iterator iter = b.begin();
    iter != b.end(); ++iter){
    size_t id = (*iter).first;
    double val = (*iter).second;
    if(a.find(id) != a.end()){
      a[id] += b.at(id);
    }else{
      a.insert(std::make_pair(id, val));
    }
  }
  return a;
}

// helper function to scale a vector vectors
std::vector<double>& scale_vector(std::vector<double>& a, double d) {
  for (size_t i = 0;i < a.size(); ++i) {
    a[i] *= d;
  }
  return a;
}

// helper function to scale a vector vectors
std::map<size_t, double>& scale_vector(std::map<size_t, double>& a, double d) {
  for(std::map<size_t, double>::iterator iter = a.begin();
    iter != a.end(); ++iter){
  size_t id = (*iter).first;
  double val = (*iter).second;
  a[id] = val*d;
//    (*iter).second *= d;
  }
  return a;
}

typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

graphlab::atomic<graphlab::vertex_id_type> NEXT_VID;

// Read a line from a file and creates a vertex
bool vertex_loader(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;
  namespace qi = boost::spirit::qi;
  namespace ascii = boost::spirit::ascii;
  namespace phoenix = boost::phoenix;
  vertex_data vtx;
  const bool success = qi::phrase_parse
    (line.begin(), line.end(),
     //  Begin grammar
     (
      (qi::double_[phoenix::push_back(phoenix::ref(vtx.point), qi::_1)] % -qi::char_(",") )
      )
     ,
     //  End grammar
     ascii::space);

  if (!success) return false;
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(NEXT_VID.inc_ret_last(1), vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_sparse(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;

  vertex_data vtx;
  boost::char_separator<char> sep(" ");
  boost::tokenizer< boost::char_separator<char> > tokens(line, sep);
  BOOST_FOREACH (const std::string& t, tokens) {
    std::string::size_type pos = t.find(":");
    if(pos > 0){
      size_t id = (size_t)std::atoi(t.substr(0, pos).c_str());
      double val = std::atof(t.substr(pos+1, t.length() - pos -1).c_str());
      vtx.point_sparse.insert(std::make_pair(id, val));
    }
  }
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(NEXT_VID.inc_ret_last(1), vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_with_id(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;
  size_t id = 0;
  namespace qi = boost::spirit::qi;
  namespace ascii = boost::spirit::ascii;
  namespace phoenix = boost::phoenix;
  vertex_data vtx;
  const bool success = qi::phrase_parse
    (line.begin(), line.end(),
     //  Begin grammar
     (
      qi::ulong_[phoenix::ref(id) = qi::_1] >> -qi::char_(",") >>
      (qi::double_[phoenix::push_back(phoenix::ref(vtx.point), qi::_1)] % -qi::char_(",") )
      )
     ,
     //  End grammar
     ascii::space);

  if (!success) return false;
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(id, vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_with_id_sparse(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;

  vertex_data vtx;
  size_t id = 0;
  boost::char_separator<char> sep(" ");
  boost::tokenizer<boost::char_separator<char> > tokens(line, sep);
  bool first = true;
  BOOST_FOREACH (const std::string& t, tokens) {
    if(first){
      id = (size_t)std::atoi(t.c_str());
      first = false;
    }else{
      std::string::size_type pos = t.find(":");
      if(pos > 0){
        size_t id = (size_t)std::atoi(t.substr(0, pos).c_str());
        double val = std::atof(t.substr(pos+1, t.length() - pos -1).c_str());
        vtx.point_sparse.insert(std::make_pair(id, val));
      }
    }
  }
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(id, vtx);
  return true;
}

//call this when edge weight file is given.
//each line should be [source id] [target id] [weight].
//directions of edges are ignored.
bool edge_loader(graph_type& graph, const std::string& filename,
    const std::string& textline) {
  if (textline.empty())
    return true;
  std::stringstream strm(textline);
  size_t source_vid = 0;
  size_t target_vid = 0;
  double weight = 0.0;
  strm >> source_vid;
  strm.ignore(1);
  strm >> target_vid;
  strm.ignore(1);
  strm >> weight;
  if(source_vid != target_vid)
    graph.add_edge(source_vid, target_vid, edge_data(weight));
  return true;
}

// A set of Map Reduces to compute the maximum and minimum vector sizes
// to ensure that all vectors have the same length
struct max_point_size_reducer: public graphlab::IS_POD_TYPE {
  size_t max_point_size;

  static max_point_size_reducer get_max_point_size(const graph_type::vertex_type& v) {
    max_point_size_reducer r;
    r.max_point_size = v.data().point.size();
    return r;
  }

  max_point_size_reducer& operator+=(const max_point_size_reducer& other) {
    max_point_size = std::max(max_point_size, other.max_point_size);
    return *this;
  }
};

struct min_point_size_reducer: public graphlab::IS_POD_TYPE {
  size_t min_point_size;

  static min_point_size_reducer get_min_point_size(const graph_type::vertex_type& v) {
    min_point_size_reducer r;
    r.min_point_size = v.data().point.size();
    return r;
  }

  min_point_size_reducer& operator+=(const min_point_size_reducer& other) {
    min_point_size = std::min(min_point_size, other.min_point_size);
    return *this;
  }
};

/*
 * This transform vertices call is only used during
 * the initialization phase. It computes distance to
 * cluster[KMEANS_INITIALIZATION] and assigns itself
 * to the new cluster KMEANS_INITIALIZATION if the new distance
 * is smaller that its previous cluster assignment
 */
void kmeans_pp_initialization(graph_type::vertex_type& v) {
  double d = sqr_distance(v.data().point,
                          CLUSTERS[KMEANS_INITIALIZATION].center);
  if (v.data().best_distance > d) {
    v.data().best_distance = d;
    v.data().best_cluster = KMEANS_INITIALIZATION;
  }
}

void kmeans_pp_initialization_sparse(graph_type::vertex_type& v) {
  double d = sqr_distance(v.data().point_sparse,
                          CLUSTERS[KMEANS_INITIALIZATION].center_sparse);
  if (v.data().best_distance > d) {
    v.data().best_distance = d;
    v.data().best_cluster = KMEANS_INITIALIZATION;
  }
}

/*
 * Draws a random sample from the data points that is
 * proportionate to the "best distance" stored in the vertex.
 */
struct random_sample_reducer {
  std::vector<double> vtx;
  double weight;

  random_sample_reducer():weight(0) { }
  random_sample_reducer(const std::vector<double>& vtx,
                        double weight):vtx(vtx),weight(weight) { }

  static random_sample_reducer get_weight(const graph_type::vertex_type& v) {
    if (v.data().best_cluster == (size_t)(-1)) {
      return random_sample_reducer(v.data().point, 1);
    }
    else {
      return random_sample_reducer(v.data().point,
                                   v.data().best_distance);
    }
  }

  random_sample_reducer& operator+=(const random_sample_reducer& other) {
    double totalweight = weight + other.weight;
    // if any weight is too small, just quit
    if (totalweight <= 0) return *this;

    double myp = weight / (weight + other.weight);
    if (graphlab::random::bernoulli(myp)) {
      weight += other.weight;
      return *this;
    }
    else {
      vtx = other.vtx;
      weight += other.weight;
      return *this;
    }
  }

  void save(graphlab::oarchive &oarc) const {
    oarc << vtx << weight;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> vtx >> weight;
  }
};

struct random_sample_reducer_sparse{
  std::map<size_t, double> vtx;
  double weight;

  random_sample_reducer_sparse():weight(0) { }
  random_sample_reducer_sparse(const std::map<size_t, double>& vtx,
                        double weight):vtx(vtx),weight(weight) { }

  static random_sample_reducer_sparse get_weight(const graph_type::vertex_type& v) {
    if (v.data().best_cluster == (size_t)(-1)) {
      return random_sample_reducer_sparse(v.data().point_sparse, 1);
    }
    else {
      return random_sample_reducer_sparse(v.data().point_sparse,
                                   v.data().best_distance);
    }
  }

  random_sample_reducer_sparse& operator+=(const random_sample_reducer_sparse& other) {
    double totalweight = weight + other.weight;
    // if any weight is too small, just quit
    if (totalweight <= 0) return *this;

    double myp = weight / (weight + other.weight);
    if (graphlab::random::bernoulli(myp)) {
      weight += other.weight;
      return *this;
    }
    else {
      vtx = other.vtx;
      weight += other.weight;
      return *this;
    }
  }

  void save(graphlab::oarchive &oarc) const {
    oarc << vtx << weight;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> vtx >> weight;
  }
};

/*
 * This transform vertices call is used during the
 * actual k-means iteration. It computes distance to
 * all "changed" clusters and reassigns itself if necessary
 */
void kmeans_iteration(graph_type::vertex_type& v) {
  // if current vertex's cluster was modified, we invalidate the distance.
  // and we need to recompute to all existing clusters
  // otherwise, we just need to recompute to changed cluster centers.
  size_t prev_asg = v.data().best_cluster;
  if (CLUSTERS[v.data().best_cluster].changed) {
    // invalidate. recompute to all
    v.data().best_cluster = (size_t)(-1);
    v.data().best_distance = std::numeric_limits<double>::infinity();
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(v.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d = sqr_distance(v.data().point, CLUSTERS[i].center);
        if (d < v.data().best_distance) {
          v.data().best_distance = d;
          v.data().best_cluster = i;
        }
      }
    }
  }
  else {
    // just compute distance to what has changed
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].changed &&
          (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0)) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(v.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d= sqr_distance(v.data().point, CLUSTERS[i].center);
        if (d < v.data().best_distance) {
          v.data().best_distance = d;
          v.data().best_cluster = i;
        }
      }
    }
  }
  v.data().changed = (prev_asg != v.data().best_cluster);
}

//gathered information
//used when edge weight file is given
struct neighbor_info {
  std::map<size_t, double> cw_map;

  neighbor_info() :
      cw_map() {
  }
  neighbor_info(size_t clst, double weight) :
      cw_map() {
    cw_map.insert(std::make_pair(clst, weight));
  }

  neighbor_info& operator+=(const neighbor_info& other) {
    for (std::map<size_t, double>::const_iterator iter = other.cw_map.begin();
        iter != other.cw_map.end(); iter++) {
      size_t clst = iter->first;
      if (cw_map.find(clst) == cw_map.end()) {
        cw_map.insert(std::make_pair(clst, iter->second));
      } else {
        cw_map[clst] += iter->second;
      }
    }
    return *this;
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << cw_map;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> cw_map;
  }
};

//used when edge weight file is given
class cluster_assignment: public graphlab::ivertex_program<graph_type,
    neighbor_info>, public graphlab::IS_POD_TYPE {
public:
  //gather on all the edges
  edge_dir_type gather_edges(icontext_type& context,
      const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }

  //for each edge gather the weights and the assigned clusters of the neighbors
  neighbor_info gather(icontext_type& context, const vertex_type& vertex,
      edge_type& edge) const {
    if (edge.source().id() == vertex.id()) { //out edge
      return neighbor_info(edge.target().data().best_cluster,
          edge.data().weight);
    } else { //in edge
      return neighbor_info(edge.source().data().best_cluster,
          edge.data().weight);
    }
  }

  //assign a cluster, considering the clusters of neighbors
  void apply(icontext_type& context, vertex_type& vertex,
      const gather_type& total) {
    size_t past_clst = vertex.data().best_cluster;
    vertex.data().best_cluster = (size_t) (-1);
    vertex.data().best_distance = std::numeric_limits<double>::infinity();
    for (size_t i = 0; i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(vertex.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d = sqr_distance(vertex.data().point, CLUSTERS[i].center);
        //consider neighbors
        const std::map<size_t, double>& cw_map = total.cw_map;
        for (std::map<size_t, double>::const_iterator iter = cw_map.begin();
            iter != cw_map.end(); iter++) {
          size_t neighbor_cluster = iter->first;
          double total_wieght = iter->second;
          if (i == neighbor_cluster)
            d -= total_wieght;
        }
        if (d < vertex.data().best_distance) {
          vertex.data().best_distance = d;
          vertex.data().best_cluster = i;
        }
      }
    }
    vertex.data().changed = (past_clst != vertex.data().best_cluster);
  }

  //send signals to the neighbors when the cluster assignment has changed
  edge_dir_type scatter_edges(icontext_type& context,
      const vertex_type& vertex) const {
    if (vertex.data().changed)
      return graphlab::ALL_EDGES;
    else
      return graphlab::NO_EDGES;
  }

  void scatter(icontext_type& context, const vertex_type& vertex,
      edge_type& edge) const {
  }
};

/*
 * computes new cluster centers
 * Also accumulates a counter counting the number of vertices which
 * assignments changed.
 */
struct cluster_center_reducer {
  std::vector<cluster> new_clusters;
  size_t num_changed;
  double cost;

  cluster_center_reducer():new_clusters(NUM_CLUSTERS), num_changed(0), cost(0) { }

  static cluster_center_reducer get_center(const graph_type::vertex_type& v) {
    cluster_center_reducer cc;
    ASSERT_NE(v.data().best_cluster, (size_t)(-1));

    if(IS_SPARSE == true)
      cc.new_clusters[v.data().best_cluster].center_sparse = v.data().point_sparse;
    else
      cc.new_clusters[v.data().best_cluster].center = v.data().point;
    cc.new_clusters[v.data().best_cluster].count = 1;
    cc.num_changed = v.data().changed;
    cc.cost = v.data().best_distance;
    return cc;
  }

  cluster_center_reducer& operator+=(const cluster_center_reducer& other) {
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (new_clusters[i].count == 0) new_clusters[i] = other.new_clusters[i];
      else if (other.new_clusters[i].count > 0) {
        if(IS_SPARSE == true)
          plus_equal_vector(new_clusters[i].center_sparse, other.new_clusters[i].center_sparse);
        else
          plus_equal_vector(new_clusters[i].center, other.new_clusters[i].center);
        new_clusters[i].count += other.new_clusters[i].count;
      }
    }
    num_changed += other.num_changed;
    cost += other.cost;
    return *this;
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << new_clusters << num_changed <<cost;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> new_clusters >> num_changed >> cost;
  }
};

struct vertex_writer {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    for (size_t i = 0;i < v.data().point.size(); ++i) {
      strm << v.data().point[i] << "\t";
    }
    strm << v.data().best_distance << "\t";
    strm << v.data().best_cluster << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

struct vertex_writer_sparse {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    for(std::map<size_t, double>::iterator iter = v.data().point_sparse.begin();
        iter != v.data().point_sparse.end();++iter){
      strm << (*iter).first << ":" << (*iter).second << " ";
    }
    strm << v.data().best_cluster << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

struct vertex_writer_with_id {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    strm << v.id() << "\t";
    strm << v.data().best_cluster+1 << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

int main(int argc, char** argv) {
  std::cout << "Computes a K-means clustering of data.\n\n";

  graphlab::command_line_options clopts
    ("K-means clustering. The input data file is provided by the "
     "--data argument which is non-optional. The format of the data file is a "
     "collection of lines, where each line contains a comma or white-space "
     "separated lost of numeric values representing a vector. Every line "
     "must have the same number of values. The required --clusters=N "
     "argument denotes the number of clusters to generate. To store the output "
     "see the --output-cluster and --output-data arguments");

  std::string datafile;
  std::string outcluster_file;
  std::string outdata_file;
  std::string edgedata_file;
  size_t MAX_ITERATION = 0;
  bool use_id = false;
  clopts.attach_option("data", datafile,
                       "Input file. Each line holds a white-space or comma separated numeric vector");
  clopts.attach_option("clusters", NUM_CLUSTERS,
                       "The number of clusters to create.");
  clopts.attach_option("output-clusters", outcluster_file,
                       "If set, will write a file containing cluster centers "
                       "to this filename. This must be on the local filesystem "
                       "and must be accessible to the root node.");
  clopts.attach_option("output-data", outdata_file,
                       "If set, will output a copy of the input data with an additional "
                       "two columns. The first added column is the distance to assigned "
		       "center and the last is the assigned cluster centers. The output "
                       "will be written to a sequence of filenames where each file is "
                       "prefixed by this value. This may be on HDFS.");
  clopts.attach_option("sparse", IS_SPARSE,
                       "If set to true, will use a sparse vector representation."
                       "The file format is [feature id]:[value] [feature id]:[value] ..."
                       ", where [feature id] must be positive integer or zero.");
  clopts.attach_option("id", use_id,
                       "If set to true, will use ids for data points. The id of a data point "
                       "must be written at the head of each line of the input data. "
                       "The output data will consist of two columns: the first one "
                       "denotes the ids; the second one denotes the assigned clusters.");
  clopts.attach_option("pairwise-reward", edgedata_file,
                       "If set, will consider pairwise rewards when clustering. "
                       "Each line of the file beginning with the argument holds [id1] [id2] "
                       "[reward]. This mode must be used with --id option.");
  clopts.attach_option("max-iteration", MAX_ITERATION,
                       "The max number of iterations");

  if(!clopts.parse(argc, argv)) return EXIT_FAILURE;
  if (datafile == "") {
    std::cout << "--data is not optional\n";
    return EXIT_FAILURE;
  }
  if (NUM_CLUSTERS == 0) {
    std::cout << "--clusters is not optional\n";
    return EXIT_FAILURE;
  }
  if(edgedata_file.size() > 0){
    if(use_id == false){
      std::cout << "--id is not optional when you use edge data\n";
      return EXIT_FAILURE;
    }
  }

  graphlab::mpi_tools::init(argc, argv);
  graphlab::distributed_control dc;
  // load graph
  graph_type graph(dc, clopts);
  NEXT_VID = (((graphlab::vertex_id_type)1 << 31) / dc.numprocs()) * dc.procid();
  if(IS_SPARSE == true){
    if(use_id){
      graph.load(datafile, vertex_loader_with_id_sparse);
    }else{
      graph.load(datafile, vertex_loader_sparse);
    }
  }else{
    if(use_id){
      graph.load(datafile, vertex_loader_with_id);
    }else{
      graph.load(datafile, vertex_loader);
    }
  }
  if(edgedata_file.size() > 0){
    graph.load(edgedata_file, edge_loader);
  }
  graph.finalize();
  dc.cout() << "Number of datapoints: " << graph.num_vertices() << std::endl;

  if (graph.num_vertices() < NUM_CLUSTERS) {
    dc.cout() << "More clusters than datapoints! Cannot proceed" << std::endl;
    return EXIT_FAILURE;
  }

  dc.cout() << "Validating data...";

  CLUSTERS.resize(NUM_CLUSTERS);
  // make sure all have the same array length
  if(IS_SPARSE == false){
    size_t max_p_size = graph.map_reduce_vertices<max_point_size_reducer>
                                  (max_point_size_reducer::get_max_point_size).max_point_size;
    size_t min_p_size = graph.map_reduce_vertices<min_point_size_reducer>
                                  (min_point_size_reducer::get_min_point_size).min_point_size;
    if (max_p_size != min_p_size) {
      dc.cout() << "Data has dimensionality ranging from " << min_p_size << " to " << max_p_size
                << "! K-means cannot proceed!" << std::endl;
      return EXIT_FAILURE;
    }
    // allocate clusters
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      CLUSTERS[i].center.resize(max_p_size);
    }
  }

  dc.cout() << "Initializing using Kmeans++\n";
  // ok. perform kmeans++ initialization
  for (KMEANS_INITIALIZATION = 0;
       KMEANS_INITIALIZATION < NUM_CLUSTERS;
       ++KMEANS_INITIALIZATION) {

    if(IS_SPARSE == true){
      random_sample_reducer_sparse rs = graph.map_reduce_vertices<random_sample_reducer_sparse>
                                        (random_sample_reducer_sparse::get_weight);
      CLUSTERS[KMEANS_INITIALIZATION].center_sparse = rs.vtx;
      graph.transform_vertices(kmeans_pp_initialization_sparse);
    }else{
      random_sample_reducer rs = graph.map_reduce_vertices<random_sample_reducer>
                                        (random_sample_reducer::get_weight);
      CLUSTERS[KMEANS_INITIALIZATION].center = rs.vtx;
      graph.transform_vertices(kmeans_pp_initialization);
    }
  }

  // "reset" all clusters
  for (size_t i = 0; i < NUM_CLUSTERS; ++i) CLUSTERS[i].changed = true;
  // perform Kmeans iteration

  dc.cout() << "Running Kmeans...\n";
  bool clusters_changed = true;
  size_t iteration_count = 0;
  while(clusters_changed) {
		if(MAX_ITERATION > 0 && iteration_count >= MAX_ITERATION)
			break;

    cluster_center_reducer cc = graph.map_reduce_vertices<cluster_center_reducer>
                                    (cluster_center_reducer::get_center);
    // the first round (iteration_count == 0) is not so meaningful
    // since I am just recomputing the centers from the output of the KMeans++
    // initialization
    if (iteration_count > 0) {
      dc.cout() << "Kmeans iteration " << iteration_count << ": " <<
                 "# points with changed assignments = " << cc.num_changed <<
		 " total cost: " << cc.cost << std::endl;
    }
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      double d = cc.new_clusters[i].count;
      if(IS_SPARSE){
        if (d > 0) scale_vector(cc.new_clusters[i].center_sparse, 1.0 / d);
        if (cc.new_clusters[i].count == 0 && CLUSTERS[i].count > 0) {
          dc.cout() << "Cluster " << i << " lost" << std::endl;
          CLUSTERS[i].center_sparse.clear();
          CLUSTERS[i].count = 0;
          CLUSTERS[i].changed = false;
        }
        else {
          CLUSTERS[i] = cc.new_clusters[i];
          CLUSTERS[i].changed = true;
        }
      }else{
        if (d > 0) scale_vector(cc.new_clusters[i].center, 1.0 / d);
        if (cc.new_clusters[i].count == 0 && CLUSTERS[i].count > 0) {
          dc.cout() << "Cluster " << i << " lost" << std::endl;
          CLUSTERS[i].center.clear();
          CLUSTERS[i].count = 0;
          CLUSTERS[i].changed = false;
        }
        else {
          CLUSTERS[i] = cc.new_clusters[i];
          CLUSTERS[i].changed = true;
        }
      }
    }
    clusters_changed = iteration_count == 0 || cc.num_changed > 0;

    if(edgedata_file.size() > 0){
      clopts.engine_args.set_option("factorized", true);
      graphlab::omni_engine<cluster_assignment> engine(dc, graph, "async", clopts);
      engine.signal_all();
      engine.start();
    }else{
      graph.transform_vertices(kmeans_iteration);
    }

    ++iteration_count;
  }

  if (!outcluster_file.empty() && dc.procid() == 0) {
    dc.cout() << "Writing Cluster Centers..." << std::endl;
    std::ofstream fout(outcluster_file.c_str());
    if(IS_SPARSE){
      for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
        if(use_id)
          fout << i+1 << "\t";
        for (std::map<size_t, double>::iterator iter = CLUSTERS[i].center_sparse.begin();
             iter != CLUSTERS[i].center_sparse.end();++iter) {
          fout << (*iter).first << ":" << (*iter).second << " ";
        }
        fout << "\n";
      }
    }else{
      for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
        if(use_id)
          fout << i+1 << "\t";
        for (size_t j = 0; j < CLUSTERS[i].center.size(); ++j) {
          fout << CLUSTERS[i].center[j] << " ";
        }
        fout << "\n";
      }
    }
  }

  if (!outdata_file.empty()) {
    dc.cout() << "Writing Data with cluster assignments...\n" << std::endl;
    if(use_id){
      graph.save(outdata_file, vertex_writer_with_id(), false, true, false, 1);
    }else{
      if(IS_SPARSE == true)
        graph.save(outdata_file, vertex_writer_sparse(), false, true, false, 1);
      else
        graph.save(outdata_file, vertex_writer(), false, true, false, 1);
    }
  }

  graphlab::mpi_tools::finalize();
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS
  @copyright 2016 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_MINMAG_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_MINMAG_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!

    @ingroup group-ieee
    Function object implementing minmag capabilities

    Returns the input value which have the least absolute value.

    @par Semantic:

    @code
    T r = minmag(x,y);
    @endcode

    is similar to:

    @code
    T r = abs(x) < abs(y) ? x : abs(y) < abs(x) ? y : min(x, y);
    @endcode

   @see min, minnummag, minnum

  **/
  const boost::dispatch::functor<tag::minmag_> minmag = {};
} }
#endif

#include <boost/simd/function/scalar/minmag.hpp>
#include <boost/simd/function/simd/minmag.hpp>

#endif


/**
*   Point.hpp
*
*   Copyright 2015 The Climate Corporation
*   Copyright 2015 Tony Stone
*
*   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.
*
*   Created by Tony Stone on 6/9/15.
*
*   MODIFIED 2015 BY Tony Stone. Modifications licensed under Apache License, Version 2.0.
*
*/
#pragma once

#ifndef __Point_HPP_
#define __Point_HPP_

#include "Geometry.hpp"

#include <boost/mpl/int.hpp>
#include <boost/static_assert.hpp>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_type.hpp>
#include <boost/geometry/core/coordinate_system.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/cs.hpp>

#include <boost/geometry/geometries/concepts/point_concept.hpp>

namespace geofeatures {

    /**
     * @class       Point
     *
     * @brief       A simple 2d point.
     *
     * @author      Tony Stone
     * @date        6/9/15
     */
    class Point : public Geometry {

    public:
        // member types
        typedef double CoordinateType;
        typedef boost::geometry::cs::cartesian CoordinateSystem;

    private:
        static const std::size_t dimensionCount = 2;
        enum { cs_check = sizeof(CoordinateSystem) };

    public:
        // construction/destruction
        inline Point(CoordinateType const& x = 0.0, CoordinateType const& y = 0.0) noexcept : Geometry() {
            values_[0] = x;
            values_[1] = y;
        }

        inline virtual ~Point() noexcept {};

        template <std::size_t K>
        inline CoordinateType const& get() const noexcept {
            static_assert(K < dimensionCount, "Index out of bounds");
            return values_[K];
        }

        template <std::size_t K>
        inline void set(CoordinateType const& value) noexcept {
            static_assert(K < dimensionCount, "Index out of bounds");
            values_[K] = value;
        }

    private:
        CoordinateType values_[dimensionCount];
    };

}   // namespace geofeatures

namespace geofeatures_boost {
    namespace geometry {
        namespace traits {
            // Adapt geofeatures::Point to Boost.Geometry

            template<> struct tag<geofeatures::Point>
            { typedef point_tag type; };

            template<> struct coordinate_type<geofeatures::Point>
            { typedef geofeatures::Point::CoordinateType type; };

            template<> struct coordinate_system<geofeatures::Point>
            { typedef geofeatures::Point::CoordinateSystem type; };

            template<> struct dimension<geofeatures::Point> : boost::mpl::int_<2> {};

            template<>
            struct access<geofeatures::Point, 0>
            {
                static inline geofeatures::Point::CoordinateType get(geofeatures::Point const& p) {
                    return p.template get<0>();
                }

                static inline void set(geofeatures::Point& p, geofeatures::Point::CoordinateType const& value) {
                    p.template set<0>(value);
                }
            };

            template<>
            struct access<geofeatures::Point, 1>
            {
                static inline geofeatures::Point::CoordinateType get(geofeatures::Point const& p) {
                    return p.template get<1>();
                }

                static inline void set(geofeatures::Point& p, geofeatures::Point::CoordinateType const& value) {
                    p.template set<1>(value);
                }
            };
        }
    }
} // namespace boost::geometry::traits

#endif //__Point_HPP_


#include "logval.h"
#define BOOST_TEST_MODULE LogValTest
#include <boost/test/unit_test.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <iostream>

using namespace std;

BOOST_AUTO_TEST_CASE(Order) {
  LogVal<double> a(-0.3);
  LogVal<double> b(0.3);
  LogVal<double> c(2.4);
  BOOST_CHECK_LT(a,b);
  BOOST_CHECK_LT(b,c);
  BOOST_CHECK_LT(a,c);
  BOOST_CHECK(b >= a);
  BOOST_CHECK(c >= a);
  BOOST_CHECK(c >= b);
  BOOST_CHECK(c >= c);
  BOOST_CHECK(b >= b);
  BOOST_CHECK(a >= a);
}

BOOST_AUTO_TEST_CASE(Negate) {
  LogVal<double> x(-2.4);
  LogVal<double> y(2.4);
  y.negate();
  BOOST_CHECK_CLOSE(x.as_float(),y.as_float(), 1e-6);
}

BOOST_AUTO_TEST_CASE(Inverse) {
  LogVal<double> x(1/2.4);
  LogVal<double> y(2.4);
  BOOST_CHECK_CLOSE(x.as_float(),y.inverse().as_float(), 1e-6);
}

BOOST_AUTO_TEST_CASE(Minus) {
  LogVal<double> x(12);
  LogVal<double> y(2);
  LogVal<double> z1 = x - y;
  LogVal<double> z2 = x;
  z2 -= y;
  BOOST_CHECK_CLOSE(z1.as_float(), z2.as_float(), 1e-6);
  BOOST_CHECK_CLOSE(z1.as_float(), 10.0, 1e-6);
  BOOST_CHECK_CLOSE((y - x).as_float(), -10.0, 1e-6);
}

BOOST_AUTO_TEST_CASE(TestOps) {
  LogVal<double> x(-12.12);
  LogVal<double> y(x);
  cerr << x << endl;
  cerr << (x*y) << endl;
  cerr << (x*y + x) << endl;
  cerr << (x + x*y) << endl;
  cerr << log1p(-0.5) << endl;
  LogVal<double> aa(0.2);
  LogVal<double> bb(-0.3);
  cerr << (aa + bb) << endl;
  cerr << (bb + aa) << endl;
  BOOST_CHECK_CLOSE((aa + bb).as_float(), (bb + aa).as_float(), 1e-6);
  BOOST_CHECK_CLOSE((aa + bb).as_float(), -0.1, 1e-6);
}

BOOST_AUTO_TEST_CASE(TestSizes) {
  cerr << sizeof(LogVal<double>) << endl;
  cerr << sizeof(LogVal<float>) << endl;
  cerr << sizeof(void*) << endl;
}


#include <iostream>

#include <boost/spirit/include/qi.hpp>

#include "objects/object.hpp"
#include "objects/system/stream.hpp"
#include "objects/core/symbol.hpp"
#include "objects/core/char.hpp"

#include "objects/numerics/float.hpp"
#include "objects/numerics/integer.hpp"

#include "objects/collections/list.hpp"
#include "objects/collections/string.hpp"
#include "objects/collections/ast.hpp"

#include "objects/invokable/inbuilt_function.hpp"
#include "objects/invokable/user_function.hpp"
#include "objects/invokable/effect.hpp"

#include "globals/system.hpp"
#include "globals/collections.hpp"
#include "globals/core.hpp"

using obj::Object;

using sys::cons::read;

Object *to_value(std::string &token) {
  // character literal
  if (token[0] == '#') {
    obj::Char *ch;
    if (token.size() == 2) {
      ch = new obj::Char(token[1]);
    } else {
      if (token == "#newline") {
        ch = new obj::Char('\n');
      } else {
        std::string err = "non-regognized character: " + token;
        throw err;
      }
    }
    return ch;
  }
  try {
    if (token.find(".") != std::string::npos) {
      double num = stod(token);
      return new obj::Float(num);
    } else {
      int num = stoi(token);
      return new obj::Integer(num);
    }
  } catch (std::invalid_argument &) {}
  // it is a symbol, first split based on ':'
  mod::vector<std::string> path;
  std:: string current = "";
  for (char c : token) {
    if (c == ':') {
      // get string as symbols
      path.push_back(current);
      current = "";
    } else {
      current += c;
    }
  }
  path.push_back(current);

  if (path.size() == 1) {
    return std::get<Object *>(
        (*core::mk_symbol)({new obj::String(path.front())}, *context::empty));
  } else {
    return new obj::MultiSymbol(path);
  }
}

Object *to_ast(mod::list<Object *> list) {
  if (list.size() > 0) {
    mod::optional<Object*> op = list.front();
    list.pop_front();
    mod::vector<Object*> vec;
    std::copy(list.begin(), list.end(), std::back_inserter(vec));
    return new obj::AST(op, vec);
  } else {
  } return new obj::AST;
}

Object *mac_lsquare(std::istream &stream, char c) {
  // continue to add tokens to list until we hit a ')'
  mod::vector<Object *> list;
  while (!stream.eof()) {
    c = stream.peek();
    switch (c) {
    case ' ':
    case '\n':
    case '\t':
      stream.read(&c, 1);
      break;
    case ']': {
      stream.read(&c, 1);
      if (list.empty()) {
        return coll::nil;
      } else {
        return new obj::AST(coll::list, list);
      }
      break;
    } break;
    default:
      list.push_back(read(stream));
      break;
    }
  }
  throw "read error: [ with no matching ]!";
}

Object *mac_lparen(std::istream& stream, char c) {
  // continue to add tokens to list until we hit a ')'
  mod::list<obj::Object *> list;

  while (!stream.eof()) {
    c = stream.peek();
    switch (c) {
    case ' ':
    case '\n':
    case '\t':
      stream.read(&c, 1);
      break;
    case ')': {
      stream.read(&c, 1);
      if (list.empty())
        return core::unit;
      else
        return to_ast(list);
      break;
    } break;
    default:
      list.push_back(read(stream));
      break;
    }
  }
  throw "( or [ with no matching ) or ]!";
}

Object *mac_string(std::istream& stream, char c) {
  stream.read(&c, 1);
  std::string str;
  while (!stream.eof()) {
    if (c == '"') {
      return new obj::String(str);
    }
    if (c == '\\') {
      stream.read(&c, 1);
      str += c;
      stream.read(&c, 1);
    } else {
      str += c;
      stream.read(&c, 1);
    }
  }
  std::string err = "non-terminated string";
  throw err;
}

Object *mac_token(std::istream& stream, char c) {
  std::string token;
  bool cond = false;
  while (!stream.eof()) {
    switch (c) {
    case ' ':
    case '\n':
    case '\t':
    case '(':
    case ')':
    case '[':
    case ']':
      return to_value(token);
      break;
      // escape character!
    case '\\': // treat the next character as normal
      token += stream.peek();
      stream.read(&c, 1);
      break;
      // the hash is used for various literals

    default:
      token += c;
      break;
    }

    if (cond) {
      stream.read(&c, 1);
    }
    cond = true;
    int ch = stream.peek();
    if (ch == EOF)
      return to_value(token);
    c = ch;
  }

  return to_value(token);
}

obj::AST *fold_ast(obj::AST *ast) {
  if (ast->args.size() == 2) {
    return ast;
  } else {
    size_t n = ast->args.size() - 1;
    return new obj::AST(
        ast->op.value(),
        {fold_ast(new obj::AST(ast->op.value(),
                               ast->args.take(n))),
         ast->args[n]});
  }
}

/*! Take an AST that has an arbitrary number of args, e.g. (+ 1 2 3 4) and turn
 *  it into one of only binary operations, i.e. (+ (+ (+ 1 2) 3) 4).
 */
Object *mac_fold_oper(std::istream& stream, char c) {
  using obj::AST;

  Object* obj = read(stream);
  AST* ast = dynamic_cast<AST*>(obj);
  if (ast == nullptr) {
    std::string err = "error in mac_fold_oper: must provide an ast";
    throw err;
  }
  if (ast->args.size() > 1) {
    return fold_ast(ast);
  } else {
    std::string err = "error in mac_fold_oper: must provide an ast with at least 2 args";
    throw err;
  }

}

Object *mac_quote(std::istream& stream, char c) {
  Object* obj = read(stream);
  return new obj::AST(core::quote, {obj});
}

/*! The backquote macro is quite complicated, and itself allows the use of
 *  several other macros - notably eval (,) and list-splice (,@)
 */
Object *mac_backquote(std::istream& stream, char c) {
  while (!stream.eof()) {
    c = stream.peek();
    switch (c) {
    case '\t':
    case '\n':
    case ' ':
      stream.read(&c, 1);
      break;
    case ',':
      stream.read(&c, 1);
      return read(stream);
    default:
      Object *obj = read(stream);
      return new obj::AST(core::quote, {obj});
    }
  }

  return ::core::unit;
}

Object *sys::cons::read(std::istream& stream) {
  // strings work too!
  char c = '\0';
  stream.read(&c, 1);
  while (!stream.eof()) {
    switch (c) {
      // whitespace characters
    case '\t':
    case '\n':
    case ' ':
      break;

    case '(':
      return mac_lparen(stream, c);
      break;
    case '[':
      return mac_lsquare(stream, c);
      break;

    case ']':
    case ')': {
      std::string err = "error: ') or ]' with no matching '( or ['!";
      throw err;
    } break;
    case '"': {
      return mac_string(stream, c);
    }
    case '!': {
      return mac_fold_oper(stream, c);
    }
    case '\'': {
      return mac_quote(stream, c);
    }
    case '`': {
      return mac_backquote(stream, c);
    }

    default:
      return mac_token(stream, c);
      break;
    }

    stream.read(&c, 1);
  }

  // todo: more advanced read options!
  return ::core::unit;
}


/*
 * Copyright (c) 2015-2021 Agalmic Ventures LLC (www.agalmicventures.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

#include <cmath>

#include <boost/test/unit_test.hpp>

#include "Werk/Math/ContinuousEma.hpp"

BOOST_AUTO_TEST_SUITE(ContinuousEmaTest)

BOOST_AUTO_TEST_CASE(testEmpty)
{
	werk::ContinuousEma ema(10.0);

	BOOST_CHECK_EQUAL(ema.halfLife(), 10.0);
	BOOST_CHECK(ema.factor() < 1.0 / 10.0);
	BOOST_CHECK(std::isnan(ema.value()));
}

BOOST_AUTO_TEST_CASE(testConstant)
{
	werk::ContinuousEma ema(10.0);

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(4.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(17.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);

	ema.add(-5.0);
	BOOST_CHECK_EQUAL(ema.value(), 5.0);
}

BOOST_AUTO_TEST_CASE(testDecreasing)
{
	werk::ContinuousEma ema(60.0);

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);

	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(13.0, 0.0);
	BOOST_CHECK(ema.value() < 10.0);

	double value = ema.value();
	ema.sample(19.0, 0.0);
	BOOST_CHECK(ema.value() < value);
}

BOOST_AUTO_TEST_CASE(testIncreasing)
{
	werk::ContinuousEma ema(60.0);

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(7.0, 13.0);
	BOOST_CHECK(ema.value() > 10);

	double value = ema.value();
	ema.sample(19.0, 20.0);
	BOOST_CHECK(ema.value() > value);
}

BOOST_AUTO_TEST_CASE(testOscillating)
{
	werk::ContinuousEma ema(10.0);

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(10.0, 0.0);
	BOOST_CHECK_CLOSE(ema.value(), 5.0, 0.5);
	ema.sample(20.0, 10.0);
	BOOST_CHECK_CLOSE(ema.value(), 7.5, 0.5);
	ema.sample(30.0, 0.0);
	BOOST_CHECK_CLOSE(ema.value(), 3.75, 0.5);
}

BOOST_AUTO_TEST_CASE(testNan)
{
	werk::ContinuousEma ema(60.0);

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(3.0, std::numeric_limits<double>::quiet_NaN());
	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(5.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
}

BOOST_AUTO_TEST_CASE(testInfiniteHalfLife)
{
	werk::ContinuousEma ema(std::numeric_limits<double>::infinity());

	BOOST_CHECK(std::isnan(ema.value()));
	ema.sample(0.0, 10.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(3.0, 15.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
	ema.sample(15.0, 2310.0);
	BOOST_CHECK_EQUAL(ema.value(), 10.0);
}

BOOST_AUTO_TEST_SUITE_END()


#include <cfloat>
#include <cmath>

#include <boost/math/distributions/normal.hpp>

#include <DataTypes/DataTypeTuple.h>
#include <DataTypes/DataTypesDecimal.h>
#include <DataTypes/DataTypesNumber.h>
#include <Columns/ColumnTuple.h>
#include <Columns/ColumnsNumber.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/IFunction.h>
#include <Functions/castTypeToEither.h>
#include <Interpreters/castColumn.h>

namespace DB
{

namespace ErrorCodes
{
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
}

template <typename Impl>
class FunctionMinSampleSize : public IFunction
{
public:
    static constexpr auto name = Impl::name;

    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionMinSampleSize<Impl>>(); }

    String getName() const override { return name; }

    size_t getNumberOfArguments() const override { return Impl::num_args; }
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override
    {
        return ColumnNumbers(std::begin(Impl::const_args), std::end(Impl::const_args));
    }

    bool useDefaultImplementationForNulls() const override { return false; }
    bool useDefaultImplementationForConstants() const override { return true; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }

    static DataTypePtr getReturnType()
    {
        auto float_64_type = std::make_shared<DataTypeNumber<Float64>>();

        DataTypes types{
            float_64_type,
            float_64_type,
            float_64_type,
        };

        Strings names{
            "minimum_sample_size",
            "detect_range_lower",
            "detect_range_upper",
        };

        return std::make_shared<DataTypeTuple>(std::move(types), std::move(names));
    }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        Impl::validateArguments(arguments);
        return getReturnType();
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
    {
        return Impl::execute(arguments, input_rows_count);
    }
};

static bool isBetweenZeroAndOne(Float64 v)
{
    return v >= 0.0 && v <= 1.0 && fabs(v - 0.0) >= DBL_EPSILON && fabs(v - 1.0) >= DBL_EPSILON;
}

struct ContinousImpl
{
    static constexpr auto name = "minSampleSizeContinous";
    static constexpr size_t num_args = 5;
    static constexpr size_t const_args[] = {2, 3, 4};

    static void validateArguments(const DataTypes & arguments)
    {
        for (size_t i = 0; i < arguments.size(); ++i)
        {
            if (!isNativeNumber(arguments[i]))
            {
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The {}th Argument of function {} must be a number.", i + 1, name);
            }
        }
    }

    static ColumnPtr execute(const ColumnsWithTypeAndName & arguments, size_t input_rows_count)
    {
        auto float_64_type = std::make_shared<DataTypeFloat64>();
        auto baseline_argument = arguments[0];
        baseline_argument.column = baseline_argument.column->convertToFullColumnIfConst();
        auto baseline_column_untyped = castColumnAccurate(baseline_argument, float_64_type);
        const auto * baseline_column = checkAndGetColumn<ColumnVector<Float64>>(*baseline_column_untyped);
        const auto & baseline_column_data = baseline_column->getData();

        auto sigma_argument = arguments[1];
        sigma_argument.column = sigma_argument.column->convertToFullColumnIfConst();
        auto sigma_column_untyped = castColumnAccurate(sigma_argument, float_64_type);
        const auto * sigma_column = checkAndGetColumn<ColumnVector<Float64>>(*sigma_column_untyped);
        const auto & sigma_column_data = sigma_column->getData();

        const IColumn & col_mde = *arguments[2].column;
        const IColumn & col_power = *arguments[3].column;
        const IColumn & col_alpha = *arguments[4].column;

        auto res_min_sample_size = ColumnFloat64::create();
        auto & data_min_sample_size = res_min_sample_size->getData();
        data_min_sample_size.reserve(input_rows_count);

        auto res_detect_lower = ColumnFloat64::create();
        auto & data_detect_lower = res_detect_lower->getData();
        data_detect_lower.reserve(input_rows_count);

        auto res_detect_upper = ColumnFloat64::create();
        auto & data_detect_upper = res_detect_upper->getData();
        data_detect_upper.reserve(input_rows_count);

        /// Minimal Detectable Effect
        const Float64 mde = col_mde.getFloat64(0);
        /// Sufficient statistical power to detect a treatment effect
        const Float64 power = col_power.getFloat64(0);
        /// Significance level
        const Float64 alpha = col_alpha.getFloat64(0);

        boost::math::normal_distribution<> nd(0.0, 1.0);

        for (size_t row_num = 0; row_num < input_rows_count; ++row_num)
        {
            /// Mean of control-metric
            Float64 baseline = baseline_column_data[row_num];
            /// Standard deviation of conrol-metric
            Float64 sigma = sigma_column_data[row_num];

            if (!std::isfinite(baseline) || !std::isfinite(sigma) || !isBetweenZeroAndOne(mde) || !isBetweenZeroAndOne(power)
                || !isBetweenZeroAndOne(alpha))
            {
                data_min_sample_size.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_lower.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_upper.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                continue;
            }

            Float64 delta = baseline * mde;

            using namespace boost::math;
            /// https://towardsdatascience.com/required-sample-size-for-a-b-testing-6f6608dd330a
            /// \frac{2\sigma^{2} * (Z_{1 - alpha /2} + Z_{power})^{2}}{\Delta^{2}}
            Float64 min_sample_size
                = 2 * std::pow(sigma, 2) * std::pow(quantile(nd, 1.0 - alpha / 2) + quantile(nd, power), 2) / std::pow(delta, 2);

            data_min_sample_size.emplace_back(min_sample_size);
            data_detect_lower.emplace_back(baseline - delta);
            data_detect_upper.emplace_back(baseline + delta);
        }

        return ColumnTuple::create(Columns{std::move(res_min_sample_size), std::move(res_detect_lower), std::move(res_detect_upper)});
    }
};

struct ConversionImpl
{
    static constexpr auto name = "minSampleSizeConversion";
    static constexpr size_t num_args = 4;
    static constexpr size_t const_args[] = {1, 2, 3};

    static void validateArguments(const DataTypes & arguments)
    {
        size_t arguments_size = arguments.size();
        for (size_t i = 0; i < arguments_size; ++i)
        {
            if (!isFloat(arguments[i]))
            {
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The {}th argument of function {} must be a float.", i + 1, name);
            }
        }
    }

    static ColumnPtr execute(const ColumnsWithTypeAndName & arguments, size_t input_rows_count)
    {
        auto first_argument_column = castColumnAccurate(arguments[0], std::make_shared<DataTypeFloat64>());

        if (const ColumnConst * const col_p1_const = checkAndGetColumnConst<ColumnVector<Float64>>(first_argument_column.get()))
        {
            const Float64 left_value = col_p1_const->template getValue<Float64>();
            return process<true>(arguments, &left_value, input_rows_count);
        }
        else if (const ColumnVector<Float64> * const col_p1 = checkAndGetColumn<ColumnVector<Float64>>(first_argument_column.get()))
        {
            return process<false>(arguments, col_p1->getData().data(), input_rows_count);
        }
        else
        {
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The first argument of function {} must be a float.", name);
        }
    }

    template <bool const_p1>
    static ColumnPtr process(const ColumnsWithTypeAndName & arguments, const Float64 * col_p1, const size_t input_rows_count)
    {
        const IColumn & col_mde = *arguments[1].column;
        const IColumn & col_power = *arguments[2].column;
        const IColumn & col_alpha = *arguments[3].column;

        auto res_min_sample_size = ColumnFloat64::create();
        auto & data_min_sample_size = res_min_sample_size->getData();
        data_min_sample_size.reserve(input_rows_count);

        auto res_detect_lower = ColumnFloat64::create();
        auto & data_detect_lower = res_detect_lower->getData();
        data_detect_lower.reserve(input_rows_count);

        auto res_detect_upper = ColumnFloat64::create();
        auto & data_detect_upper = res_detect_upper->getData();
        data_detect_upper.reserve(input_rows_count);

        /// Minimal Detectable Effect
        const Float64 mde = col_mde.getFloat64(0);
        /// Sufficient statistical power to detect a treatment effect
        const Float64 power = col_power.getFloat64(0);
        /// Significance level
        const Float64 alpha = col_alpha.getFloat64(0);

        boost::math::normal_distribution<> nd(0.0, 1.0);

        for (size_t row_num = 0; row_num < input_rows_count; ++row_num)
        {
            /// Proportion of control-metric
            Float64 p1;

            if constexpr (const_p1)
            {
                p1 = col_p1[0];
            }
            else if constexpr (!const_p1)
            {
                p1 = col_p1[row_num];
            }

            if (!std::isfinite(p1) || !isBetweenZeroAndOne(mde) || !isBetweenZeroAndOne(power) || !isBetweenZeroAndOne(alpha))
            {
                data_min_sample_size.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_lower.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_upper.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                continue;
            }

            Float64 q1 = 1.0 - p1;
            Float64 p2 = p1 + mde;
            Float64 q2 = 1.0 - p2;
            Float64 p_bar = (p1 + p2) / 2.0;
            Float64 q_bar = 1.0 - p_bar;

            using namespace boost::math;
            /// https://towardsdatascience.com/required-sample-size-for-a-b-testing-6f6608dd330a
            /// \frac{(Z_{1-alpha/2} * \sqrt{2*\bar{p}*\bar{q}} + Z_{power} * \sqrt{p1*q1+p2*q2})^{2}}{\Delta^{2}}
            Float64 min_sample_size
                = std::pow(
                      quantile(nd, 1.0 - alpha / 2.0) * std::sqrt(2.0 * p_bar * q_bar) + quantile(nd, power) * std::sqrt(p1 * q1 + p2 * q2),
                      2)
                / std::pow(mde, 2);

            data_min_sample_size.emplace_back(min_sample_size);
            data_detect_lower.emplace_back(p1 - mde);
            data_detect_upper.emplace_back(p1 + mde);
        }

        return ColumnTuple::create(Columns{std::move(res_min_sample_size), std::move(res_detect_lower), std::move(res_detect_upper)});
    }
};

void registerFunctionMinSampleSize(FunctionFactory & factory)
{
    factory.registerFunction<FunctionMinSampleSize<ContinousImpl>>();
    factory.registerFunction<FunctionMinSampleSize<ConversionImpl>>();
}

}


#include <pluginlib/class_list_macros.h>
#include <kdl_parser/kdl_parser.hpp>
#include <math.h>
#include <Eigen/LU>

#include <utils/pseudo_inversion.h>
#include <utils/skew_symmetric.h>
#include <lwr_controllers/multi_task_priority_inverse_dynamics.h>

namespace lwr_controllers
{
	MultiTaskPriorityInverseDynamics::MultiTaskPriorityInverseDynamics() {}
	MultiTaskPriorityInverseDynamics::~MultiTaskPriorityInverseDynamics() {}

	bool MultiTaskPriorityInverseDynamics::init(hardware_interface::EffortJointInterface *robot, ros::NodeHandle &n)
	{
        KinematicChainControllerBase<hardware_interface::EffortJointInterface>::init(robot, n);

		jnt_to_jac_solver_.reset(new KDL::ChainJntToJacSolver(kdl_chain_));
		id_solver_.reset(new KDL::ChainDynParam(kdl_chain_,gravity_));
		fk_pos_solver_.reset(new KDL::ChainFkSolverPos_recursive(kdl_chain_));

		qdot_last_.resize(kdl_chain_.getNrOfJoints());
		tau_.resize(kdl_chain_.getNrOfJoints());
		J_.resize(kdl_chain_.getNrOfJoints());
		J_dot_.resize(kdl_chain_.getNrOfJoints());
		Kp_.resize(kdl_chain_.getNrOfJoints());
		Kd_.resize(kdl_chain_.getNrOfJoints());
		M_.resize(kdl_chain_.getNrOfJoints());
		C_.resize(kdl_chain_.getNrOfJoints());
		G_.resize(kdl_chain_.getNrOfJoints());

		sub_command_ = nh_.subscribe("command", 1, &MultiTaskPriorityInverseDynamics::command, this);

		pub_error_ = nh_.advertise<std_msgs::Float64MultiArray>("error", 1000);
		pub_marker_ = nh_.advertise<visualization_msgs::MarkerArray>("marker",1000);

		return true;
	}

	void MultiTaskPriorityInverseDynamics::starting(const ros::Time& time)
	{
		// get joint positions
  		for(int i=0; i < joint_handles_.size(); i++)
  		{
    		joint_msr_states_.q(i) = joint_handles_[i].getPosition();
    		joint_msr_states_.qdot(i) = joint_handles_[i].getVelocity();
    		joint_des_states_.q(i) = joint_msr_states_.q(i);
    		joint_des_states_.qdot(i) = joint_msr_states_.qdot(i);
    		Kp_(i) = 100;
  			Kd_(i) = 20;
    	}

    	I_ = Eigen::Matrix<double,7,7>::Identity(7,7);

    	first_step_ = 0;
    	cmd_flag_ = 0;

	}

	void MultiTaskPriorityInverseDynamics::update(const ros::Time& time, const ros::Duration& period)
	{
		// get joint positions
  		for(int i=0; i < joint_handles_.size(); i++)
  		{
    		joint_msr_states_.q(i) = joint_handles_[i].getPosition();
    		joint_msr_states_.qdot(i) = joint_handles_[i].getVelocity();
    	}

    	// clearing error msg before publishing
    	msg_err_.data.clear();

    	if (cmd_flag_)
    	{
			// computing Inertia, Coriolis and Gravity matrices
	    	id_solver_->JntToMass(joint_msr_states_.q, M_);
	    	id_solver_->JntToCoriolis(joint_msr_states_.q, joint_msr_states_.qdot, C_);
	    	id_solver_->JntToGravity(joint_msr_states_.q, G_);
	    	// deleting gravity contribute
	    	G_.data.setZero();

	    	// computing the inverse of M_ now, since it will be used often
	    	pseudo_inverse(M_.data,M_inv_,false);

    		// resetting N and tau(t=0) for the highest priority task
    		N_trans_ = I_;
    		SetToZero(tau_);

    		for (int index = 0; index < ntasks_; index++)
    		{
		    	// computing Jacobian J(q)
		    	jnt_to_jac_solver_->JntToJac(joint_msr_states_.q,J_,links_index_[index]);

		    	if (first_step_)
		    	{
		    		J_last_[index] = J_;
		    		first_step_ = 0;
		    		return;
		    	}

		    	// computing the derivative of Jacobian J_dot(q) through numerical differentiation
		    	J_dot_.data = (J_.data - J_last_[index].data)/period.toSec();

		    	// computing forward kinematics
		    	fk_pos_solver_->JntToCart(joint_msr_states_.q,x_,links_index_[index]);

		    	// setting marker parameters
		    	set_marker(x_,index,msg_id_);

		    	// computing end-effector position/orientation error w.r.t. desired frame
		    	x_err_ = diff(x_,x_des_[index]);

		    	x_dot_ = J_.data*joint_msr_states_.qdot.data;

		    	// setting error reference
		    	for(int i = 0; i < e_ref_.size(); i++)
 		    	{
		    		// e = -Kd*x_dot + Kp*(x_des - x)
		    		e_ref_(i) =  (-Kd_(i)*(x_dot_(i)) + Kp_(i)*x_err_(i));
	    			msg_err_.data.push_back(e_ref_(i));
		    	}

		    	// computing b = J*M^-1*(c+g) - J_dot*q_dot
		    	b_ = J_.data*M_inv_*(C_.data + G_.data) - J_dot_.data*joint_msr_states_.qdot.data;

		    	// computing omega = J*M^-1*N^T*J
		    	omega_ = J_.data*M_inv_*N_trans_*J_.data.transpose();

		    	// computing lambda = omega^-1
		    	pseudo_inverse(omega_,lambda_);

		    	// finally, computing the torque tau
		    	tau_.data = tau_.data + J_.data.transpose()*lambda_*(e_ref_ + b_ - J_.data*M_inv_*tau_.data);

		    	// stop condition
		    	if (!on_target_flag_[index])
		    	{
			    	if (Equal(x_,x_des_[index],0.01))
			    	{
			    		ROS_INFO("Task %d on target",index);
			    		on_target_flag_[index] = true;
			    		if (index == (ntasks_ - 1))
			    			cmd_flag_ = 0;
			    	}
			    }

		    	// updating N_^T
		    	N_trans_ = N_trans_ - J_.data.transpose()*lambda_*J_.data*M_inv_;

		    	// saving J_ of the last step
		    	J_last_[index] = J_;
		    }
    	}

    	// set controls for joints
    	for (int i = 0; i < joint_handles_.size(); i++)
    	{
    		if(cmd_flag_)
    			joint_handles_[i].setCommand(tau_(i));
    		else
    			joint_handles_[i].setCommand(PIDs_[i].computeCommand(joint_des_states_.q(i) - joint_msr_states_.q(i),period));
    	}

    	// publishing markers for visualization in rviz
    	pub_marker_.publish(msg_marker_);
    	msg_id_++;

	    // publishing error for all tasks as an array of ntasks*6
	    pub_error_.publish(msg_err_);
	    ros::spinOnce();

	}

	void MultiTaskPriorityInverseDynamics::command(const lwr_controllers::MultiPriorityTask::ConstPtr &msg)
	{
		if (msg->links.size() == msg->tasks.size()/6)
		{
			ntasks_ = msg->links.size();
			ROS_INFO("Number of tasks: %d",ntasks_);
			// Dynamically resize desired postures and links index when a message arrives
			x_des_.resize(ntasks_);
			links_index_.resize(ntasks_);
			on_target_flag_.resize(ntasks_);
			msg_marker_.markers.resize(ntasks_);
			J_last_.resize(ntasks_);
			for (int i = 0; i < ntasks_; i++)
				J_last_[i].resize(kdl_chain_.getNrOfJoints());
			msg_id_ = 0;

			for (int i = 0; i < ntasks_; i++)
			{
					if (msg->links[i] == -1)	// adjust index
						links_index_[i] = msg->links[i];
					else if (msg->links[i] >= 1 && msg->links[i] <=joint_handles_.size())
						links_index_[i] = msg->links[i] + 1;
					else
					{
						ROS_INFO("Links index must be within 1 and %ld. (-1 is end-effector)",joint_handles_.size());
						return;
					}

					x_des_[i] = KDL::Frame(
									KDL::Rotation::RPY(msg->tasks[i*6 + 3],
													   msg->tasks[i*6 + 4],
													   msg->tasks[i*6 + 5]),
									KDL::Vector(msg->tasks[i*6],
												msg->tasks[i*6 + 1],
												msg->tasks[i*6 + 2]));

					on_target_flag_[i] = false;
			}

			first_step_ = 1;
			cmd_flag_ = 1;
		}
		else
		{
			ROS_INFO("The number of links index and tasks must be the same");
			ROS_INFO("Tasks parameters are [x,y,x,roll,pitch,yaw]");
			return;
		}

	}

	void MultiTaskPriorityInverseDynamics::set_marker(KDL::Frame x, int index, int id)
	{
				sstr_.str("");
				sstr_.clear();

				if (links_index_[index] == -1)
					sstr_<<"end_effector";
				else
					sstr_<<"link_"<<(links_index_[index]-1);

				msg_marker_.markers[index].header.frame_id = "world";
				msg_marker_.markers[index].header.stamp = ros::Time();
				msg_marker_.markers[index].ns = sstr_.str();
				msg_marker_.markers[index].id = id;
				msg_marker_.markers[index].type = visualization_msgs::Marker::SPHERE;
				msg_marker_.markers[index].action = visualization_msgs::Marker::ADD;
				msg_marker_.markers[index].pose.position.x = x.p(0);
				msg_marker_.markers[index].pose.position.y = x.p(1);
				msg_marker_.markers[index].pose.position.z = x.p(2);
				msg_marker_.markers[index].pose.orientation.x = 0.0;
				msg_marker_.markers[index].pose.orientation.y = 0.0;
				msg_marker_.markers[index].pose.orientation.z = 0.0;
				msg_marker_.markers[index].pose.orientation.w = 1.0;
				msg_marker_.markers[index].scale.x = 0.01;
				msg_marker_.markers[index].scale.y = 0.01;
				msg_marker_.markers[index].scale.z = 0.01;
				msg_marker_.markers[index].color.a = 1.0;
				msg_marker_.markers[index].color.r = 0.0;
				msg_marker_.markers[index].color.g = 1.0;
				msg_marker_.markers[index].color.b = 0.0;
	}
}

PLUGINLIB_EXPORT_CLASS(lwr_controllers::MultiTaskPriorityInverseDynamics, controller_interface::ControllerBase)


////////////////////////////////////////////////////////////////////////////////
/// Copyright 2018-present Xinyan DAI<xinyan.dai@outlook.com>
///
/// permission is hereby granted, free of charge, to any person obtaining a copy
/// of this software and associated documentation files (the "Software"), to
/// deal in the Software without restriction, including without limitation the
/// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
/// sell copies of the Software, and to permit persons to whom the Software is
/// furnished to do so, subject to the following conditions:
///
/// The above copyright notice and this permission notice shall be included in
/// all copies or substantial portions ofthe Software.
///
/// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
/// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
/// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
/// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
/// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
/// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
/// IN THE SOFTWARE.

/// @version 0.1
/// @author  Xinyan DAI
/// @contact xinyan.dai@outlook.com
//////////////////////////////////////////////////////////////////////////////

#pragma once

#include <utility>
#include <vector>
#include <random>
#include <unordered_map>
#include <boost/progress.hpp>

#include "matrix.hpp"
#include "rptree.hpp"
#include "kmeans.hpp"
#include "../index.hpp"

namespace ss {

    using std::vector;
    template <typename DataType >
    vector<vector<vector<DataType > > >  GroupCenters(const Visitor<DataType> & data, const parameter& para) {
        KMeansIndex<DataType > kmeansIndex(para);
        // number of leaf node: Ks
        // size of leaf: L
        int Ks = para.kmeans_centers;
        int L = para.forest_size;
        kmeansIndex.reset(Ks * L);
        kmeansIndex.Iterate(data);
        const vector<vector<DataType > > & list_centers= kmeansIndex.get_centers();
        Matrix<DataType > centers(list_centers);

        parameter rpt_parameters = para;
        { // depth of rpt should be log2(K)
            rpt_parameters.num_bit = 0;
            int temp_Ks = Ks;
            while (temp_Ks >>= 1)
                ++rpt_parameters.num_bit;
        }
        RPTIndex<DataType > rpt(rpt_parameters);
        rpt.Train(centers);
        rpt.Add(centers);

        vector<vector<int > * > leafIdx = rpt.LeafIdx();
        vector<vector<vector<DataType > > > grouped_center(
                L, vector<vector<DataType > >(Ks, vector<DataType >(para.dim)));
        for (int l = 0; l < L; ++l) {
            for (int k = 0; k < Ks; ++k) {
                grouped_center[l][k] = list_centers[ (*(leafIdx[k]))[l] ];
            }
        }
        return grouped_center;
    }

    template<class DataType>
    class JointInvertedIndex : public Index<DataType > {
    protected:
        vector<KMeansIndex<DataType > * >   _joint_index;
    public:
        explicit JointInvertedIndex(const parameter& para) :
                Index<DataType >(para), _joint_index(para.forest_size) {
        }

        ~JointInvertedIndex() {
            for (KMeansIndex<DataType> * pointer : _joint_index) {
                if (pointer)
                    delete pointer;
            }
        }

        void Train(const Matrix<DataType > & data) override {

            vector<vector<vector<DataType > > > centers = ss::GroupCenters(
                    Visitor<DataType >(data, 0, data.getDim()), this->_para);
            for (int l = 0; l < this->_para.forest_size; ++l) {
                _joint_index[l] = new KMeansIndex<DataType >(this->_para);
                _joint_index[l]->set_centers(centers[l]);
            }
        }

        void Add(const Matrix<DataType > & data) override {
            for (int l = 0; l < this->_para.forest_size; ++l) {
                _joint_index[l]->Add(data);
            }
        }

        void Search(const DataType *query, const std::function<void (int)>& prober) override {
            return SearchByProbe(query, prober);
        }

        void SearchByProbe(const DataType *query, const std::function<void (int)>& prober) {
            for (int l = 0; l < this->_para.forest_size; ++l) {
                _joint_index[l]->Search(query, prober);
            }
        }

        void SearchByRanking(const DataType *query, const std::function<void (int)>& prober) {
            int Ks = this->_para.kmeans_centers;
            int L = this->_para.forest_size;
            std::vector<std::pair<float, int > > KxLdistances(L * Ks);
            for (int l = 0; l < L; ++l) {
                std::vector<std::pair<float, int > > dis = _joint_index[l]->ClusterDistance(query, this->_para.dim);
                for (int i = 0; i < Ks; ++i) {
                    KxLdistances[i + l * Ks].first = dis[i].first;
                    KxLdistances[i + l * Ks].second = dis[i].second + Ks * l;
                }
            }
            ss::SortPairByFirst(&KxLdistances);
            for (int i = 0; i < L; i++) {
                int l = KxLdistances[i].second / Ks;
                int k = KxLdistances[i].second % Ks;
                const vector<int > &points = this->_joint_index[l]->get_points()[k];
                for (int id : points) {
                    prober(id);
                }
            }
        }

    };

} // namespace ss


#include "world.h"

#include <glm/gtc/matrix_transform.hpp>
#include <boost/log/core.hpp>
#include <boost/log/trivial.hpp>

#include "adh/camera.h"

World::World(const std::shared_ptr<adh::Camera> & camera):
  _camera(camera),
  _deltaLook(0, 0),
  _look(0, 0),
  _deltaMove(0, 0),
  _position(0.0, 0.0, -1.0),
  _maxSpeed(5),
  _maxStrafe(5)
{
}

void World::onChangeView(const glm::vec2 & deltaLook)
{
  _deltaLook = deltaLook;
}

void World::onChangePosition(const glm::vec2 & deltaMove)
{
  _deltaMove = -deltaMove;
}

void World::run(float deltaT)
{
  _look += _deltaLook * deltaT;
  _deltaLook = glm::vec2(0, 0);

  // Attitude
  glm::mat4 base(glm::vec4(0.0f, 0.0f, 1.0f, 0.0f), // Front
                 glm::vec4(1.0f, 0.0f, 0.0f, 0.0f), // Right
                 glm::vec4(0.0f, 1.0f, 0.0f, 0.0f), // Up
                 glm::vec4(0.0f, 0.0f, 0.0f, 0.0f)); // Empty

  glm::mat4 heading_matrix = glm::rotate(glm::mat4(1.0f), -_look[0], glm::vec3(0.0f, 1.0f, 0.0f));
  glm::mat4 pitch_matrix = glm::rotate(glm::mat4(1.0f), -_look[1], glm::vec3(1.0f, 0.0f, 0.0f));

  glm::mat4 look = heading_matrix * pitch_matrix * base;

  // Position
  _position +=
    glm::vec3(look[0]) * _deltaMove[0] * deltaT * _maxSpeed +
    glm::vec3(look[1]) * _deltaMove[1] * deltaT * _maxStrafe;

  _camera->setViewMatrix(glm::lookAt(_position,
                                     _position + glm::vec3(look[0]),
                                     glm::vec3(look[2])));
}


#include <chrono>
#include <string>
#include <iostream>
#include <algorithm>
#include <pngwriter.h>
#include <boost/program_options.hpp>
#include "mandelbrot.hpp"

namespace po = boost::program_options;

int main(int argc, char* argv[]) {
    std::string out_file;
    double scale, cx, cy;
    size_t max_n, width, height, n_colors, n_threads, nx, ny;
    po::options_description desc("Allowed options");
    desc.add_options()
            ("help,h", "Print this message")
            ("out,o", po::value(&out_file)->default_value("result.png"), "Output file")
            ("scale,s", po::value(&scale)->default_value(1), "Scale")
            ("cx", po::value(&cx)->default_value(-0.28676842048), "X center coordinate")
            ("cy", po::value(&cy)->default_value(0), "Y center coordinate")
            ("iter,i", po::value(&max_n)->default_value(5000), "Maximum number of iterations")
            ("width", po::value(&width)->default_value(1920), "Image width")
            ("height", po::value(&height)->default_value(1080), "Image height")
            ("colors", po::value(&n_colors)->default_value(10), "Number of colors")
            ("threads,t", po::value(&n_threads)->default_value(1), "Number of threads")
            ("nx", po::value(&nx)->default_value(16), "Split into nx columns")
            ("ny", po::value(&ny)->default_value(16), "Split into ny rows");
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);
    if (vm.count("help")) {
        std::cout << desc << std::endl;
        return 0;
    }
    auto m = mtms::mandelbrot_set<double>(width, height, n_colors);
    const auto start = std::chrono::high_resolution_clock::now();
    m.run(n_threads, nx, ny, scale * std::max(3. / width, 2. / height), {cx, cy}, max_n);
    const auto run_end = std::chrono::high_resolution_clock::now();
    pngwriter png(static_cast<int>(width), static_cast<int>(height), 0, out_file.c_str());
    for (size_t i = 0; i < height; ++i)
        for (size_t j = 0; j < width; ++j) {
            const auto color = m.color(i, j);
            png.plot(static_cast<int>(j), static_cast<int>(i), color.r, color.g, color.b);
        }
    png.close();
    const auto write_end = std::chrono::high_resolution_clock::now();
    std::cout << "Run time: "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(run_end - start).count() << " ns\n"
              << "Write time : "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(write_end - run_end).count() << " ns\n"
              << "All time: "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(write_end - start).count() << " ns" << std::endl;
    return 0;
}


#define BOOST_TEST_MODULE "test_omp_3spn2_base_pair_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <test/util/utility.hpp>

#include <mjolnir/math/math.hpp>
#include <mjolnir/math/constants.hpp>
#include <mjolnir/util/make_unique.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/core/VerletList.hpp>
#include <mjolnir/omp/System.hpp>
#include <mjolnir/omp/ThreeSPN2BasePairInteraction.hpp>

BOOST_AUTO_TEST_CASE(omp_ThreeSPN2BasePairIntearction_basepair)
{
    namespace test = mjolnir::test;

    constexpr double tol = 1e-4;
    mjolnir::LoggerManager::set_default_logger(
            "test_omp_3spn2_base_pair_interaction.log");

    using seq_traits_type  = mjolnir::SimulatorTraits      <double, mjolnir::UnlimitedBoundary>;
    using omp_traits_type  = mjolnir::OpenMPSimulatorTraits<double, mjolnir::UnlimitedBoundary>;

    using real_type        = typename omp_traits_type::real_type;
    using coordinate_type  = typename omp_traits_type::coordinate_type;
    using boundary_type    = typename omp_traits_type::boundary_type;
    using topology_type    = mjolnir::Topology;

    using potential_type       = mjolnir::ThreeSPN2BasePairPotential<real_type>;

    using omp_system_type      = mjolnir::System<omp_traits_type>;
    using omp_interaction_type = mjolnir::ThreeSPN2BasePairInteraction<omp_traits_type>;

    using seq_system_type      = mjolnir::System<seq_traits_type>;
    using seq_interaction_type = mjolnir::ThreeSPN2BasePairInteraction<seq_traits_type>;

    using omp_parameter_list_type = mjolnir::ThreeSPN2BasePairParameterList<omp_traits_type>;
    using seq_parameter_list_type = mjolnir::ThreeSPN2BasePairParameterList<seq_traits_type>;
    using omp_parameter_type      = typename omp_parameter_list_type::parameter_type;
    using seq_parameter_type      = typename seq_parameter_list_type::parameter_type;

    using omp_partition_type      = mjolnir::VerletList<omp_traits_type, potential_type>;
    using seq_partition_type      = mjolnir::VerletList<seq_traits_type, potential_type>;

    // those actually does not depends on traits.
    using base_kind            = typename omp_parameter_list_type::base_kind;
    using ignore_group_type    = typename omp_parameter_list_type::ignore_group_type;
    using ignore_molecule_type = typename omp_parameter_list_type::ignore_molecule_type;

    constexpr real_type pi = mjolnir::math::constants<real_type>::pi();

    const int max_number_of_threads = omp_get_max_threads();
    BOOST_TEST_WARN(max_number_of_threads > 2);
    BOOST_TEST_MESSAGE("maximum number of threads = " << omp_get_max_threads());

    {
        using unit_type = mjolnir::unit::constants<real_type>;
        using phys_type = mjolnir::physics::constants<real_type>;
        const std::string energy = "kcal/mol";
        const std::string length = "angstrom";
        phys_type::set_kB(phys_type::kB() * (unit_type::J_to_cal() / 1000.0) *
                          unit_type::avogadro_constant());
        phys_type::set_eps0(phys_type::eps0() * (1000.0 / unit_type::J_to_cal()) /
                            unit_type::avogadro_constant());
        phys_type::set_energy_unit(energy);

        phys_type::set_eps0(phys_type::eps0() / unit_type::m_to_angstrom());

        phys_type::set_m_to_length(unit_type::m_to_angstrom());
        phys_type::set_length_to_m(unit_type::angstrom_to_m());

        phys_type::set_L_to_volume(1e-3 * std::pow(unit_type::m_to_angstrom(), 3));
        phys_type::set_volume_to_L(1e+3 * std::pow(unit_type::angstrom_to_m(), 3));

        phys_type::set_length_unit(length);
    }

    // Here, this test case checks base pairing interaction only.
    //
    //  theta1    theta2
    //       |    |
    // Si o  v    v o Sj
    //     \-.   .-/
    //      o === o
    //     Bi  r  Bj
    //
    // Required test cases are the combination of the following conditions
    //
    // BasePairing:
    // rij:
    //  1. (r < r0)
    //  2. (r0 < r)
    // theta1:
    //  1. theta < pi/2K
    //  2. pi/2K < theta < pi/K
    //  3. pi/K  < theta
    // theta2:
    //  1. theta < pi/2K
    //  2. pi/2K < theta < pi/K
    //  3. pi/K  < theta

    std::mt19937 rng(123456789);

    for(int num_thread=1; num_thread<=max_number_of_threads; ++num_thread)
    {
        omp_set_num_threads(num_thread);
        BOOST_TEST_MESSAGE("maximum number of threads = " << omp_get_max_threads());

        const std::vector<base_kind> bases{base_kind::A, base_kind::T};
        //  theta1    theta2
        //       |    |
        //  0 o  v    v o 2
        //     \-.   .-/
        //      o === o
        //     1       3
        //
        constexpr auto invalid = potential_type::invalid();

        omp_parameter_list_type omp_parameter_list(
            mjolnir::ThreeSPN2BasePairGlobalPotentialParameter<real_type>{}, {
                {1, omp_parameter_type{bases.at(0), 0}},
                {3, omp_parameter_type{bases.at(1), 2}}
            }, {}, ignore_molecule_type("Nothing"), ignore_group_type({})
        );
        seq_parameter_list_type seq_parameter_list(
            mjolnir::ThreeSPN2BasePairGlobalPotentialParameter<real_type>{}, {
                {1, seq_parameter_type{bases.at(0), 0}},
                {3, seq_parameter_type{bases.at(1), 2}}
            }, {}, ignore_molecule_type("Nothing"), ignore_group_type({})
        );

        omp_interaction_type omp_interaction(potential_type{},
            omp_parameter_list_type(omp_parameter_list),
            mjolnir::SpatialPartition<omp_traits_type, potential_type>(
                mjolnir::make_unique<omp_partition_type>()));

        seq_interaction_type seq_interaction(potential_type{},
            seq_parameter_list_type(seq_parameter_list),
            mjolnir::SpatialPartition<seq_traits_type, potential_type>(
                mjolnir::make_unique<seq_partition_type>()));

        omp_system_type omp_sys(4, boundary_type{});
        seq_system_type seq_sys(4, boundary_type{});

        test::clear_everything(omp_sys);
        test::clear_everything(seq_sys);

        topology_type topol(4);
        topol.add_connection(0, 1, "bond");
        topol.add_connection(2, 3, "bond");
        topol.construct_molecules();

        omp_interaction.initialize(omp_sys, topol);
        seq_interaction.initialize(seq_sys, topol);

        // paramter_lists in the interactions will be initialized via
        // interaction.initialize(), but parameter_list here will not.
        potential_type pot;
        omp_parameter_list.initialize(omp_sys, topol, pot);

        const auto bp_kind      = omp_parameter_list.bp_kind(bases.at(0), bases.at(1));
        const auto rbp_0        = omp_parameter_list.r0(bp_kind);
        const auto theta1_0     = omp_parameter_list.theta1_0(bp_kind);
        const auto theta2_0     = omp_parameter_list.theta2_0(bp_kind);
        const auto pi_over_K_BP = omp_parameter_list.pi_over_K_BP();

        for(const auto rbp  : {rbp_0 - 0.2, rbp_0 + 0.5})
        {
        for(const auto theta1 : {theta1_0 + pi_over_K_BP * 0.2,
                                 theta1_0 + pi_over_K_BP * 0.7,
                                 theta1_0 + pi_over_K_BP * 1.2,
                                 theta1_0 - pi_over_K_BP * 1.2,
                                 theta1_0 - pi_over_K_BP * 0.7,
                                 theta1_0 - pi_over_K_BP * 0.2})
        {
        for(const auto theta2 : {theta2_0 + pi_over_K_BP * 0.2,
                                 theta2_0 + pi_over_K_BP * 0.7,
                                 theta2_0 + pi_over_K_BP * 1.2,
                                 theta2_0 - pi_over_K_BP * 1.2,
                                 theta2_0 - pi_over_K_BP * 0.7,
                                 theta2_0 - pi_over_K_BP * 0.2})
        {
            BOOST_TEST_MESSAGE("===========================================");
            BOOST_TEST_MESSAGE("rbp    = " << rbp);
            BOOST_TEST_MESSAGE("theta1 = " << theta1);
            BOOST_TEST_MESSAGE("theta2 = " << theta2);

            //  theta1    theta2
            //       |    |
            //  0 o  v    v o 2
            //     \-.   .-/
            //      o === o
            //     1       3
            //
            omp_sys.position(1) = coordinate_type(0.0, 0.0, 0.0);
            omp_sys.position(3) = coordinate_type(rbp, 0.0, 0.0);

            omp_sys.position(0) = coordinate_type(std::cos(theta1),          std::sin(theta1),    0.0);
            omp_sys.position(2) = coordinate_type(std::cos(pi-theta2) + rbp, std::sin(pi-theta2), 0.0);

            test::apply_random_rotation(omp_sys, rng);

            // check configuration is okay
            {
                const auto v13 = omp_sys.position(3) - omp_sys.position(1);
                BOOST_TEST_REQUIRE(mjolnir::math::length(v13) == rbp,
                                   boost::test_tools::tolerance(1e-3));

                const auto v10 = omp_sys.position(0) - omp_sys.position(1);
                const auto cos_013 = mjolnir::math::dot_product(v13, v10) /
                    (mjolnir::math::length(v13) * mjolnir::math::length(v10));
                BOOST_TEST_REQUIRE(std::acos(cos_013) == theta1,
                                   boost::test_tools::tolerance(1e-3));

                const auto v23 = omp_sys.position(3) - omp_sys.position(2);
                const auto cos_132 = mjolnir::math::dot_product(v13, v23) /
                    (mjolnir::math::length(v13) * mjolnir::math::length(v23));
                BOOST_TEST_REQUIRE(std::acos(cos_132) == theta2,
                                   boost::test_tools::tolerance(1e-3));
            }

            test::apply_random_perturbation(omp_sys, rng, 0.01);

            for(std::size_t i=0; i<seq_sys.size(); ++i)
            {
                seq_sys.position(i) = omp_sys.position(i);
            }
            test::clear_force(omp_sys);
            test::clear_force(seq_sys);

            test::check_force_consistency              (omp_sys, omp_interaction, seq_sys, seq_interaction, tol);
            test::check_force_and_energy_consistency   (omp_sys, omp_interaction, seq_sys, seq_interaction, tol);
            test::check_force_and_virial_consistency   (omp_sys, omp_interaction, seq_sys, seq_interaction, tol);
            test::check_force_energy_virial_consistency(omp_sys, omp_interaction, seq_sys, seq_interaction, tol);
            test::check_energy_consistency             (omp_sys, omp_interaction, seq_sys, seq_interaction, tol);

        } // theta2
        } // theta1
        } // rbp
    }
}


#include <gtest/gtest.h>

#include <RKTools.h>
#include <EigenMatrixTypedefs.h>
#include <RKMatrixEigenTransformations.h>

namespace genfit {

    class RKMatrixEigenTransformations : public::testing::Test {
        protected:
    };

    TEST_F(RKMatrixEigenTransformations, VectorTransformation)
    {
        RKMatrix<1, 4> myVector4 = {{ 1, 2, 3, 4 }};

        auto myEigenMatrix(RKMatrixToEigenMatrix<1, 4>(myVector4));
        EXPECT_EQ(myVector4[0], myEigenMatrix[0]);
        EXPECT_EQ(myVector4[1], myEigenMatrix[1]);
        EXPECT_EQ(myVector4[2], myEigenMatrix[2]);
        EXPECT_EQ(myVector4[3], myEigenMatrix[3]);

        auto myRKMatrix(eigenMatrixToRKMatrix<1, 4>(myEigenMatrix));
        EXPECT_EQ(myVector4[0], myRKMatrix[0]);
        EXPECT_EQ(myVector4[1], myRKMatrix[1]);
        EXPECT_EQ(myVector4[2], myRKMatrix[2]);
        EXPECT_EQ(myVector4[3], myRKMatrix[3]);
    }

    TEST_F(RKMatrixEigenTransformations, MatrixTransformation)
    {
        RKMatrix<2, 3> myMatrix2x3 = {
                1, 2, 3,
                4, 5, 6,
        };

        auto myEigenMatrix(RKMatrixToEigenMatrix<2, 3>(myMatrix2x3));
        EXPECT_EQ(myMatrix2x3(0, 0), myEigenMatrix(0, 0));
        EXPECT_EQ(myMatrix2x3(0, 1), myEigenMatrix(0, 1));
        EXPECT_EQ(myMatrix2x3(0, 2), myEigenMatrix(0, 2));
        EXPECT_EQ(myMatrix2x3(1, 0), myEigenMatrix(1, 0));
        EXPECT_EQ(myMatrix2x3(1, 1), myEigenMatrix(1, 1));
        EXPECT_EQ(myMatrix2x3(1, 2), myEigenMatrix(1, 2));

        auto myRKMatrix(eigenMatrixToRKMatrix<2, 3>(myEigenMatrix));
        EXPECT_EQ(myMatrix2x3(0, 0), myRKMatrix(0, 0));
        EXPECT_EQ(myMatrix2x3(0, 1), myRKMatrix(0, 1));
        EXPECT_EQ(myMatrix2x3(0, 2), myRKMatrix(0, 2));
        EXPECT_EQ(myMatrix2x3(1, 0), myRKMatrix(1, 0));
        EXPECT_EQ(myMatrix2x3(1, 1), myRKMatrix(1, 1));
        EXPECT_EQ(myMatrix2x3(1, 2), myRKMatrix(1, 2));
    }

}