// Boost.Polygon library voronoi_builder_test.cpp file

//          Copyright Andrii Sydorchuk 2010-2012.
// 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)

// See http://www.boost.org for updates, documentation, and revision history.

#include "voronoi_test_helper.hpp"
#include <boost/core/lightweight_test.hpp>
#include <boost/polygon/polygon.hpp>
#include <boost/polygon/voronoi.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <limits>
#include <list>
#include <vector>
#include <ctime>

using boost::polygon::voronoi_builder;
using boost::polygon::voronoi_diagram;

typedef voronoi_diagram<double> vd_type;
typedef vd_type::coordinate_type coordinate_type;
typedef vd_type::edge_type voronoi_edge_type;
typedef vd_type::const_cell_iterator const_cell_iterator;
typedef vd_type::const_vertex_iterator const_vertex_iterator;

#define CHECK_OUTPUT_SIZE(output, cells, vertices, edges) \
    BOOST_TEST_EQ(output.num_cells(), (std::size_t)cells); \
    BOOST_TEST_EQ(output.num_vertices(), (std::size_t)vertices); \
    BOOST_TEST_EQ(output.num_edges(), (std::size_t)edges)

#define VERIFY_OUTPUT(output) \
    BOOST_TEST(voronoi_test_helper::verify_output(output, \
        voronoi_test_helper::CELL_CONVEXITY)); \
    BOOST_TEST(voronoi_test_helper::verify_output(output, \
        voronoi_test_helper::INCIDENT_EDGES_CCW_ORDER)); \
    BOOST_TEST(voronoi_test_helper::verify_output(output, \
        voronoi_test_helper::NO_HALF_EDGE_INTERSECTIONS))

#define VERIFY_NO_HALF_EDGE_INTERSECTIONS(output) \
    BOOST_TEST(voronoi_test_helper::verify_output(output, \
        voronoi_test_helper::NO_HALF_EDGE_INTERSECTIONS))

// Sites: (0, 0).
void single_site_test()
{
  std::vector< point_data<int> > points;
  points.push_back(point_data<int>(0, 0));
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);

  BOOST_TEST(test_output.cells().size() == 1);
  CHECK_OUTPUT_SIZE(test_output, 1, 0, 0);

  const_cell_iterator it = test_output.cells().begin();
  BOOST_TEST(it->incident_edge() == NULL);
}

// Sites: (0, 0), (0, 1).
void collinear_sites_test1()
{
  std::vector< point_data<int> > points;
  points.push_back(point_data<int>(0, 0));
  points.push_back(point_data<int>(0, 1));
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);
  CHECK_OUTPUT_SIZE(test_output, 2, 0, 2);

  const_cell_iterator cell_it = test_output.cells().begin();
  cell_it++;

  const voronoi_edge_type* edge1_1 = cell_it->incident_edge();
  const voronoi_edge_type* edge1_2 = edge1_1->twin();

  BOOST_TEST(edge1_1->twin() == edge1_2);
  BOOST_TEST(edge1_2->twin() == edge1_1);

  BOOST_TEST(edge1_1->next() == edge1_1);
  BOOST_TEST(edge1_1->prev() == edge1_1);
  BOOST_TEST(edge1_1->rot_next() == edge1_2);
  BOOST_TEST(edge1_1->rot_prev() == edge1_2);

  BOOST_TEST(edge1_2->next() == edge1_2);
  BOOST_TEST(edge1_2->prev() == edge1_2);
  BOOST_TEST(edge1_2->rot_next() == edge1_1);
  BOOST_TEST(edge1_2->rot_prev() == edge1_1);
}

// Sites: (0, 0), (1, 1), (2, 2).
void collinear_sites_test2()
{
  std::vector< point_data<int> > points;
  points.push_back(point_data<int>(0, 0));
  points.push_back(point_data<int>(1, 1));
  points.push_back(point_data<int>(2, 2));
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);
  CHECK_OUTPUT_SIZE(test_output, 3, 0, 4);

  const_cell_iterator cell_it = test_output.cells().begin();
  const voronoi_edge_type* edge1_1 = cell_it->incident_edge();
  const voronoi_edge_type* edge1_2 = edge1_1->twin();
  cell_it++;
  cell_it++;
  const voronoi_edge_type* edge2_2 = cell_it->incident_edge();
  const voronoi_edge_type* edge2_1 = edge2_2->twin();

  BOOST_TEST(edge1_1->twin() == edge1_2 && edge1_2->twin() == edge1_1);
  BOOST_TEST(edge2_1->twin() == edge2_2 && edge2_2->twin() == edge2_1);

  BOOST_TEST(edge1_1->next() == edge1_1 && edge1_1->prev() == edge1_1);
  BOOST_TEST(edge1_2->next() == edge2_1 && edge1_2->prev() == edge2_1);
  BOOST_TEST(edge2_1->next() == edge1_2 && edge2_1->prev() == edge1_2);
  BOOST_TEST(edge2_2->next() == edge2_2 && edge2_2->prev() == edge2_2);

  BOOST_TEST(edge1_1->rot_next() == edge1_2 && edge1_1->rot_prev() == edge2_1);
  BOOST_TEST(edge1_2->rot_next() == edge2_2 && edge1_2->rot_prev() == edge1_1);
  BOOST_TEST(edge2_1->rot_next() == edge1_1 && edge2_1->rot_prev() == edge2_2);
  BOOST_TEST(edge2_2->rot_next() == edge2_1 && edge2_2->rot_prev() == edge1_2);

  BOOST_TEST(edge1_2->next() == edge2_1 && edge1_2->prev() == edge2_1);
  BOOST_TEST(edge2_1->next() == edge1_2 && edge2_1->prev() == edge1_2);
}

// Sites: (0, 0), (0, 4), (2, 1).
void triangle_test1()
{
  point_data<int> point1(0, 0);
  point_data<int> point2(0, 4);
  point_data<int> point3(2, 1);
  std::vector< point_data<int> > points;
  points.push_back(point1);
  points.push_back(point2);
  points.push_back(point3);
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);
  CHECK_OUTPUT_SIZE(test_output, 3, 1, 6);

  const_vertex_iterator it = test_output.vertices().begin();
  BOOST_TEST_EQ(it->x(), 0.25);
  BOOST_TEST_EQ(it->y(), 2.0);

  const voronoi_edge_type* edge1_1 = it->incident_edge();
  const voronoi_edge_type* edge1_2 = edge1_1->twin();
  BOOST_TEST(edge1_1->cell()->source_index() == 1);
  BOOST_TEST(edge1_2->cell()->source_index() == 2);

  const voronoi_edge_type* edge2_1 = edge1_1->rot_prev();
  const voronoi_edge_type* edge2_2 = edge2_1->twin();
  BOOST_TEST(edge2_1->cell()->source_index() == 2);
  BOOST_TEST(edge2_2->cell()->source_index() == 0);

  const voronoi_edge_type* edge3_1 = edge2_1->rot_prev();
  const voronoi_edge_type* edge3_2 = edge3_1->twin();
  BOOST_TEST(edge3_1->cell()->source_index() == 0);
  BOOST_TEST(edge3_2->cell()->source_index() == 1);

  BOOST_TEST(edge1_2->twin() == edge1_1);
  BOOST_TEST(edge2_2->twin() == edge2_1);
  BOOST_TEST(edge3_2->twin() == edge3_1);

  BOOST_TEST(edge1_1->prev() == edge3_2 && edge1_1->next() == edge3_2);
  BOOST_TEST(edge2_1->prev() == edge1_2 && edge2_1->next() == edge1_2);
  BOOST_TEST(edge3_1->prev() == edge2_2 && edge3_1->next() == edge2_2);

  BOOST_TEST(edge1_2->next() == edge2_1 && edge1_2->prev() == edge2_1);
  BOOST_TEST(edge2_2->next() == edge3_1 && edge2_2->prev() == edge3_1);
  BOOST_TEST(edge3_2->next() == edge1_1 && edge3_2->prev() == edge1_1);

  BOOST_TEST(edge1_1->rot_next() == edge3_1);
  BOOST_TEST(edge3_1->rot_next() == edge2_1);
  BOOST_TEST(edge2_1->rot_next() == edge1_1);

  BOOST_TEST(edge1_2->rot_next() == edge2_2);
  BOOST_TEST(edge2_2->rot_next() == edge3_2);
  BOOST_TEST(edge3_2->rot_next() == edge1_2);
}

// Sites: (0, 1), (2, 0), (2, 4).
void triangle_test2()
{
  point_data<int> point1(0, 1);
  point_data<int> point2(2, 0);
  point_data<int> point3(2, 4);
  std::vector< point_data<int> > points;
  points.push_back(point1);
  points.push_back(point2);
  points.push_back(point3);
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);
  CHECK_OUTPUT_SIZE(test_output, 3, 1, 6);

  const_vertex_iterator it = test_output.vertices().begin();
  BOOST_TEST_EQ(it->x(), 1.75);
  BOOST_TEST_EQ(it->y(), 2.0);

  const voronoi_edge_type* edge1_1 = it->incident_edge();
  const voronoi_edge_type* edge1_2 = edge1_1->twin();
  BOOST_TEST(edge1_1->cell()->source_index() == 2);
  BOOST_TEST(edge1_2->cell()->source_index() == 1);

  const voronoi_edge_type* edge2_1 = edge1_1->rot_prev();
  const voronoi_edge_type* edge2_2 = edge2_1->twin();
  BOOST_TEST(edge2_1->cell()->source_index() == 1);
  BOOST_TEST(edge2_2->cell()->source_index() == 0);

  const voronoi_edge_type* edge3_1 = edge2_1->rot_prev();
  const voronoi_edge_type* edge3_2 = edge3_1->twin();
  BOOST_TEST(edge3_1->cell()->source_index() == 0);
  BOOST_TEST(edge3_2->cell()->source_index() == 2);

  BOOST_TEST(edge1_2->twin() == edge1_1);
  BOOST_TEST(edge2_2->twin() == edge2_1);
  BOOST_TEST(edge3_2->twin() == edge3_1);

  BOOST_TEST(edge1_1->prev() == edge3_2 && edge1_1->next() == edge3_2);
  BOOST_TEST(edge2_1->prev() == edge1_2 && edge2_1->next() == edge1_2);
  BOOST_TEST(edge3_1->prev() == edge2_2 && edge3_1->next() == edge2_2);

  BOOST_TEST(edge1_2->next() == edge2_1 && edge1_2->prev() == edge2_1);
  BOOST_TEST(edge2_2->next() == edge3_1 && edge2_2->prev() == edge3_1);
  BOOST_TEST(edge3_2->next() == edge1_1 && edge3_2->prev() == edge1_1);

  BOOST_TEST(edge1_1->rot_next() == edge3_1);
  BOOST_TEST(edge3_1->rot_next() == edge2_1);
  BOOST_TEST(edge2_1->rot_next() == edge1_1);

  BOOST_TEST(edge1_2->rot_next() == edge2_2);
  BOOST_TEST(edge2_2->rot_next() == edge3_2);
  BOOST_TEST(edge3_2->rot_next() == edge1_2);
}

// Sites: (0, 0), (0, 1), (1, 0), (1, 1).
void square_test1()
{
  point_data<int> point1(0, 0);
  point_data<int> point2(0, 1);
  point_data<int> point3(1, 0);
  point_data<int> point4(1, 1);
  std::vector< point_data<int> > points;
  points.push_back(point1);
  points.push_back(point2);
  points.push_back(point3);
  points.push_back(point4);
  vd_type test_output;
  construct_voronoi(points.begin(), points.end(), &test_output);
  VERIFY_OUTPUT(test_output);
  CHECK_OUTPUT_SIZE(test_output, 4, 1, 8);

  // Check voronoi vertex.
  const_vertex_iterator it = test_output.vertices().begin();
  BOOST_TEST_EQ(it->x(), 0.5);
  BOOST_TEST_EQ(it->y(), 0.5);

  // Check voronoi edges.
  const voronoi_edge_type* edge1_1 = it->incident_edge();
  const voronoi_edge_type* edge1_2 = edge1_1->twin();
  BOOST_TEST(edge1_1->cell()->source_index() == 3);
  BOOST_TEST(edge1_2->cell()->source_index() == 2);

  const voronoi_edge_type* edge2_1 = edge1_1->rot_prev();
  const voronoi_edge_type* edge2_2 = edge2_1->twin();
  BOOST_TEST(edge2_1->cell()->source_index() == 2);
  BOOST_TEST(edge2_2->cell()->source_index() == 0);

  const voronoi_edge_type* edge3_1 = edge2_1->rot_prev();
  const voronoi_edge_type* edge3_2 = edge3_1->twin();
  BOOST_TEST(edge3_1->cell()->source_index() == 0);
  BOOST_TEST(edge3_2->cell()->source_index() == 1);

  const voronoi_edge_type* edge4_1 = edge3_1->rot_prev();
  const voronoi_edge_type* edge4_2 = edge4_1->twin();
  BOOST_TEST(edge4_1->cell()->source_index() == 1);
  BOOST_TEST(edge4_2->cell()->source_index() == 3);

  BOOST_TEST(edge1_2->twin() == edge1_1);
  BOOST_TEST(edge2_2->twin() == edge2_1);
  BOOST_TEST(edge3_2->twin() == edge3_1);
  BOOST_TEST(edge4_2->twin() == edge4_1);

  BOOST_TEST(edge1_1->prev() == edge4_2 && edge1_1->next() == edge4_2);
  BOOST_TEST(edge2_1->prev() == edge1_2 && edge2_1->next() == edge1_2);
  BOOST_TEST(edge3_1->prev() == edge2_2 && edge3_1->next() == edge2_2);
  BOOST_TEST(edge4_1->prev() == edge3_2 && edge4_1->next() == edge3_2);

  BOOST_TEST(edge1_2->next() == edge2_1 && edge1_2->prev() == edge2_1);
  BOOST_TEST(edge2_2->next() == edge3_1 && edge2_2->prev() == edge3_1);
  BOOST_TEST(edge3_2->next() == edge4_1 && edge3_2->prev() == edge4_1);
  BOOST_TEST(edge4_2->next() == edge1_1 && edge4_2->prev() == edge1_1);

  BOOST_TEST(edge1_1->rot_next() == edge4_1);
  BOOST_TEST(edge4_1->rot_next() == edge3_1);
  BOOST_TEST(edge3_1->rot_next() == edge2_1);
  BOOST_TEST(edge2_1->rot_next() == edge1_1);

  BOOST_TEST(edge1_2->rot_next() == edge2_2);
  BOOST_TEST(edge2_2->rot_next() == edge3_2);
  BOOST_TEST(edge3_2->rot_next() == edge4_2);
  BOOST_TEST(edge4_2->rot_next() == edge1_2);
}

#ifdef NDEBUG
void grid_test()
{
  vd_type test_output_small, test_output_large;
  std::vector< point_data<int> > point_vec_small, point_vec_large;
  int grid_size[] = {10, 33, 101};
  int max_value[] = {10, 33, 101};
  int array_length = sizeof(grid_size) / sizeof(int);
  for (int k = 0; k < array_length; k++) {
    test_output_small.clear();
    test_output_large.clear();
    point_vec_small.clear();
    point_vec_large.clear();
    int koef = (std::numeric_limits<int>::max)() / max_value[k];
    for (int i = 0; i < grid_size[k]; i++) {
      for (int j = 0; j < grid_size[k]; j++) {
        point_vec_small.push_back(point_data<int>(i, j));
        point_vec_large.push_back(point_data<int>(koef * i, koef * j));
      }
    }
    construct_voronoi(point_vec_small.begin(), point_vec_small.end(), &test_output_small);
    construct_voronoi(point_vec_large.begin(), point_vec_large.end(), &test_output_large);
    VERIFY_OUTPUT(test_output_small);
    VERIFY_OUTPUT(test_output_large);
    unsigned int num_cells = grid_size[k] * grid_size[k];
    unsigned int num_vertices = num_cells - 2 * grid_size[k] + 1;
    unsigned int num_edges = 4 * num_cells - 4 * grid_size[k];
    CHECK_OUTPUT_SIZE(test_output_small, num_cells, num_vertices, num_edges);
    CHECK_OUTPUT_SIZE(test_output_large, num_cells, num_vertices, num_edges);
  }
}
#endif

#ifdef NDEBUG
void random_test()
{
  boost::mt19937 gen(static_cast<unsigned int>(time(NULL)));
  vd_type test_output_small, test_output_large;
  std::vector< point_data<int> > point_vec_small, point_vec_large;
  int num_points[] = {10, 100, 1000, 10000};
  int num_runs[] = {1000, 100, 10, 1};
  int mod_koef[] = {10, 100, 100, 1000};
  int max_value[] = {5, 50, 50, 5000};
  int array_length = sizeof(num_points) / sizeof(int);
  for (int k = 0; k < array_length; k++) {
    int koef = (std::numeric_limits<int>::max)() / max_value[k];
    for (int i = 0; i < num_runs[k]; i++) {
      test_output_small.clear();
      test_output_large.clear();
      point_vec_small.clear();
      point_vec_large.clear();
      for (int j = 0; j < num_points[k]; j++) {
        T x = gen() % mod_koef[k] - mod_koef[k] / 2;
        T y = gen() % mod_koef[k] - mod_koef[k] / 2;
        point_vec_small.push_back(point_data<int>(x, y));
        point_vec_large.push_back(point_data<int>(koef * x, koef * y));
      }
      construct_voronoi(point_vec_small.begin(), point_vec_small.end(), &test_output_small);
      construct_voronoi(point_vec_large.begin(), point_vec_large.end(), &test_output_large);
      VERIFY_OUTPUT(test_output_small);
      VERIFY_OUTPUT(test_output_large);
      BOOST_TEST_EQ(test_output_small.num_cells(), test_output_large.num_cells());
      BOOST_TEST_EQ(test_output_small.num_vertices(), test_output_large.num_vertices());
      BOOST_TEST_EQ(test_output_small.num_edges(), test_output_large.num_edges());
    }
  }
}
#endif

void segment_sites_test1()
{
  vd_type test_output;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(1, 1);
  segments.push_back(segment_data<int>(point1, point2));
  construct_voronoi(segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 3, 0, 4);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_sites_test2()
{
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(4, 4);
  point_data<int> point3(3, 1);
  point_data<int> point4(1, 3);
  segments.push_back(segment_data<int>(point1, point2));
  points.push_back(point3);
  points.push_back(point4);
  construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 5, 4, 16);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_sites_test3()
{
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(4, 0);
  point_data<int> point2(0, 4);
  point_data<int> point3(3, 3);
  point_data<int> point4(1, 1);
  segments.push_back(segment_data<int>(point1, point2));
  points.push_back(point3);
  points.push_back(point4);
  construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 5, 4, 16);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_sites_test4()
{
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(4, 0);
  point_data<int> point2(0, 4);
  point_data<int> point3(3, 2);
  point_data<int> point4(2, 3);
  segments.push_back(segment_data<int>(point1, point2));
  points.push_back(point3);
  points.push_back(point4);
  construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 5, 3, 14);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_site_test5()
{
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(0, 8);
  point_data<int> point3(-2, -2);
  point_data<int> point4(-2, 4);
  point_data<int> point5(-2, 10);
  segments.push_back(segment_data<int>(point1, point2));
  points.push_back(point3);
  points.push_back(point4);
  points.push_back(point5);
  construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 6, 4, 18);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_site_test6()
{
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(-1, 1);
  point_data<int> point2(1, 0);
  point_data<int> point3(1, 2);
  segments.push_back(segment_data<int>(point2, point3));
  points.push_back(point1);
  construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 4, 2, 10);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_site_test7()
{
  vd_type test_output;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(4, 0);
  point_data<int> point3(0, 4);
  point_data<int> point4(4, 4);
  segments.push_back(segment_data<int>(point1, point2));
  segments.push_back(segment_data<int>(point2, point3));
  segments.push_back(segment_data<int>(point3, point4));
  construct_voronoi(segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 7, 6, 24);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_site_test8()
{
  vd_type test_output;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(4, 0);
  point_data<int> point3(4, 4);
  point_data<int> point4(0, 4);
  segments.push_back(segment_data<int>(point1, point2));
  segments.push_back(segment_data<int>(point2, point3));
  segments.push_back(segment_data<int>(point3, point4));
  segments.push_back(segment_data<int>(point4, point1));
  construct_voronoi(segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 8, 5, 24);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

void segment_site_test9()
{
  vd_type test_output;
  std::vector< segment_data<int> > segments;
  point_data<int> point1(0, 0);
  point_data<int> point2(2, 0);
  point_data<int> point3(4, 0);
  segments.push_back(segment_data<int>(point1, point2));
  segments.push_back(segment_data<int>(point2, point3));
  construct_voronoi(segments.begin(), segments.end(), &test_output);
  CHECK_OUTPUT_SIZE(test_output, 5, 0, 8);
  VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
}

#ifdef NDEBUG
void segment_grid_test()
{
  vd_type test_output_small, test_output_large;
  std::vector< segment_data<int> > segments_small, segments_large;
  int grid_size[] = {10, 27, 53};
  int max_value[] = {100, 330, 1000};
  int array_length = sizeof(grid_size) / sizeof(int);
  for (int k = 0; k < array_length; k++) {
    test_output_small.clear();
    test_output_large.clear();
    segments_small.clear();
    segments_large.clear();
    int cur_sz = grid_size[k];
    int koef = (std::numeric_limits<int>::max)() / max_value[k];
    for (int i = 0; i < cur_sz + 1; i++)
      for (int j = 0; j < cur_sz; j++) {
        point_data<int> point1_1(10 * i, 10 * j);
        point_data<int> point1_2(koef * 10 * i, koef * 10 * j);
        point_data<int> point2_1(10 * i, 10 * j + 10);
        point_data<int> point2_2(koef * 10 * i, koef * (10 * j + 10));
        segments_small.push_back(segment_data<int>(point1_1, point2_1));
        segments_large.push_back(segment_data<int>(point1_2, point2_2));
        point_data<int> point3_1(10 * j, 10 * i);
        point_data<int> point3_2(koef * 10 * j, koef * 10 * i);
        point_data<int> point4_1(10 * j + 10, 10 * i);
        point_data<int> point4_2(koef * (10 * j + 10), koef * 10 * i);
        segments_small.push_back(segment_data<int>(point3_1, point4_1));
        segments_large.push_back(segment_data<int>(point3_2, point4_2));
      }
    construct_voronoi(segments_small.begin(), segments_small.end(), &test_output_small);
    construct_voronoi(segments_large.begin(), segments_large.end(), &test_output_large);
    VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output_small);
    VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output_large);
    BOOST_TEST_EQ(test_output_small.num_cells(), test_output_large.num_cells());
    BOOST_TEST_EQ(test_output_small.num_vertices(), test_output_large.num_vertices());
    BOOST_TEST_EQ(test_output_small.num_edges(), test_output_large.num_edges());
  }
}
#endif

#ifdef NDEBUG
void segment_random_test1()
{
  boost::mt19937 gen(static_cast<unsigned int>(time(NULL)));
  vd_type test_output;
  std::vector< point_data<int> > points;
  std::vector< segment_data<int> > segments;
  int num_runs = 1000;
  int num_segments = 10;
  points.push_back(point_data<int>(-100, -100));
  points.push_back(point_data<int>(-100, 100));
  points.push_back(point_data<int>(100, -100));
  points.push_back(point_data<int>(100, 100));
  for (int i = 0; i < num_runs; i++) {
    test_output.clear();
    segments.clear();
    for (int j = 0; j < num_segments; j++) {
      T x1 = 0, y1 = 0, x2 = 0, y2 = 0;
      while (x1 == x2 && y1 == y2) {
        x1 = (gen() % 100) - 50;
        y1 = (gen() % 100) - 50;
        x2 = (gen() % 100) - 50;
        y2 = (gen() % 100) - 50;
      }
      point_data<int> point1(x1, y1);
      point_data<int> point2(x2, y2);
      segments.push_back(segment_data<int>(point1, point2));
    }
    voronoi_test_helper::clean_segment_set(segments);
    construct_voronoi(points.begin(), points.end(), segments.begin(), segments.end(), &test_output);
    VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output);
  }
}
#endif

#ifdef NDEBUG
void segment_random_test2()
{
  boost::mt19937 gen(static_cast<unsigned int>(time(NULL)));
  vd_type test_output_small, test_output_large;
  std::vector< segment_data<int> > segments_small, segments_large;
  int num_segments[] = {5, 25, 125, 625};
  int num_runs[] = {1000, 100, 10, 1};
  int mod_koef1[] = {10, 100, 200, 300};
  int mod_koef2[] = {10, 20, 50, 100};
  int max_value[] = {10, 60, 125, 200};
  int array_length = sizeof(num_segments) / sizeof(int);
  for (int k = 0; k < array_length; k++) {
    int koef = (std::numeric_limits<int>::max)() / max_value[k];
    for (int i = 0; i < num_runs[k]; i++) {
      test_output_small.clear();
      test_output_large.clear();
      segments_small.clear();
      segments_large.clear();
      for (int j = 0; j < num_segments[k]; j++) {
        T x1 = (gen() % mod_koef1[k]) - mod_koef1[k] / 2;
        T y1 = (gen() % mod_koef1[k]) - mod_koef1[k] / 2;
        T dx = 0, dy = 0;
        while (dx == 0 && dy == 0) {
          dx = (gen() % mod_koef2[k]) - mod_koef2[k] / 2;
          dy = (gen() % mod_koef2[k]) - mod_koef2[k] / 2;
        }
        T x2 = x1 + dx;
        T y2 = y1 + dy;
        point_data<int> point1_small(x1, y1);
        point_data<int> point2_small(x2, y2);
        segments_small.push_back(segment_data<int>(point1_small, point2_small));
      }
      voronoi_test_helper::clean_segment_set(segments_small);
      for (typename std::vector< segment_data<int> >::iterator it = segments_small.begin();
           it != segments_small.end(); ++it) {
        T x1 = it->low().x() * koef;
        T y1 = it->low().y() * koef;
        T x2 = it->high().x() * koef;
        T y2 = it->high().y() * koef;
        point_data<int> point1_large(x1, y1);
        point_data<int> point2_large(x2, y2);
        segments_large.push_back(segment_data<int>(point1_large, point2_large));
      }
      construct_voronoi(segments_small.begin(), segments_small.end(), &test_output_small);
      construct_voronoi(segments_large.begin(), segments_large.end(), &test_output_large);
      VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output_small);
      VERIFY_NO_HALF_EDGE_INTERSECTIONS(test_output_large);
      BOOST_TEST_EQ(test_output_small.num_cells(), test_output_large.num_cells());
      BOOST_TEST_EQ(test_output_small.num_vertices(), test_output_large.num_vertices());
      BOOST_TEST_EQ(test_output_small.num_edges(), test_output_large.num_edges());
    }
  }
}
#endif

int main()
{
    single_site_test();
    collinear_sites_test1();
    collinear_sites_test2();
    triangle_test1();
    triangle_test2();
    square_test1();
#ifdef NDEBUG
    grid_test();
    random_test();
#endif
    segment_sites_test1();
    segment_sites_test2();
    segment_sites_test3();
    segment_sites_test4();
    segment_site_test5();
    segment_site_test6();
    segment_site_test7();
    segment_site_test8();
    segment_site_test9();
#ifdef NDEBUG
    segment_grid_test();
    segment_random_test1();
    segment_random_test2();
#endif
    return boost::report_errors();
}


/*
 * pyramid.cpp
 *
 *  Created on: Apr 10, 2019
 *      Author: Gregory Kramida
 *   Copyright: 2019 Gregory Kramida
 *
 *   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.
 */

//libraries
#include <Eigen/Dense>

//local
#include "pyramid.tpp"
#include "../../math/typedefs.hpp"

namespace eig = Eigen;

namespace nonrigid_optimization {
namespace hierarchical{

template class Pyramid<eig::MatrixXf>;
template class Pyramid<math::MatrixXv2f>;
template class Pyramid<math::Tensor3f>;
template class Pyramid<math::Tensor3v3f>;

} //namespace hierarchical
} //namespace nonrigid_optimization


#include <iostream>
#include <queue>
#include <string>

#include <boost/random.hpp>
#include <boost/generator_iterator.hpp>

#include <glog/logging.h>

using boost::variate_generator;
using boost::mt19937;
using boost::exponential_distribution;

#define ONLY_EVENT_TYPE     0
#define NUMBER_EVENT_TYPES  2   //NEED TO MAKE SURE THIS IS 1 MORE THAN THE LAST DEFINED EVENT

const std::string event_names[NUMBER_EVENT_TYPES] = {
    "ARRIVE",
    "DEPART"
};

/**
 *  A class for an event, which holds an int for the event type, a double for simulation time and possibly other data.
 *  We may want to subclass this with our own events
 */
class Event {
    public:
        const double time;
        const double server_start_time;
        const int type;

        Event(double time, int type) : time(time), server_start_time(0), type(type) {
            //cout << "created an event with simulation time: " << this->time << endl;
            //cout << "created an event with event type: " << this->type << endl;
        }

        Event(double time, double server_start_time, int type) : time(time), server_start_time(server_start_time), type(type) {
            //cout << "created an event with simulation time: " << this->time << endl;
            //cout << "created an event with event type: " << this->type << endl;
        }

        bool operator<( const Event& e ) const {
            return time < e.time;
        }

        bool operator>( const Event& e ) const {
            return time > e.time;
        }

        bool less( const Event& e) const {
            return time < e.time;
        }

        friend std::ostream& operator<< (std::ostream& out, Event& event);
};

class CompareEvent {
    public:
        bool operator()(Event* &e1, Event* &e2) {
            return e1->time > e2->time;
            return false;
        }
};

// Print event
std::ostream& operator<< ( std::ostream& out, Event& event) {
    out << "[sim_time: " << std::setw(10) << std::setprecision(3) << std::fixed << event.time << ", type: " << std::setw(4) << event.type;
    if (event.type < NUMBER_EVENT_TYPES) {
        out << " - " << std::left << std::setw(50) << event_names[event.type];
    } else {
        out << std::left << std::setw(50) << " - UNKNOWN";
    }
    out << std::right << "]";
    return out;
}

int main(int argc, char **argv) {
    // Initialize Google Logging
    google::InitGoogleLogging(argv[0]);
    // Log to Stderr
    FLAGS_logtostderr = 1;

    int seed = time(0);
    double distribution_mean = 3.0;
    double simulation_time_s = 0;
    double end_simulation_time_s = atoi(argv[1]);
    double previous_time_s;

    bool server_idle = true;
    double server_work_time = 0;

    double total_queue_time = 0;
    double total_departures = 0;

    double sum_of_time_queue_length = 0;

    variate_generator< mt19937, exponential_distribution<> > rand_generator(mt19937(seed), exponential_distribution<>(1/distribution_mean));

    std::priority_queue<Event*, std::vector<Event*>, CompareEvent> heap;
    std::queue<double> queue;

    // Put initial event in the heap
    heap.push(new Event(simulation_time_s + rand_generator(), 0));

    while (simulation_time_s  < end_simulation_time_s) {
        Event *current_event = heap.top();
        heap.pop();

        if (current_event == NULL) {
            LOG(ERROR) << "Simulation not complete and there are no events in the min heap.";
            LOG(ERROR) << "\tsimulation time: " << simulation_time_s;
            exit(0);
        }

        previous_time_s = simulation_time_s;
        simulation_time_s = current_event->time;

        // Calculate average time stuff
        double time_since_last = simulation_time_s - previous_time_s;
        sum_of_time_queue_length += queue.size() * time_since_last;

        switch (current_event->type) {
            case 0: // ARRIVE
                // If the server is busy add new arrival to waiting queue
                // otherwise set the server as busy and add new departure time.
                if (!server_idle) {
                    queue.push(simulation_time_s);
                } else {
                    server_idle = false;
                    heap.push(new Event(simulation_time_s + rand_generator(), simulation_time_s, 1));
                }
                // Add next arrival event
                heap.push(new Event(simulation_time_s + rand_generator(), 0));
                break;
            case 1: // DEPART
                // Add server busy time.
                server_work_time += simulation_time_s - current_event->server_start_time;
                // If the queue is empty set the server to idle otherwise
                // get the next person from the queue and set their departure
                // time.
                if (queue.empty()) {
                    server_idle = true;
                } else {
                    // Add current simulation time minus time stored in the
                    // queue to the total queue time.
                    total_queue_time += simulation_time_s - queue.front();
                    queue.pop();
                    // Add a new depature to the heap.
                    heap.push(new Event(simulation_time_s + rand_generator(), simulation_time_s, 1));
                }
                total_departures++;
                break;
            default:
                LOG(ERROR) << "Simulation had an event with an unknown type: " << current_event->type;
                LOG(ERROR) << "\tsimulation time: " << simulation_time_s;
                exit(0);
        }

        VLOG(2) << *current_event << ", h: " << heap.size() << ", q: " << queue.size();

        delete current_event; // Event's are created with new, so we need to delete them when we're done with them
    }

    LOG(INFO) << "The simulation ended at time: " << end_simulation_time_s;
    VLOG(1) << "The server was busy for time: " << server_work_time;
    LOG(INFO) << "The server utilization was: " << server_work_time / end_simulation_time_s;
    LOG(INFO) << "The average length of the queue was: " << sum_of_time_queue_length / end_simulation_time_s;
    VLOG(1) << "Total time spent in the queue: " << total_queue_time;
    VLOG(1) << "Total departures: " << total_departures;
    LOG(INFO) << "The average time spent in queue: " << total_queue_time / total_departures;
}


/*
 [auto_generated]
 libs/numeric/odeint/test_external/eigen/runge_kutta_dopri5.cpp

 [begin_description]
 tba.
 [end_description]

 Copyright 2009-2012 Karsten Ahnert
 Copyright 2009-2012 Mario Mulansky

 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/config.hpp>
#ifdef BOOST_MSVC
    #pragma warning(disable:4996)
#endif

#define BOOST_TEST_MODULE odeint_eigen_runge_kutta4

#include <boost/test/unit_test.hpp>

#include <boost/numeric/odeint/algebra/vector_space_algebra.hpp>
#include <boost/numeric/odeint/stepper/runge_kutta_dopri5.hpp>
#include <boost/numeric/odeint/stepper/runge_kutta4.hpp>
#include <boost/numeric/odeint/stepper/controlled_runge_kutta.hpp>
#include <boost/numeric/odeint/stepper/dense_output_runge_kutta.hpp>

// #include <boost/numeric/odeint/external/eigen/eigen_resize.hpp>
#include <boost/numeric/odeint/external/eigen/eigen_algebra.hpp>

using namespace boost::unit_test;
using namespace boost::numeric::odeint;

struct sys
{
    template< class State , class Deriv >
    void operator()( const State &x , Deriv &dxdt , double t ) const
    {
        dxdt[0] = 1.0;
    }
};

template< class State >
struct stepper
{
    typedef runge_kutta_dopri5< State , double , State , double , vector_space_algebra > type;
};

template< class State >
struct controlled_stepper
{
    typedef controlled_runge_kutta< typename stepper< State >::type > type;
};

template< class State >
struct dense_output_stepper
{
    typedef dense_output_runge_kutta< typename controlled_stepper< State >::type > type;
};

BOOST_AUTO_TEST_SUITE( eigen_runge_kutta_dopri5 )

BOOST_AUTO_TEST_CASE( compile_time_matrix )
{
    typedef Eigen::Matrix< double , 1 , 1 > state_type;
    state_type x;
    x[0] = 10.0;
    double t = 0.0 , dt = 0.1 ;

    // dense_output_stepper< state_type >::type s;
    // s.initialize( x , t , dt );

    // controlled_stepper< state_type >::type s;
    // s.try_step( sys() , x , t , dt );

    stepper< state_type >::type s;
    s.do_step( sys() , x , t , dt );

    // runge_kutta4< state_type , double , state_type , double , vector_space_algebra > rk4;
    // rk4.do_step( sys() , x , 0.0 , 0.1 );
    // BOOST_CHECK_CLOSE( x[0] , 10.1 , 1.0e-13 );

}

BOOST_AUTO_TEST_CASE( runtime_matrix )
{
    typedef Eigen::Matrix< double , Eigen::Dynamic , 1 > state_type;
    state_type x( 1 );
    x[0] = 10.0;

    // runge_kutta4< state_type , double , state_type , double , vector_space_algebra > rk4;
    // rk4.do_step( sys() , x , 0.0 , 0.1 );
    // BOOST_CHECK_CLOSE( x[0] , 10.1 , 1.0e-13 );
}

BOOST_AUTO_TEST_CASE( compile_time_array )
{
    typedef Eigen::Array< double , 1 , 1 > state_type;
    state_type x;
    x[0] = 10.0;
    // runge_kutta4< state_type , double , state_type , double , vector_space_algebra > rk4;
    // rk4.do_step( sys() , x , 0.0 , 0.1 );
    // BOOST_CHECK_CLOSE( x[0] , 10.1 , 1.0e-13 );

}

BOOST_AUTO_TEST_CASE( runtime_array )
{
    typedef Eigen::Array< double , Eigen::Dynamic , 1 > state_type;
    state_type x( 1 );
    x[0] = 10.0;
    // runge_kutta4< state_type , double , state_type , double , vector_space_algebra > rk4;
    // rk4.do_step( sys() , x , 0.0 , 0.1 );
    // BOOST_CHECK_CLOSE( x[0] , 10.1 , 1.0e-13 );
}

BOOST_AUTO_TEST_SUITE_END()


#include <algorithm>
#include <vector>
#include <cmath>

#include "google/protobuf/descriptor.h"
#include "caffe/layer.hpp"
#include "caffe/layers/data_augmentation_layer.hpp"
#include "caffe/util/rng.hpp"
#include "caffe/util/math_functions.hpp"

#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

#include <boost/random.hpp>
#include <boost/random/normal_distribution.hpp>

#include <iostream>
#include <fstream>
#include <omp.h>

using std::max;

namespace caffe {

template <typename T> int sgn(T val) {
  return (T(0) < val) - (val < T(0));
}

template <typename Dtype>
Dtype caffe_rng_generate(const RandomGeneratorParameter& param, Dtype discount_coeff = 1) {
  float spread;
  if (param.apply_schedule())
    spread = param.spread() * discount_coeff;
  else
    spread = param.spread();
  const std::string rand_type =  param.rand_type();
  //std::cout << rand_type << " " << rand_type.compare("uniform") << " " << rand_type.compare("gaussian") << " " << rand_type.compare("bernoulli");
  Dtype rand;
  if (rand_type.compare("uniform") == 0) {
    float tmp;
    if (spread > 0.)
      caffe_rng_uniform(1, param.mean() - spread, param.mean() + spread, &tmp);
    else
      tmp = param.mean();
    if (param.exp())
      tmp = exp(tmp);
    rand = static_cast<Dtype>(tmp);
  }
  else if (rand_type.compare("gaussian") == 0) {
    float tmp;
    if (spread > 0.)
      caffe_rng_gaussian(1, param.mean(), spread, &tmp);
    else
      tmp = param.mean();
    if (param.exp())
      tmp = exp(tmp);
    rand = static_cast<Dtype>(tmp);
  }
  else if (rand_type.compare("bernoulli") == 0) {
    int tmp;
    if (param.prob() > 0.)
      caffe_rng_bernoulli(1, param.prob(), &tmp);
    else
      tmp = 0;
    rand = static_cast<Dtype>(tmp);
  }
  else if (rand_type.compare("uniform_bernoulli") == 0) {
    float tmp1;
    int tmp2;

    if (spread > 0.)
      caffe_rng_uniform(1, param.mean() - spread, param.mean() + spread, &tmp1);
    else
      tmp1 = param.mean();

    if (param.prob() > 0.)
      caffe_rng_bernoulli(1, param.prob(), &tmp2);
    else
      tmp2 = 0;

    tmp1 = tmp1 * static_cast<float>(tmp2);

    if (param.exp())
      tmp1 = exp(tmp1);

    rand = static_cast<Dtype>(tmp1);
  }
  else if (rand_type.compare("gaussian_bernoulli") == 0) {
    float tmp1;
    int tmp2;

    if (spread > 0.)
      caffe_rng_gaussian(1, param.mean(), spread, &tmp1);
    else
      tmp1 = param.mean();

    if (param.prob() > 0.)
      caffe_rng_bernoulli(1, param.prob(), &tmp2);
    else
      tmp2 = 0;

    tmp1 = tmp1 * static_cast<float>(tmp2);

    if (param.exp())
      tmp1 = exp(tmp1);

    rand = static_cast<Dtype>(tmp1);
  }
  else {
    LOG(ERROR) << "Unknown random type " << rand_type;
    rand = NAN;
  }
  return rand;
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  CHECK_GE(bottom.size(), 1) << "Data aumentation layer takes one or two input blobs.";
  CHECK_LE(bottom.size(), 2) << "Data aumentation layer takes one or two input blobs.";
  CHECK_GE(top.size(), 1) << "Data Layer takes one or two output blobs.";
  CHECK_LE(top.size(), 2) << "Data Layer takes one or two output blobs.";

  if (top.size() == 1) {
    output_params_ = false;
  } else {
    output_params_ = true;
  }

  if (bottom.size() == 1) {
    input_params_ = false;
  } else {
    input_params_ = true;
  }

  if (output_params_) {
    AugmentationCoeff coeff;
    num_params_ = coeff.GetDescriptor()->field_count();
    LOG(INFO) << "Writing " << num_params_ << " augmentation params";
  }
  if (input_params_) {
    AugmentationCoeff coeff;
    num_params_ = coeff.GetDescriptor()->field_count();
    LOG(INFO) << "Reading " << num_params_ << " augmentation params";
  }

  num_iter_ = 0;
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
  const int num = bottom[0]->num();
  const int channels = bottom[0]->channels();
  const int height = bottom[0]->height();
  const int width = bottom[0]->width();

  aug_ = this->layer_param_.augmentation_param();

  discount_coeff_schedule_ = this->layer_param_.coeff_schedule_param();

  if (!aug_.has_crop_size())
    LOG(ERROR) << "Please enter crop_size if you want to perform augmentation";
  int crop_size = aug_.crop_size();
  CHECK_GE(height, crop_size) << "crop size greater than original";
  CHECK_GE(width, crop_size) << "crop size greater than original";

  cropped_height_ = crop_size;
  cropped_width_ = crop_size;

  (top)[0]->Reshape(num, channels, crop_size, crop_size);

  if (aug_.recompute_mean()) {
    data_mean_.Reshape(1, channels, crop_size, crop_size);
  }

  if (output_params_) {
    (top)[1]->Reshape(num, num_params_, 1, 1);
  }
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {
   num_iter_++;

  const int num = bottom[0]->num();
  const int channels = bottom[0]->channels();
  const int height = bottom[0]->height();
  const int width = bottom[0]->width();

  const Dtype* bottom_data = bottom[0]->cpu_data();
  Dtype* top_data = (top)[0]->mutable_cpu_data();

  std::string write_augmented;
  if (aug_.has_write_augmented())
    write_augmented = aug_.write_augmented();
  else
    write_augmented = std::string("");

  bool augment_during_test = aug_.augment_during_test();
  bool train_phase = (this->phase_ == TRAIN);
  //LOG(INFO) <<  " === train_phase " << train_phase;

  Dtype mean_eig [3];
  Dtype mean_rgb[3] = {0., 0., 0.};
  Dtype max_abs_eig[3] = {0., 0., 0.};
  Dtype max_rgb[3] = {0., 0., 0.};
  Dtype min_rgb[3] = {0., 0., 0.};
  Dtype max_l;
  Dtype rgb [3];
  Dtype eig [3];
//   const Dtype eigvec [9] = {0.5579, 0.5859, 0.5878, 0.8021, -0.1989, -0.5631, -0.2130, 0.7856, -0.5809};
//   const Dtype eigvec [9] = {1., 0., 0., 0., 1., 0., 0., 0., 1.};
//   const Dtype eigvec [9] = {0.57, 0.58, 0.57, -0.72, 0.03, 0.68, -0.38, 0.81, -0.44};
//  const Dtype eigvec [9] = {0.5878, 0.5859, 0.5579, -0.5631, -0.1989, 0.8021, -0.5809, 0.7856, -0.2130};
  const Dtype eigvec [9] = {0.51, 0.56, 0.65, 0.79, 0.01, -0.62, 0.35, -0.83, 0.44};
  if (channels == 3) {
    for (int item_id = 0; item_id < num; ++item_id) {
      for (int x=0; x < width; x++) {
        for (int y=0; y < height; y++) {
          for (int c=0; c<channels; c++) {
            rgb[c] = bottom_data[((item_id*channels + c)*width + x)*height + y];
          }
          //if (x==0 && y ==0)
          //    LOG(INFO) << "item_id=" << item_id << ", x=" << x << ", y=" << y << ", rgb[0]=" << rgb[0] << ", rgb[1]=" << rgb[1] << ", rgb[2]=" << rgb[2];
          for (int c=0; c<channels; c++) {
            eig[c] = eigvec[3*c] * rgb[0] + eigvec[3*c+1] * rgb[1] + eigvec[3*c+2] * rgb[2];
            if (fabs(eig[c]) > max_abs_eig[c])
              max_abs_eig[c] = fabs(eig[c]);
            if (rgb[c] > max_rgb[c])
              max_rgb[c] = rgb[c];
            if (rgb[c] < min_rgb[c])
              min_rgb[c] = rgb[c];
            mean_rgb[c] = mean_rgb[c] + rgb[c]/width/height;
          }
        }
      }
    }
    for (int c=0; c<channels; c++)
      mean_rgb[c] = mean_rgb[c] / num;

    for (int c=0; c<channels; c++) {
      mean_eig[c] = eigvec[3*c] * mean_rgb[0] + eigvec[3*c+1] * mean_rgb[1] + eigvec[3*c+2] * mean_rgb[2];
      if ( max_abs_eig[c] > 1e-2 )
        mean_eig[c] = mean_eig[c] / max_abs_eig[c];
    }
    max_l = sqrt(max_abs_eig[0]*max_abs_eig[0] + max_abs_eig[1]*max_abs_eig[1] + max_abs_eig[2]*max_abs_eig[2]);
  }

  //LOG(INFO) << "mean_rgb[0]=" << mean_rgb[0] << ", mean_rgb[1]=" << mean_rgb[1] << ", mean_rgb[2]=" << mean_rgb[2];
  //LOG(INFO) << "max_abs_eig[0]=" << max_abs_eig[0];

  AugmentationParameter aug = aug_;
  Dtype* top_params;
  Dtype* bottom_params;
  if (output_params_)
    top_params = (top)[1]->mutable_cpu_data();
  if (input_params_)
    bottom_params = bottom[1]->mutable_cpu_data();
  bool output_params = output_params_;
  bool input_params = input_params_;
  int num_params = num_params_;

  Dtype discount_coeff = discount_coeff_schedule_.initial_coeff() +
      (discount_coeff_schedule_.final_coeff() - discount_coeff_schedule_.initial_coeff()) * (Dtype(2) /
      (Dtype(1) + exp(-discount_coeff_schedule_.gamma() * num_iter_)) - Dtype(1));

//   LOG(INFO) << "num_iter=" << num_iter_ << ", discount_coeff=" << discount_coeff;

#pragma omp parallel for shared(aug, train_phase, write_augmented, augment_during_test, mean_rgb, mean_eig, max_abs_eig, max_rgb, min_rgb, max_l, eigvec,  output_params, input_params, num_params, discount_coeff)  private(rgb, eig)
  for (int item_id = 0; item_id < num; ++item_id) {
    int x, y, c, top_idx, bottom_idx, h_off, w_off;
    Dtype x1, y1, x2, y2;
    bool do_spatial_transform, do_chromatic_transform;

    AugmentationCoeff coeff;

    //   We only do transformations during training or if specifically asked to do them during testing.
    if (!(train_phase || aug.augment_during_test())) {
      do_spatial_transform   = false;
      do_chromatic_transform = false;
    }
    else {
      if (input_params) {
        array_to_coeff(bottom_params + item_id * num_params, coeff);
      }
      do_spatial_transform   = true;
      do_chromatic_transform = true;
    }

      // sample the parameters of the transformations
    if (do_spatial_transform) {
      int counter = 0;
      int max_num_tries = 20;
      int good_params = 0;

      // try to sample parameters for which transformed image doesn't go outside the borders of the original one
      // in order to check this, just apply the transformations to 4 corners
      while (good_params < 4 && counter < max_num_tries) {
        good_params = 0;
        generate_spatial_coeffs(aug, coeff, discount_coeff);

        //LOG(INFO) << "angle: " << angle << ", zoom: " << zoom_coeff << ", dx: " << dx << ", dy: " << dy << ", mirror: " << mirror;

        for (x = 0; x < cropped_width_; x += cropped_width_-1) {
          for (y = 0; y < cropped_height_; y += cropped_height_-1) {
            // move the origin and mirror
            if (coeff.mirror()) {
              x1 =  static_cast<Dtype>(x) - .5 * static_cast<Dtype>(cropped_width_);
              y1 = -static_cast<Dtype>(y) + .5 * static_cast<Dtype>(cropped_height_);
            }
            else {
              x1 = static_cast<Dtype>(x) - .5 * static_cast<Dtype>(cropped_width_);
              y1 = static_cast<Dtype>(y) - .5 * static_cast<Dtype>(cropped_height_);
            }
            // rotate
            x2 =  cos(coeff.angle()) * x1 - sin(coeff.angle()) * y1;
            y2 =  sin(coeff.angle()) * x1 + cos(coeff.angle()) * y1;
            // translate
            x2 = x2 + coeff.dx() * static_cast<Dtype>(cropped_width_);
            y2 = y2 + coeff.dy() * static_cast<Dtype>(cropped_height_);
            // zoom
            x2 = x2 / coeff.zoom_x();
            y2 = y2 / coeff.zoom_y();
            // move the origin back
            x2 = x2 + .5 * static_cast<Dtype>(width);
            y2 = y2 + .5 * static_cast<Dtype>(height);

            if (!(floor(x2) < 0 || floor(x2) > static_cast<Dtype>(width - 2) || floor(y2) < 0 || floor(y2) > static_cast<Dtype>(height - 2)))
                good_params++;
            //mexPrintf(" (%f,%f) ", x2, y2);
          }
        }
        //mexPrintf("\n");
        counter++;
      }
      if (counter >= max_num_tries) {
        clear_spatial_coeffs(coeff);
        do_spatial_transform = false;
      }
      else {
        clear_defaults(coeff);
        do_spatial_transform  = coeff.has_mirror() || coeff.has_dx() || coeff.has_dy() || coeff.has_angle() || coeff.has_zoom_x() || coeff.has_zoom_y();
      }
    }

    if (do_chromatic_transform) {
      generate_chromatic_coeffs(aug, coeff, discount_coeff);
      clear_defaults(coeff);
      do_chromatic_transform =  coeff.has_pow_nomean0()     || coeff.has_pow_nomean1()     || coeff.has_pow_nomean2()    ||
                                coeff.has_add_nomean0()     || coeff.has_add_nomean1()     || coeff.has_add_nomean2()    ||
                                coeff.has_mult_nomean0()    || coeff.has_mult_nomean1()    || coeff.has_mult_nomean2()   ||
                                coeff.has_pow_withmean0()   || coeff.has_pow_withmean1()   || coeff.has_pow_withmean2()  ||
                                coeff.has_add_withmean0()   || coeff.has_add_withmean1()   || coeff.has_add_withmean2()  ||
                                coeff.has_mult_withmean0()  || coeff.has_mult_withmean1()  || coeff.has_mult_withmean2() ||
                                coeff.has_lmult_pow()       || coeff.has_lmult_add()       || coeff.has_lmult_mult()     ||
                                coeff.has_col_angle();

    }

    // Augment chromatically only 3-channel inputs
    if (do_chromatic_transform) {
      CHECK_EQ(channels, 3) << "Augment chromatically only 3-channel inputs";
    }

    if (write_augmented.size()) {
      if (do_spatial_transform)
        LOG(INFO) << "Augmenting " << item_id
                  << ", mirror: "  << coeff.mirror()  << ", angle: " << coeff.angle()
                  << ", zoom_x: "  << coeff.zoom_x()  << ", zoom_y: "  << coeff.zoom_y()
                  << ", dx: "      << coeff.dx()      << ", dy: " << coeff.dy() ;
      else
        LOG(INFO) << "Not augmenting " << item_id << " spatially";
      if (do_chromatic_transform)
        LOG(INFO) << "Augmenting "   << item_id
                  << ", pow_nm0: "   << coeff.pow_nomean0()   << ", mult_nm0: "   << coeff.mult_nomean0()   << ", add_nm0: "   << coeff.add_nomean0()
                  << ", pow_nm1: "   << coeff.pow_nomean1()   << ", mult_nm1: "   << coeff.mult_nomean1()   << ", add_nm1: "   << coeff.add_nomean1()
                  << ", pow_nm2: "   << coeff.pow_nomean2()   << ", mult_nm2: "   << coeff.mult_nomean2()   << ", add_nm2: "   << coeff.add_nomean2()
                  << ", pow_wm0: "   << coeff.pow_withmean0() << ", mult_wm0: "   << coeff.mult_withmean0() << ", add_wm0: "   << coeff.add_withmean0()
                  << ", pow_wm1: "   << coeff.pow_withmean1() << ", mult_wm1: "   << coeff.mult_withmean1() << ", add_wm1: "   << coeff.add_withmean1()
                  << ", pow_wm2: "   << coeff.pow_withmean2() << ", mult_wm2: "   << coeff.mult_withmean2() << ", add_wm2: "   << coeff.add_withmean2()
                  << ". lmult_pow: " << coeff.lmult_pow()     << ", lmult_mult: " <<  coeff.lmult_mult()    << ", lmult_add: " << coeff.lmult_add()
                  << ", col_angle: " << coeff.col_angle();
      else
        LOG(INFO) << "Not augmenting " << item_id << " chromatically";
    }

    if (output_params) {
//       LOG(INFO) << "coeff_to_array item_id=" << item_id << ", num_params=" << num_params;
      coeff_to_array(coeff, top_params + item_id * num_params);
    }

    // actually apply the transformation
    if (do_spatial_transform) {
      int i00,i01,i10,i11;
      for (x = 0; x < cropped_width_; x++) {
        for (y = 0; y < cropped_height_; y++) {
          // move the origin and mirror
          if (coeff.mirror()) {
            x1 =  static_cast<Dtype>(x) - .5 * static_cast<Dtype>(cropped_width_);
            y1 = -static_cast<Dtype>(y) + .5 * static_cast<Dtype>(cropped_height_);
          }
          else {
            x1 = static_cast<Dtype>(x) - .5 * static_cast<Dtype>(cropped_width_);
            y1 = static_cast<Dtype>(y) - .5 * static_cast<Dtype>(cropped_height_);
          }
          // rotate
          if (coeff.has_angle()) {
            x2 =  cos(coeff.angle()) * x1 - sin(coeff.angle()) * y1;
            y2 =  sin(coeff.angle()) * x1 + cos(coeff.angle()) * y1;
          }
          else {
            x2 = x1;
            y2 = y1;
          }
          // translate
          if (coeff.has_dx())
            x2 = x2 + coeff.dx() * static_cast<Dtype>(cropped_width_);
          if (coeff.has_dy())
            y2 = y2 + coeff.dy() * static_cast<Dtype>(cropped_height_);
          // zoom
          if (coeff.has_zoom_x())
            x2 = x2 / coeff.zoom_x();
          if (coeff.has_zoom_y())
            y2 = y2 / coeff.zoom_y();
          // move the origin back
          x2 = x2 + .5 * static_cast<Dtype>(width);
          y2 = y2 + .5 * static_cast<Dtype>(height);

          for (c = 0; c < channels; c++) {
            top_idx = ((item_id*channels + c)*cropped_width_ + x)*cropped_height_ + y;
            if (floor(x2) < 0. || floor(x2) > static_cast<Dtype>(width - 2) || floor(y2) < 0. || floor(y2) > static_cast<Dtype>(height - 2))
              top_data[top_idx] = 0.;
            else {
              if (coeff.has_angle() || coeff.has_zoom_x() || coeff.has_zoom_y()) {
                i00 = static_cast<int>(((item_id*channels + c) * width +  floor(x2)) *height + floor(y2));
                i01 = i00 + 1;
                i10 = i00 + height;
                i11 = i00 + height + 1;

                top_data[top_idx] = bottom_data[i00] * ((floor(x2)+1)  - x2) * ((floor(y2)+1)  - y2) +
                                    bottom_data[i01] * ((floor(x2)+1)  - x2) * (y2 - floor(y2))      +
                                    bottom_data[i10] * (x2 - floor(x2))      * ((floor(y2)+1)  - y2) +
                                    bottom_data[i11] * (x2 - floor(x2))      * (y2 - floor(y2));
              }
              else {
                i00 = static_cast<int>(((item_id*channels + c) * width +  floor(x2)) *height + floor(y2));
                top_data[top_idx] = bottom_data[i00];
              }
            }
            // TODO Alexey check this
            //top_data[i] = (top_data[i] - 127.5) * scale;
          }
          //mexPrintf(" (%f,%f) ", x2, y2);
        }
      }
    }
    else {
      h_off = (height - cropped_height_)/2;
      w_off = (width - cropped_width_)/2;
      for (x = 0; x < cropped_width_; x++) {
        for (y = 0; y < cropped_height_; y++) {
          for (c = 0; c < channels; c++) {
            top_idx = ((item_id*channels + c)*cropped_width_ + x)*cropped_height_ + y;
            bottom_idx = ((item_id*channels + c)*width + x + w_off)*height + y + h_off;
            top_data[top_idx] = bottom_data[bottom_idx];
          }
        }
      }
    }

    if (do_chromatic_transform) {
//       LOG(INFO) << " >>> do chromatic transform " << item_id;
      Dtype s, s1, l, l1;

      for (x=0; x<cropped_width_; x++) {
        for (y=0; y<cropped_height_; y++) {
          // subtracting the mean
          for (c=0; c<channels; c++) {
            rgb[c] = top_data[((item_id*channels + c)*cropped_width_ + x)*cropped_height_ + y];
            rgb[c] = rgb[c] - mean_rgb[c];
          }
          // doing the nomean stuff
          for (c=0; c<channels; c++) {
            eig[c] = eigvec[3*c] * rgb[0] + eigvec[3*c+1] * rgb[1] + eigvec[3*c+2] * rgb[2];
            if ( max_abs_eig[c] > 1e-2 ) {
              eig[c] = eig[c] / max_abs_eig[c];
              if (c==0) {
                if (coeff.has_pow_nomean0())
                  eig[c] = static_cast<float>(sgn(eig[c])) * pow(fabs(eig[c]), coeff.pow_nomean0());
                if (coeff.has_add_nomean0())
                  eig[c] = eig[c] + coeff.add_nomean0();
                if (coeff.has_mult_nomean0())
                  eig[c] = eig[c] * coeff.mult_nomean0();
              }
              else if (c==1) {
                if (coeff.has_pow_nomean1())
                  eig[c] = static_cast<float>(sgn(eig[c])) * pow(fabs(eig[c]), coeff.pow_nomean1());
                if (coeff.has_add_nomean1())
                  eig[c] = eig[c] + coeff.add_nomean1();
                if (coeff.has_mult_nomean1())
                  eig[c] = eig[c] * coeff.mult_nomean1();
              } else if (c==2) {
                if (coeff.has_pow_nomean2())
                  eig[c] = static_cast<float>(sgn(eig[c])) * pow(fabs(eig[c]), coeff.pow_nomean2());
                if (coeff.has_add_nomean2())
                  eig[c] = eig[c] + coeff.add_nomean2();
                if (coeff.has_mult_nomean2())
                  eig[c] = eig[c] * coeff.mult_nomean2();
              }
            }
          }
          // re-adding the mean
          for (c=0; c<channels; c++)
            eig[c] = eig[c] + mean_eig[c];

          // doing the withmean stuff
          if ( max_abs_eig[c] > 1e-2 && (coeff.has_pow_withmean0() || coeff.has_add_withmean0() || coeff.has_mult_withmean0())) {
            if (coeff.has_pow_withmean0())
              eig[0] = static_cast<float>(sgn(eig[0])) * pow(fabs(eig[0]), coeff.pow_withmean0());
            if (coeff.has_add_withmean0())
              eig[0] = eig[0] + coeff.add_withmean0();
            if (coeff.has_mult_withmean0())
              eig[0] = eig[0] * coeff.mult_withmean0();
          }
          if (coeff.has_pow_withmean1() || coeff.has_add_withmean1() || coeff.has_mult_withmean1()) {
            s = sqrt(eig[1]*eig[1] + eig[2]*eig[2]);
            s1 = s;
            if (s > 1e-2) {
              if (coeff.has_pow_withmean1())
                s1 = pow(s1, coeff.pow_withmean1());
              if (coeff.has_add_withmean1())
                s1 = fmax(s1 + coeff.add_withmean1(), 0.);
              if (coeff.has_mult_withmean1())
                s1 = s1 * coeff.mult_withmean1();
            }
          }
          if (coeff.has_col_angle()) {
            Dtype temp1, temp2;
            temp1 =  cos(coeff.col_angle()) * eig[1] - sin(coeff.col_angle()) * eig[2];
            temp2 =  sin(coeff.col_angle()) * eig[1] + cos(coeff.col_angle()) * eig[2];
            eig[1] = temp1;
            eig[2] = temp2;
          }
          for (c=0; c<channels; c++) {
            if ( max_abs_eig[c] > 1e-2 ) {
              eig[c] = eig[c] * max_abs_eig[c];
            }
          }
          if (max_l > 1e-2 && (coeff.has_lmult_pow() || coeff.has_lmult_add() || coeff.has_lmult_mult()) ||
                (coeff.has_pow_withmean1() || coeff.has_add_withmean1() || coeff.has_mult_withmean1())) {
            l1 = sqrt(eig[0]*eig[0] + eig[1]*eig[1] + eig[2]*eig[2]);
            l1 = l1 / max_l;
          }
          if (s > 1e-2 && (coeff.has_pow_withmean1() || coeff.has_add_withmean1() || coeff.has_mult_withmean1())) {
            eig[1] = eig[1] / s * s1;
            eig[2] = eig[2] / s * s1;
          }
          if ( max_l > 1e-2 && (coeff.has_lmult_pow() || coeff.has_lmult_add() || coeff.has_lmult_mult()) ||
                (coeff.has_pow_withmean1() || coeff.has_add_withmean1() || coeff.has_mult_withmean1())) {
            l = sqrt(eig[0]*eig[0] + eig[1]*eig[1] + eig[2]*eig[2]);
            if (coeff.has_lmult_pow())
              l1 = pow(l1, coeff.lmult_pow());
            if (coeff.has_lmult_add())
              l1 = fmax(l1 + coeff.lmult_add(), 0.);
            if (coeff.has_lmult_mult())
              l1 = l1 * coeff.lmult_mult();
            l1 = l1 * max_l;
            if (l > 1e-2)
              for (c=0; c<channels; c++) {
                eig[c] = eig[c] / l * l1;
                if (eig[c] > max_abs_eig[c])
                  eig[c] = max_abs_eig[c];
              }
          }
          for (c=0; c<channels; c++) {
            rgb[c] = eigvec[c] * eig[0] + eigvec[3+c] * eig[1] + eigvec[6+c] * eig[2];
//             if (rgb[c] > aug.max_multiplier()*max_rgb[c])
//               rgb[c] = aug.max_multiplier()*max_rgb[c];
//             if (rgb[c] < aug.max_multiplier()*min_rgb[c])
//               rgb[c] = aug.max_multiplier()*min_rgb[c];
            if (rgb[c] > aug.max_multiplier())
              rgb[c] = aug.max_multiplier();
            if (rgb[c] < 0.)
              rgb[c] = 0.;
            top_data[((item_id*channels + c)*cropped_width_ + x)*cropped_height_ + y] = rgb[c];
          }
        }
      }
    }
  }

  if(aug.recompute_mean() > 0 && num_iter_ <= aug.recompute_mean() ) {
    Dtype* data_mean_cpu = data_mean_.mutable_cpu_data();
    int count = cropped_width_*cropped_height_*channels;
    for (int c = 0; c < count; ++c) {
      data_mean_cpu[c] = data_mean_cpu[c]*(static_cast<Dtype>(num_iter_)-1);
      for (int item_id = 0; item_id < num; ++item_id)
        data_mean_cpu[c] = data_mean_cpu[c] + top_data[item_id*count + c] / num;
      data_mean_cpu[c] = data_mean_cpu[c] / static_cast<Dtype>(num_iter_);
    }
  }

  if(aug.recompute_mean() > 0) {
    Dtype* data_mean_cpu = data_mean_.mutable_cpu_data();
    int count = cropped_width_*cropped_height_*channels;
    for (int item_id = 0; item_id < num; ++item_id) {
      for (int c = 0; c < count; ++c) {
        top_data[item_id*count + c] = top_data[item_id*count + c] - data_mean_cpu[c];
      }
    }
  }

  if (write_augmented.size()) {
    std::ofstream out_file (write_augmented.data(), std::ios::out | std::ios::binary);
    if (out_file.is_open()) {
      uint32_t imsize[4];
      imsize[0] = num;
      imsize[1] = channels;
      imsize[2] = cropped_width_;
      imsize[3] = cropped_height_;
      LOG(INFO) << "Writing blob size " << imsize[0] << "x" << imsize[1] << "x" << imsize[2] << "x" << imsize[3];
      out_file.write(reinterpret_cast<char*>(&imsize[0]), 4*4);
      out_file.write(reinterpret_cast<const char*>(top_data), imsize[0]*imsize[1]*imsize[2]*imsize[3]*sizeof(float));
      out_file.close();
      LOG(INFO) << " finished augmenting a batch. train=" << train_phase << " === PAUSED === ";
      std::cout << " finished augmenting a batch. train=" << train_phase << " === PAUSED === ";
      std::cin.get();
    }
    else
      LOG(INFO) << "WARNING: Could not open the file" << write_augmented;
  }

  if (aug.write_mean().size()) {
    std::ofstream out_file (aug.write_mean().data(), std::ios::out | std::ios::binary);
    if (out_file.is_open()) {
      uint32_t imsize[4];
      imsize[0] = 1;
      imsize[1] = channels;
      imsize[2] = cropped_width_;
      imsize[3] = cropped_height_;
      Dtype* data_mean_cpu = data_mean_.mutable_cpu_data();
      LOG(INFO) << "Writing blob size " << imsize[0] << "x" << imsize[1] << "x" << imsize[2] << "x" << imsize[3];
      out_file.write(reinterpret_cast<char*>(&imsize[0]), 4*4);
      out_file.write(reinterpret_cast<const char*>(data_mean_cpu), imsize[0]*imsize[1]*imsize[2]*imsize[3]*sizeof(float));
      out_file.close();
      LOG(INFO) << " finished writing the mean. num_iter_=" << num_iter_ << " === PAUSED === ";
      std::cout << " finished writing the mean. num_iter_=" << num_iter_ << " === PAUSED === ";
      std::cin.get();
    }
    else
      LOG(INFO) << "WARNING: Could not open the file" << write_augmented;
  }
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::generate_spatial_coeffs(const AugmentationParameter& aug, AugmentationCoeff& coeff, Dtype discount_coeff) {
  if (aug.has_mirror())
    coeff.set_mirror(static_cast<float>(caffe_rng_generate<bool>(aug.mirror())));
  if (aug.has_translate()) {
    coeff.set_dx(caffe_rng_generate<float>(aug.translate(), discount_coeff));
    coeff.set_dy(caffe_rng_generate<float>(aug.translate(), discount_coeff));
  }
  if (aug.has_rotate())
    coeff.set_angle(caffe_rng_generate<float>(aug.rotate(), discount_coeff));
  if (aug.has_zoom()) {
    coeff.set_zoom_x(caffe_rng_generate<float>(aug.zoom(), discount_coeff));
    coeff.set_zoom_y(coeff.zoom_x());
  }
  if (aug.has_squeeze()) {
    float squeeze_coeff = caffe_rng_generate<float>(aug.squeeze(), discount_coeff);
    coeff.set_zoom_x(coeff.zoom_x() * squeeze_coeff);
    coeff.set_zoom_y(coeff.zoom_y() / squeeze_coeff);
  }
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::clear_spatial_coeffs(AugmentationCoeff& coeff) {
  coeff.clear_mirror();
  coeff.clear_dx();
  coeff.clear_dy();
  coeff.clear_angle();
  coeff.clear_zoom_x();
  coeff.clear_zoom_y();
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::generate_chromatic_coeffs(const AugmentationParameter& aug, AugmentationCoeff& coeff, Dtype discount_coeff) {
  if (aug.has_ladd_pow())
    coeff.set_pow_nomean0(caffe_rng_generate<float>(aug.ladd_pow(), discount_coeff));
  if (aug.has_col_pow()) {
    coeff.set_pow_nomean1(caffe_rng_generate<float>(aug.col_pow(), discount_coeff));
    coeff.set_pow_nomean2(caffe_rng_generate<float>(aug.col_pow(), discount_coeff));
  }

  if (aug.has_ladd_add())
    coeff.set_add_nomean0(caffe_rng_generate<float>(aug.ladd_add(), discount_coeff));
  if (aug.has_col_add()) {
    coeff.set_add_nomean1(caffe_rng_generate<float>(aug.col_add(), discount_coeff));
    coeff.set_add_nomean2(caffe_rng_generate<float>(aug.col_add(), discount_coeff));
  }

  if (aug.has_ladd_mult())
    coeff.set_mult_nomean0(caffe_rng_generate<float>(aug.ladd_mult(), discount_coeff));
  if (aug.has_col_mult()) {
    coeff.set_mult_nomean1(caffe_rng_generate<float>(aug.col_mult(), discount_coeff));
    coeff.set_mult_nomean2(caffe_rng_generate<float>(aug.col_mult(), discount_coeff));
  }

  if (aug.has_sat_pow()) {
    coeff.set_pow_withmean1(caffe_rng_generate<float>(aug.sat_pow(), discount_coeff));
    coeff.set_pow_withmean2(coeff.pow_withmean1());
  }

  if (aug.has_sat_add()) {
    coeff.set_add_withmean1(caffe_rng_generate<float>(aug.sat_add(), discount_coeff));
    coeff.set_add_withmean2(coeff.add_withmean1());
  }

  if (aug.has_sat_mult()) {
    coeff.set_mult_withmean1(caffe_rng_generate<float>(aug.sat_mult(), discount_coeff));
    coeff.set_mult_withmean2(coeff.mult_withmean1());
  }

  if (aug.has_lmult_pow())
    coeff.set_lmult_pow(caffe_rng_generate<float>(aug.lmult_pow(), discount_coeff));
  if (aug.has_lmult_mult())
    coeff.set_lmult_mult(caffe_rng_generate<float>(aug.lmult_mult(), discount_coeff));
  if (aug.has_lmult_add())
    coeff.set_lmult_add(caffe_rng_generate<float>(aug.lmult_add(), discount_coeff));
  if (aug.has_col_rotate())
    coeff.set_col_angle(caffe_rng_generate<float>(aug.col_rotate(), discount_coeff));
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::clear_defaults(AugmentationCoeff& coeff) {
  const google::protobuf::Reflection* ref = coeff.GetReflection();
  const google::protobuf::Descriptor* desc = coeff.GetDescriptor();
  for (int fn = 0 ; fn < desc->field_count(); fn++) {
    const google::protobuf::FieldDescriptor* field_desc = desc->field(fn);
    float field_val = ref->GetFloat(coeff, field_desc);
    if (fabs(field_desc->default_value_float() - field_val) < 1e-3)
      ref->ClearField(&coeff, field_desc);
//     LOG(INFO) << "field " << fn << " cleared";
  }
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::coeff_to_array(const AugmentationCoeff& coeff, Dtype* out) {
  const google::protobuf::Reflection* ref = coeff.GetReflection();
  const google::protobuf::Descriptor* desc = coeff.GetDescriptor();
  for (int fn = 0 ; fn < desc->field_count(); fn++) {
    const google::protobuf::FieldDescriptor* field_desc = desc->field(fn);
    float field_val = ref->GetFloat(coeff, field_desc);
    if (fabs(field_desc->default_value_float()) < 1e-3)
      out[fn] = field_val;
    else
      out[fn] = log(field_val);
//     LOG(INFO) << "writing field " << fn << ": " << out[fn];
  }
//   LOG(INFO) << "Finished writing params";
}

template <typename Dtype>
void DataAugmentationLayer<Dtype>::array_to_coeff(Dtype* in, AugmentationCoeff& coeff) {
  const google::protobuf::Reflection* ref = coeff.GetReflection();
  const google::protobuf::Descriptor* desc = coeff.GetDescriptor();
  for (int fn = 0 ; fn < desc->field_count(); fn++) {
    const google::protobuf::FieldDescriptor* field_desc = desc->field(fn);
    if (fabs(field_desc->default_value_float()) < 1e-3)
      ref->SetFloat(&coeff, field_desc, in[fn]);
    else
      ref->SetFloat(&coeff, field_desc, exp(in[fn]));
//     LOG(INFO) << "reading field " << fn << ": " << ref->GetFloat(coeff, field_desc);
  }
}

#ifdef CPU_ONLY
STUB_GPU(DataAugmentationLayer);
#endif

INSTANTIATE_CLASS(DataAugmentationLayer);
REGISTER_LAYER_CLASS(DataAugmentation);

}  // namespace caffe


#ifndef STAN_MATH_PRIM_SCAL_PROB_GUMBEL_RNG_HPP
#define STAN_MATH_PRIM_SCAL_PROB_GUMBEL_RNG_HPP

#include <stan/math/prim/scal/err/check_consistent_sizes.hpp>
#include <stan/math/prim/scal/err/check_finite.hpp>
#include <stan/math/prim/scal/err/check_not_nan.hpp>
#include <stan/math/prim/scal/err/check_positive.hpp>
#include <stan/math/prim/scal/meta/length.hpp>
#include <stan/math/prim/scal/meta/VectorBuilder.hpp>
#include <stan/math/prim/scal/meta/return_type.hpp>
#include <stan/math/prim/scal/fun/constants.hpp>
#include <stan/math/prim/scal/meta/include_summand.hpp>
#include <stan/math/prim/scal/fun/value_of.hpp>
#include <boost/random/uniform_01.hpp>
#include <boost/random/variate_generator.hpp>

namespace stan {
  namespace math {

    /**
     * Return a pseudorandom Gumbel variate with the given location and scale
     * using the specified random number generator.
     *
     * @tparam RNG type of random number generator
     * @param mu location parameter
     * @param beta positive scale parameter
     * @param rng random number generator
     * @return Gumbel random variate
     * @throw std::domain_error if mu is infinite or beta is nonpositive.
     */
    template <class RNG>
    inline double
    gumbel_rng(double mu,
               double beta,
               RNG& rng) {
      using boost::variate_generator;
      using boost::uniform_01;

      static const char* function("gumbel_rng");

      check_finite(function, "Location parameter", mu);
      check_positive(function, "Scale parameter", beta);

      variate_generator<RNG&, uniform_01<> >
        uniform01_rng(rng, uniform_01<>());
      return mu - beta * std::log(-std::log(uniform01_rng()));
    }

  }
}
#endif


#include "math.hpp"

#include <boost/integer.hpp>

#include <cstring>
#include <limits>

namespace base
{

template <typename Float>
bool AlmostEqualULPs(Float x, Float y, uint32_t maxULPs)
{
  static_assert(std::is_floating_point<Float>::value, "");
  static_assert(std::numeric_limits<Float>::is_iec559, "");

  // Make sure maxUlps is non-negative and small enough that the
  // default NaN won't compare as equal to anything.
  ASSERT_LESS(maxULPs, 4 * 1024 * 1024, ());

  int constexpr bits = CHAR_BIT * sizeof(Float);
  typedef typename boost::int_t<bits>::exact IntType;
  typedef typename boost::uint_t<bits>::exact UIntType;

  // Same as *reinterpret_cast<IntType const *>(&x), but without warnings.
  IntType xInt, yInt;
  static_assert(sizeof(xInt) == sizeof(x), "bit_cast impossible");
  std::memcpy(&xInt, &x, sizeof(x));
  std::memcpy(&yInt, &y, sizeof(y));

  // Make xInt and yInt lexicographically ordered as a twos-complement int.
  IntType const highestBit = IntType(1) << (bits - 1);
  if (xInt < 0)
    xInt = highestBit - xInt;
  if (yInt < 0)
    yInt = highestBit - yInt;

  // Calculate diff with special case to avoid IntType overflow.
  UIntType diff;
  if ((xInt >= 0) == (yInt >= 0))
    diff = Abs(xInt - yInt);
  else
    diff = UIntType(Abs(xInt)) + UIntType(Abs(yInt));

  return diff <= maxULPs;
}

template bool AlmostEqualULPs<float>(float x, float y, uint32_t maxULPs);
template bool AlmostEqualULPs<double>(double x, double y, uint32_t maxULPs);

} // namespace base


#pragma once

#include <cstdint>
#include <Eigen/Dense>

class RobotController {
private:
    // X Y W bot velocities (m/s)
    const static int numStates = 3;

    // Wheel 1-4 duty cycle % max
    const static int numWheels = 4;

public:
    RobotController(uint32_t dt_us);

    /**
     * Updates the controller with the latest input and calculates
     * the correct wheel velocities to reach this target
     *
     * @param pv Current state (XYW vel in m/s or rad/s)
     * @param sp Current target (XYW vel in m/s or rad/s)
     * @param output Motor targets (rad/s)
     */
    void calculateBody(Eigen::Matrix<double, numStates, 1> pv,
                       Eigen::Matrix<double, numStates, 1> sp,
                       Eigen::Matrix<double, numWheels, 1>& outputs);

    /**
     * Updates the wheels such that they try to follow the target
     * Outputs the correct motor commands to do this
     *
     * @param pv Current state (W1-4 in rad/s)
     * @param sp Current target (W1-4 in rad/s)
     * @param output Motor duty cycles in % max (-1 to 1)
     */
    void calculateWheel(Eigen::Matrix<double, numWheels, 1> pv,
                         Eigen::Matrix<double, numWheels, 1> sp,
                         Eigen::Matrix<double, numWheels, 1>& outputs);

private:
    /**
     * Limits the difference between the previous target and the new target
     * such that the acceleration limits below are never broken
     */
    bool limitBodyAccel(const Eigen::Matrix<double, numStates, 1> finalTarget,
                        Eigen::Matrix<double, numStates, 1>& dampened);

    bool limitWheelAccel(const Eigen::Matrix<double, numWheels, 1> finalTarget,
                        Eigen::Matrix<double, numWheels, 1>& dampened);

    // Body Vel
    Eigen::Matrix<double, numStates, 1> BodyKp;
    Eigen::Matrix<double, numStates, 1> BodyKi;

    Eigen::Matrix<double, numStates, 1> BodyErrorSum;

    // Restrict the i term from contributing more than
    // this amount to the output
    Eigen::Matrix<double, numStates, 1> BodyILimit;
    bool BodyUseILimit;
    bool BodyInputLimited;
    bool BodyOutputLimited;

    Eigen::Matrix<double, numStates, 1> BodyPrevTarget;

    // Wheel Vel
    Eigen::Matrix<double, numWheels, 1> WheelKp;

    Eigen::Matrix<double, numWheels, 1> WheelPrevTarget;

    double dt;

    // Max wheel acceleration rad per second^2)
    static constexpr double maxWheelAccel = 160;

    // Max acceleration (meters per second^2)
    static constexpr double maxForwardAccel = 8;
    // Max acceleration (meters per second^2)
    static constexpr double maxSideAccel = 8;

    // Max angular acceleration (rad per second^2)
    static constexpr double maxAngularAccel = 40.0;
};


#define BOOST_TEST_MODULE "test_PWMcos_potential"

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

#include <mjolnir/core/System.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/forcefield/PWMcos/PWMcosPotential.hpp>

BOOST_AUTO_TEST_CASE(PWMcos_Potential_f)
{
    mjolnir::LoggerManager::set_default_logger("test_pdns_potential.log");

    using real_type = double;
    using traits_type  = mjolnir::SimulatorTraits<real_type, mjolnir::UnlimitedBoundary>;
    using potential_type = mjolnir::PWMcosPotential<traits_type>;
    using ignore_group_type    = typename potential_type::ignore_group_type;
    using ignore_molecule_type = typename potential_type::ignore_molecule_type;

    constexpr std::size_t N = 1000;
    constexpr real_type   h = 1e-6;
    constexpr real_type  pi = mjolnir::math::constants<real_type>::pi();

    potential_type potential(
        1.0, pi / 18.0, 1.0, 0.0, 5.0, {/* contact = empty */}, {/*dnas = empty*/},
        {/* exclude = empty */}, ignore_molecule_type("Nothing"), ignore_group_type({}));

    const real_type r_0   = 7.0;
    const real_type r_min = 1.0;
    const real_type r_max = 15.0;
    const real_type dr    = (r_max - r_min) / N;

    for(std::size_t i=0; i<N; ++i)
    {
        const real_type r   = r_min + i * dr;
        const real_type f1  = potential.f(r_0, r + h);
        const real_type f2  = potential.f(r_0, r - h);
        const real_type df_numeric  = (f1 - f2) / (2 * h);
        const real_type df_analytic = potential.f_df(r_0, r).second;

        if((f1 == 0.0 || f2 == 0.0) && !(f1 == 0.0 && f2 == 0.0))
        {
            // the numeric differentiation becomes unstable here.
        }
        else
        {
            BOOST_TEST(df_numeric == df_analytic, boost::test_tools::tolerance(h));
        }

        const real_type f1_ = potential.f_df(r_0, r+h).first;
        const real_type f2_ = potential.f_df(r_0, r-h).first;
        BOOST_TEST(f1 == f1_, boost::test_tools::tolerance(h));
        BOOST_TEST(f2 == f2_, boost::test_tools::tolerance(h));
    }
}

BOOST_AUTO_TEST_CASE(PWMcos_Potential_g)
{
    mjolnir::LoggerManager::set_default_logger("test_pdns_potential.log");

    using real_type = double;
    using traits_type  = mjolnir::SimulatorTraits<real_type, mjolnir::UnlimitedBoundary>;
    using potential_type = mjolnir::PWMcosPotential<traits_type>;
    using ignore_group_type    = typename potential_type::ignore_group_type;
    using ignore_molecule_type = typename potential_type::ignore_molecule_type;

    constexpr std::size_t N = 1000;
    constexpr real_type   h = 1e-6;
    constexpr real_type  pi = mjolnir::math::constants<real_type>::pi();

    potential_type potential(
        1.0, pi / 18.0, 1.0, 0.0, 5.0, {/* contact = empty */}, {/*dnas = empty*/},
        {/* exclude = empty */}, ignore_molecule_type("Nothing"), ignore_group_type({}));

    const real_type theta_0   = pi * 0.5;
    const real_type theta_min = pi * 0.3;
    const real_type theta_max = pi * 0.7;
    const real_type dtheta    = (theta_max - theta_min) / N;

    for(std::size_t i=0; i<N; ++i)
    {
        const real_type theta = theta_min + i * dtheta;
        const real_type g1  = potential.g(theta_0, theta + h);
        const real_type g2  = potential.g(theta_0, theta - h);
        const real_type dg_numeric  = (g1 - g2) / (2 * h);
        const real_type dg_analytic = potential.g_dg(theta_0, theta).second;

        if((g1 == 0.0 || g2 == 0.0) && !(g1 == 0.0 && g2 == 0.0))
        {
            // the numeric differentiation becomes unstable here.
        }
        else
        {
            BOOST_TEST(dg_numeric == dg_analytic, boost::test_tools::tolerance(h));
        }
        const real_type g1_ = potential.g_dg(theta_0, theta+h).first;
        const real_type g2_ = potential.g_dg(theta_0, theta-h).first;
        BOOST_TEST(g1 == g1_, boost::test_tools::tolerance(h));
        BOOST_TEST(g2 == g2_, boost::test_tools::tolerance(h));
    }
}