#include <stdio.h>
#include <stdlib.h>
#include <Eigen/Dense>
#include "eigenmat.h"
#include "ziggurat.h"

extern "C" {

extern int init_random(rnd_struct* rnd_state, int seed) {
  rnd_state->seed = seed;
  r4_nor_setup(rnd_state->kn, rnd_state->fn, rnd_state->wn);
}

/* ------------------------------ Utility routines ------------------------------ */

extern int get_leading_dimension(eigenmat* mat) {
  return mat->is_trans ? mat->size[1] : mat->size[0];
}

extern int get_nonleading_dimension(eigenmat* mat) {
  return mat->is_trans ? mat->size[0] : mat->size[1];
}

extern void set_transpose(eigenmat* mat, int is_trans) {
  mat->is_trans = is_trans;
}

inline char get_transpose_char(eigenmat* mat) {
  return mat->is_trans ? 't' : 'n';
}

/* ------------------------------ Allocating/moving data ------------------------------ */

extern int allocate_memory(eigenmat* mat) {
  const int len = mat->size[0] * mat->size[1];
  mat->data = (float*)malloc(len * sizeof(float));
  return 0;
}

extern int copy_on_device(eigenmat* mat1, eigenmat* mat2) {
  const int len = mat1->size[0]*mat1->size[1];

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::VectorXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::VectorXf> eig_mat2(mat2->data, len);
  eig_mat2 = eig_mat1;
  return 0;
}

extern int get_row_slice(eigenmat* source, eigenmat* target, unsigned int start, unsigned int end) {
  int height = source->size[0];
  int width = source->size[1];

  if ((end - start) != target->size[0] || source->size[1] != target->size[1] || start >= end || end > height)
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  return 0;
}

extern int set_row_slice(eigenmat* source, eigenmat* target, unsigned int start, unsigned int end) {
  int height = target->size[0];
  int width = target->size[1];

  if ((end - start) != source->size[0] || source->size[1] != target->size[1] || start >= end || end > height)
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  return 0;
}

extern int copy_transpose(eigenmat* source, eigenmat* target) {
  unsigned int height = source->size[0];
  unsigned int width = source->size[1];

  if (source->size[0] != target->size[1] || source->size[1] != target->size[0])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::MatrixXf> eig_source(source->data, source->size[0], source->size[1]);
  Eigen::Map<Eigen::MatrixXf> eig_target(target->data, target->size[0], target->size[1]);
  eig_target = eig_source.transpose();
  return 0;
}

extern int set_shape(eigenmat* mat, unsigned int m, unsigned int n) {
  mat->size[0] = m;
  mat->size[1] = n;
  return 0;
}

extern int reshape(eigenmat* mat, unsigned int m, unsigned int n) {
  if (mat->size[0] * mat->size[1] != m * n)
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  mat->size[0] = m;
  mat->size[1] = n;
  return 0;
}

extern int get_slice(eigenmat* source, eigenmat* target, unsigned int first_col, unsigned int last_col) {
  if (source->is_trans)
    return ERROR_TRANSPOSED;

  if (last_col > source->size[1] || (first_col >= last_col))
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  int num_rows = source->size[0];

  target->data = source->data + first_col * num_rows;
  target->size[0] = source->size[0];
  target->size[1] = last_col - first_col;
  target->is_trans = 0;
  target->owns_data = 0;

  return 0;
}

extern int get_vector_slice(eigenmat* source, eigenmat* target, unsigned int first_ind, unsigned int last_ind) {
  // source must be a vector
  if (source->size[0] > 1 && source->size[1] > 1)
    return ERROR_GENERIC;

  if (source->is_trans)
    return ERROR_TRANSPOSED;

  if (first_ind >= last_ind)
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  int num_rows = source->size[0];

  target->data = source->data + first_ind * num_rows;
  target->is_trans = 0;
  target->owns_data = 0;

  if (source->size[0] > 1) {
    if (last_ind > source->size[0])
      return ERROR_INCOMPATIBLE_DIMENSIONS;

    target->size[0] = last_ind - first_ind;
    target->size[1] = 1;
  } else {
    if (last_ind > source->size[1])
      return ERROR_INCOMPATIBLE_DIMENSIONS;

    target->size[0] = 1;
    target->size[1] = last_ind - first_ind;
  }

  return 0;
}

/* ------------------------------ Initialization routines ------------------------------ */

extern void init_from_array(eigenmat* mat, float* data, int m, int n) {
  mat->data = data;
  mat->size[0] = m;
  mat->size[1] = n;
  mat->is_trans = 0;
  mat->owns_data = 1;
}

extern int init_empty(eigenmat* mat, int m, int n) {
  mat->size[0] = m;
  mat->size[1] = n;
  mat->is_trans = 0;
  mat->owns_data = 1;

  return allocate_memory(mat);
}

/* ------------------------------ Random number generation ------------------------------ */

extern float uniform(rnd_struct* rnd_state) {
  return r4_uni(&rnd_state->seed);
}

extern float normal(rnd_struct* rnd_state) {
  return r4_nor(&rnd_state->seed, rnd_state->kn, rnd_state->fn, rnd_state->wn);
}

extern int fill_with_rand(rnd_struct* rnd_state, eigenmat* mat) {
  const int len = mat->size[0] * mat->size[1];
  for (int i = 0; i < len; i++) mat->data[i] = uniform(rnd_state);
  return 0;
}

extern int fill_with_randn(rnd_struct* rnd_state, eigenmat* mat) {
  const int len = mat->size[0] * mat->size[1];
  for (int i = 0; i < len; i++) mat->data[i] = normal(rnd_state);
  return 0;
}

extern int sample_bernoulli(rnd_struct* rnd_state, eigenmat* mat, eigenmat* target) {
  int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] > uniform(rnd_state) ? 1 : 0;

  return 0;
}
extern int sample_bernoulli_tanh(rnd_struct* rnd_state, eigenmat* mat, eigenmat* target) {
  int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] > uniform(rnd_state) ? 1 : -1;

  return 0;
}

extern int sample_gaussian(rnd_struct* rnd_state, eigenmat* mat, eigenmat* target, float mult) {
  int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] + mult * normal(rnd_state);

  return 0;
}

extern int perturb_energy(rnd_struct* rnd_state, eigenmat* mat, eigenmat* target) {
  const int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] - log(-log(uniform(rnd_state)));

  return 0;
}

extern int perturb_prob(rnd_struct* rnd_state, eigenmat* mat, eigenmat* target) {
  const int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = -mat->data[i] / log(uniform(rnd_state));

  return 0;
}

extern int dropout(rnd_struct* rnd_state, eigenmat* mat, float dropprob, float val) {
  const int len = mat->size[0] * mat->size[1];

  for (int i = 0; i < len; i++)
    mat->data[i] = dropprob > uniform(rnd_state) ? 0 : mat->data[i];

  return 0;
}

/* ------------------------------ Algebraic operations ------------------------------ */

extern int add_col_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[0] != vec->size[0] || vec->size[1] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, h);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.colwise() + eig_vec;

  return 0;
}

extern int add_mult_sign(eigenmat* mat, eigenmat* mat2, float mult) {
  for (int i = 0; i < mat->size[0] * mat->size[1]; i++) {
    mat->data[i] += (mat2->data[i] == 0) ? 0 : ((mat2->data[i] > 0) ? mult:-mult);
  }
  return 0;
}

extern int add_col_mult(eigenmat* mat, eigenmat* vec, eigenmat* target, float mult) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[0] != vec->size[0] || vec->size[1] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, h);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.colwise() + eig_vec * mult;

  return 0;
}

extern int mult_diagonal_scalar(eigenmat* mat, float val, eigenmat* target) {
  unsigned int w = mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < w; i++) target->data[i*w+i] = val * mat->data[i*w+i];

  return 0;
}

extern int add_diagonal_scalar(eigenmat* mat, float val, eigenmat* target) {
  unsigned int w = mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < w; i++) target->data[i*w+i] = val * mat->data[i*w+i];

  return 0;
}

extern int mult_diagonal(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int w = mat->size[1];

  if (mat->size[0] != vec->size[1] * vec->size[0] ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < w; i++) target->data[i*w+i] = vec->data[i] * mat->data[i*w+i];

  return 0;
}

extern int add_diagonal(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int w = mat->size[1];

  if (mat->size[0] != vec->size[1] * vec->size[0] ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < w; i++) target->data[i*w+i] = vec->data[i] + mat->data[i*w+i];
  return 0;
}

extern int add_row_mult(eigenmat* mat, eigenmat* vec, eigenmat* target, float mult) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[1] != vec->size[1] || vec->size[0] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, w);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.rowwise() + eig_vec.transpose() * mult;
  return 0;
}

extern int add_row_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[1] != vec->size[1] || vec->size[0] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, w);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.rowwise() + eig_vec.transpose();
  return 0;
}

extern int mult_by_col_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[0] != vec->size[0] || vec->size[1] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, h);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.colwise() * eig_vec;

  return 0;
}

extern int mult_by_row_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[1] != vec->size[1] || vec->size[0] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, w);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.rowwise() * eig_vec.transpose();
  return 0;
}

extern int div_by_col_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[0] != vec->size[0] || vec->size[1] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, h);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.colwise() / eig_vec;
  return 0;
}

extern int div_by_row_vec(eigenmat* mat, eigenmat* vec, eigenmat* target) {
  unsigned int h = mat->size[0],
         w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (mat->size[1] != vec->size[1] || vec->size[0] != 1 ||
    mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, h, w);
  Eigen::Map<Eigen::ArrayXf> eig_vec(vec->data, w);
  Eigen::Map<Eigen::ArrayXXf> eig_target(target->data, h, w);

  eig_target = eig_mat.rowwise() / eig_vec.transpose();

  return 0;
}
extern int less_than(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  const int len = mat1->size[0] * mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat1->data[i] < mat2->data[i] ? 1 : 0;
  return 0;
}

extern int less_than_scalar(eigenmat* mat, float val, eigenmat* target) {
  const int len = mat->size[0]*mat->size[1];

  if (mat->is_trans != target->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] < val ? 1 : 0;

  return 0;
}

extern int greater_than(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat1->data[i] > mat2->data[i] ? 1 : 0;

  return 0;
}

extern int upper_bound(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat1->data[i] < mat2->data[i] ? mat1->data[i] : mat2->data[i];

  return 0;
}

extern int lower_bound(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat1->data[i] > mat2->data[i] ? mat1->data[i] : mat2->data[i];
}

extern int greater_than_scalar(eigenmat* mat, float val, eigenmat* target) {
  int len = mat->size[0]*mat->size[1];

  if (mat->is_trans != target->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] > val ? 1:0;
  return 0;
}

extern int upper_bound_scalar(eigenmat* mat, float val, eigenmat* target) {
  int len = mat->size[0]*mat->size[1];

  if (mat->is_trans != target->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] < val ? mat->data[i] : val;

  return 0;
}

extern int lower_bound_scalar(eigenmat* mat, float val, eigenmat* target) {
  const int len = mat->size[0]*mat->size[1];

  if (mat->is_trans != target->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < len; i++)
    target->data[i] = mat->data[i] > val ? mat->data[i] : val;

  return 0;
}

extern int cumsum_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  const unsigned int h = mat->size[0], w = mat->size[1];
  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      const float *mat_data = &mat->data[i * h];
      float *target_data = &target->data[i * h];
      target_data[0] = mat_data[0];
      for (int j = 1; j < h; j++)
        target_data[j] =  target_data[j-1] + mat_data[j];
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      const float *mat_data = &mat->data[i];
      float *target_data = &target->data[i];
      target_data[0] = mat_data[0];
      for (int j = 1; j < w; j++)
        target_data[j*h] =  target_data[(j-1)*h] + mat_data[j*h];
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}

extern int max_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  const unsigned int h = mat->size[0], w = mat->size[1];
  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      const float *mat_data = &mat->data[i * h];
      float max = mat_data[0];
      for (int j = 1; j < h; j++)
        if (max < mat_data[j]) max = mat_data[j];
      target->data[i] = max;
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      const float *mat_data = &mat->data[i];
      float max = mat_data[0];
      for (int j = 1; j < w; j++)
        if (max < mat_data[j*h]) max = mat_data[j*h];
      target->data[i] = max;
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}

extern int choose_max_and_accumulate(eigenmat* mat, eigenmat* acc) {
  const unsigned int h = mat->size[0], w = mat->size[1];

  if (mat->is_trans)
    return ERROR_TRANSPOSED;

  if (acc->size[0] != mat->size[0] || acc->size[1] != mat->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for (int i = 0; i < w; i++) {
    int argmax = 0;
    const float *mat_data = &mat->data[i * h];
    for (int j = 1; j < h; j++)
      if (mat_data[argmax] < mat_data[j]) argmax = j;
    acc->data[i * h + argmax] += 1;
  }

  return 0;
}

extern int choose_max_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  const unsigned int h = mat->size[0], w = mat->size[1];
  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      int argmax = 0;
      const float *mat_data = &mat->data[i * h];
      float *target_data = &target->data[i * h];
      for (int j = 1; j < h; j++)
        if (mat_data[argmax] < mat_data[j]) argmax = j;
      for (int j = 0; j < h; j++) target_data[j] = (j == argmax) ? 1 : 0;
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      int argmax = 0;
      const float *mat_data = &mat->data[i];
      float *target_data = &target->data[i];
      for (int j = 1; j < w; j++)
        if (mat_data[argmax*h] < mat_data[j*h]) argmax = j;
      for (int j = 0; j < w; j++) target_data[j*h] = (j == argmax) ? 1 : 0;
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}
extern int argmax_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  const unsigned int h = mat->size[0], w = mat->size[1];
  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      int argmax = 0;
      const float *mat_data = &mat->data[i * h];
      for (int j = 1; j < h; j++)
        if (mat_data[argmax] < mat_data[j]) argmax = j;
      target->data[i] = argmax;
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      int argmax = 0;
      const float *mat_data = &mat->data[i];
      for (int j = 1; j < w; j++)
        if (mat_data[argmax*h] < mat_data[j*h]) argmax = j;
      target->data[i] = argmax;
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}

extern int sqsum_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  const unsigned int h = mat->size[0], w = mat->size[1];

  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      float sum = 0;
      const float *mat_data = &mat->data[i * h];
      for (int j = 0; j < h; j++) sum += mat_data[j] * mat_data[j];
      target->data[i] = sum;
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      float sum = 0;
      const float *mat_data = &mat->data[i];
      for (int j = 0; j < w; j++) sum += mat_data[j*h] * mat_data[j*h];
      target->data[i] = sum;
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}

extern int add_sum_by_axis(eigenmat* mat, eigenmat* target, int axis, const float mult) {
  int len_mat = mat->size[0] * mat->size[1];
  int len_target = target->size[0] * target->size[1];
  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, mat->size[0], mat->size[1]);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len_target);
  if (mat->is_trans) axis = 1 - axis;
  if (axis == 0) {
    eig_target += mult * eig_mat.colwise().sum();
  } else {
    eig_target += mult * eig_mat.rowwise().sum();
  }
  return 0;
}

extern int sum_by_axis(eigenmat* mat, eigenmat* target, int axis) {
  int len_mat = mat->size[0] * mat->size[1];
  int len_target = target->size[0] * target->size[1];
  Eigen::Map<Eigen::ArrayXXf> eig_mat(mat->data, mat->size[0], mat->size[1]);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len_target);
  if (mat->is_trans) axis = 1 - axis;
  if (axis == 0) {
    eig_target = eig_mat.colwise().sum();
  } else {
    eig_target = eig_mat.rowwise().sum();
  }
  return 0;
}

extern int normlimit_by_axis(eigenmat* mat, eigenmat* target, int axis,
    float norm){
  const unsigned int h = mat->size[0], w = mat->size[1];

  if (axis == 0) {
    for (int i = 0; i < w; i++) {
      float sum = 0;
      const float *mat_data = &mat->data[i * h];
      float *target_data = &target->data[i * h];

      for (int j = 0; j < h; j++) sum += mat_data[j] * mat_data[j];
      float colnorm = sqrt(sum);

        if (colnorm > norm)
        {
            float factor = norm/colnorm;
            for (int j = 0; j < h; j++)  target_data[j] = factor * mat_data[j];

        }
        else
        {
            for (int j = 0; j < h; j++)  target_data[j] =  mat_data[j];
        }
    }
  } else if (axis == 1) {
    for (int i = 0; i < h; i++) {
      float sum = 0;
      const float *mat_data = &mat->data[i];
      float *target_data = &target->data[i];
      for (int j = 0; j < w; j++) sum += mat_data[j*h] * mat_data[j*h];
      float rownorm = sqrt(sum);

        if (rownorm > norm)
        {
            float factor = norm/rownorm;
            for (int j = 0; j < w; j++)  target_data[j*h] = factor * mat_data[j*h];
        }
        else
        {
            for (int j = 0; j < w; j++)  target_data[j*h] =  mat_data[j*h];
        }
    }
  } else {
    return ERROR_UNSUPPORTED;
  }
  return 0;
}

extern int sign(eigenmat* mat, eigenmat* target) {
  int len = mat->size[0]*mat->size[1];

  if (mat->is_trans != target->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < len; i++)
    target->data[i] = (mat->data[i] < 0) ? -1 : (mat->data[i] > 0 ? 1 : 0);

  return 0;
}
extern int apply_cos(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.cos();
  return 0;
}
extern int apply_sin(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.sin();

  return 0;
}

extern int apply_softmax(eigenmat* mat, eigenmat* target) {
  int width = mat->size[1], height = mat->size[0];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < width; i++) {
    float *mat_data = &mat->data[i * height];
    float *target_data = &target->data[i * height];
    float temp = mat->data[i * height];
    for (int j = 1; j < height; j++) if (temp < mat_data[j]) temp = mat_data[j];
    for (int j = 0; j < height; j++) target_data[j] = exp(mat_data[j] - temp);
    temp = 0;
    for (int j = 0; j < height; j++) temp += target_data[j];
    for (int j = 0; j < height; j++) target_data[j] /= temp;
  }
  return 0;
}

extern int apply_softmax_grad(eigenmat* mat, eigenmat* labels, eigenmat* target) {
  const int width = mat->size[1], height = mat->size[0];

  if (target != mat) {
    #pragma omp parallel for
    for (int i = 0; i < width * height; i++) target->data[i] = mat->data[i];
  }
  for (int i = 0; i < width; i++) {
    target->data[i * height + (int)labels->data[i]] -= 1.0;
  }
  return 0;

}
extern int get_softmax_cross_entropy(eigenmat* mat, eigenmat* labels, eigenmat* target, const float tiny) {
  const int width = mat->size[1], height = mat->size[0];

  for (int i = 0; i < width; i++) {
    target->data[i] = -log(mat->data[i * height + (int)labels->data[i]]);
  }
  return 0;
}

extern int get_softmax_correct(eigenmat* mat, eigenmat* labels, eigenmat* target) {
  int width = mat->size[1], height = mat->size[0];

  #pragma omp parallel for
  for (int i = 0; i < width; i++) {
    int argmax = 0;
    const int correct_label = (int)labels->data[i];
    float *mat_data = &mat->data[i * height];
    for (int j = 1; j < height; j++) if (mat_data[argmax] < mat_data[j]) argmax = j;
    target->data[i] = (correct_label == argmax) ? 1: 0;
  }
  return 0;
}

extern float sum_all(eigenmat* mat) {
  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, mat->size[0] * mat->size[1]);
  return eig_mat.sum();
}

extern int apply_sigmoid(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for(int i = 0; i < len; i++) target->data[i] = 1 / (1 + exp(-mat->data[i]));
  return 0;
}

extern int apply_tanh(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  #pragma omp parallel for
  for(int i = 0; i < len; i++) target->data[i] = 2 / (1 + exp(-mat->data[i])) - 1;
  return 0;
}

extern int apply_abs(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.abs();
  return 0;
}

extern int apply_log_1_plus_exp(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  for(int i = 0; i < len; i++) target->data[i] = log(1+exp(mat->data[i]));

  return 0;
}

extern int apply_log(eigenmat* mat, eigenmat* target, float tiny) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.log();

  return 0;
}

extern int apply_exp(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.exp();

  return 0;
}
extern int apply_ceil(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for(int i = 0; i < len; i++) target->data[i] = ceil(mat->data[i]);

  return 0;
}
extern int apply_floor(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for(int i = 0; i < len; i++) target->data[i] = floor(mat->data[i]);

  return 0;
}

extern int apply_sqrt(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.sqrt();

  return 0;
}

extern int apply_pow(eigenmat* mat, float exponent, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat.pow(exponent);

  return 0;
}

extern int apply_pow_matrix(eigenmat* mat, eigenmat* exponent, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  if (mat->size[0] != exponent->size[0] || mat->size[1] != exponent->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for(int i = 0; i < len; i++) target->data[i] = pow(mat->data[i], exponent->data[i]);
  return 0;
}

extern int compute_cross_entropy(eigenmat* mat1, eigenmat* mat2, eigenmat* target, float tiny) {
  unsigned int len = mat1->size[0] * mat1->size[1];

  if (mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  if (mat2->size[0] != mat2->size[0] || mat2->size[1] != mat2->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++) target->data[i] = - mat1->data[i] * log(mat2->data[i] + tiny);

  return 0;
}
extern int compute_cross_entropy_bernoulli(eigenmat* mat1, eigenmat* mat2, eigenmat* target, float tiny) {
  unsigned int len = mat1->size[0] * mat1->size[1];

  if (mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = - mat1->data[i] * log(mat2->data[i] + tiny) - (1-mat1->data[i]) * log(1 - mat2->data[i] + tiny);
  return 0;
}
extern int correct_preds(eigenmat* mat1, eigenmat* mat2, eigenmat* target, float cutoff) {
  unsigned int len = mat1->size[0] * mat1->size[1];

  if (mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++)
    target->data[i] = mat1->data[i] * mat2->data[i] >= cutoff + (1 - mat1->data[i]) * (mat2->data[i] < cutoff);
  return 0;
}

extern int reciprocal(eigenmat* mat, eigenmat* target) {
  unsigned int len = mat->size[0] * mat->size[1];

  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);

  eig_target = eig_mat.inverse();

  return 0;
}

extern int dot(eigenmat* mat1, eigenmat* mat2, eigenmat* target, float beta, float alpha) {

  if (get_leading_dimension(mat1) != get_leading_dimension(target) ||
    get_nonleading_dimension(mat2) != get_nonleading_dimension(target) ||
    get_nonleading_dimension(mat1) != get_leading_dimension(mat2)) {
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  }
  int m = get_leading_dimension(mat1),
    k = get_leading_dimension(mat2),
    n = get_nonleading_dimension(mat2);

  Eigen::Map<Eigen::MatrixXf> eig_mat1(mat1->data, mat1->size[0], mat1->size[1]);
  Eigen::Map<Eigen::MatrixXf> eig_mat2(mat2->data, mat2->size[0], mat2->size[1]);
  Eigen::Map<Eigen::MatrixXf> eig_target(target->data, target->size[0], target->size[1]);

  eig_target = beta * eig_target;
  if (mat1->is_trans && mat2->is_trans) {
    eig_target.noalias() += alpha * (eig_mat1.transpose() * eig_mat2.transpose());
  } else if (mat1->is_trans) {
    eig_target.noalias() += alpha * (eig_mat1.transpose() * eig_mat2);
  } else if (mat2->is_trans) {
    eig_target.noalias() += alpha * (eig_mat1 * eig_mat2.transpose());
  } else {
    eig_target.noalias() += alpha * (eig_mat1 * eig_mat2);
  }

  return 0;
}

extern float vdot(eigenmat* mat1, eigenmat* mat2, int* err_code) {
  int len = mat1->size[0]*mat1->size[1];
  float res;

  if (mat1->is_trans != mat2->is_trans) {
    *err_code = ERROR_TRANSPOSEDNESS;
    return 0;
  }

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1]) {
    *err_code = ERROR_INCOMPATIBLE_DIMENSIONS;
    return 0;
  }
  Eigen::Map<Eigen::VectorXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::VectorXf> eig_mat2(mat2->data, len);
  return eig_mat1.dot(eig_mat2);

}

/* Perform the operation mat1 = mat1 + alpha * mat2. mat1 and mat2 must
  have the same transposedness. */
extern int add_mult(eigenmat* mat1, eigenmat* mat2, float alpha) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_mat2(mat2->data, len);
  eig_mat1 += eig_mat2 * alpha;
  return 0;
}

extern int add_elementwise(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_mat2(mat2->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat1 + eig_mat2 ;
  return 0;
}

extern int subtract_elementwise(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_mat2(mat2->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat1 - eig_mat2 ;
  return 0;
}

extern int divide_elementwise(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::ArrayXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_mat2(mat2->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat1 / eig_mat2 ;

  return 0;
}

/* Elementwise multiplication of 2 matrices */
extern int mult_elementwise(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  Eigen::Map<Eigen::ArrayXf> eig_mat1(mat1->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_mat2(mat2->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat1 * eig_mat2 ;
  return 0;
}

extern int apply_sin_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  return 0;
}
extern int apply_cos_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  return 0;
}
extern int apply_logistic_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++) target->data[i] = mat1->data[i] * mat2->data[i] * (1-mat2->data[i]);

  return 0;
}

extern int apply_tanh_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++) target->data[i] = 0.5 * mat1->data[i] * (1 + mat2->data[i]) * (1 - mat2->data[i]);
  return 0;
}

extern int apply_rectified_linear_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++) target->data[i] =  mat1->data[i] * (mat2->data[i] > 0 ? 1 : 0);
  return 0;
}

extern int apply_rectified_linear_smooth_deriv(eigenmat* mat1, eigenmat* mat2, eigenmat* target) {
  int len = mat1->size[0]*mat1->size[1];

  if (mat1->is_trans != mat2->is_trans)
    return ERROR_TRANSPOSEDNESS;

  if (mat1->size[0] != mat2->size[0] || mat1->size[1] != mat2->size[1] ||
    mat1->size[0] != target->size[0] || mat1->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;

  #pragma omp parallel for
  for (int i = 0; i < len; i++) target->data[i] =  mat1->data[i] * (1 - exp(-mat2->data[i]));
  return 0;
}

extern int assign_scalar(eigenmat* mat, float alpha) {
  int len = mat->size[0]*mat->size[1];
  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  eig_mat.setConstant(alpha);
  return 0;
}

extern int mult_by_scalar(eigenmat* mat, float alpha, eigenmat* target) {
  const int len = mat->size[0]*mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat * alpha;
  return 0;
}

extern int divide_by_scalar(eigenmat* mat, float alpha, eigenmat* target) {
  const int len = mat->size[0]*mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat / alpha;
  return 0;
}

extern int add_scalar(eigenmat* mat, float alpha, eigenmat* target) {
  const int len = mat->size[0] * mat->size[1];
  if (mat->size[0] != target->size[0] || mat->size[1] != target->size[1])
    return ERROR_INCOMPATIBLE_DIMENSIONS;
  Eigen::Map<Eigen::ArrayXf> eig_mat(mat->data, len);
  Eigen::Map<Eigen::ArrayXf> eig_target(target->data, len);
  eig_target = eig_mat + alpha;
  return 0;
}

extern float euclid_norm(eigenmat* mat) {
  const int len = mat->size[0]*mat->size[1];
  Eigen::Map<Eigen::VectorXf> eig_mat(mat->data, len);
  return eig_mat.norm();
}

extern int selectCols(eigenmat* source, eigenmat* target, eigenmat* indices){
  const int n = indices->size[1] * indices->size[0];
  const int h = source->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int i = 0; i < n; i++) {
    target_offset = i * h;
    source_offset = (int)indices->data[i] * h;
    for (int j = 0; j < h; j++) {
      target->data[target_offset + j] = source->data[source_offset + j];
    }
  }
  return 0;
}

extern int selectRows(eigenmat* source, eigenmat* target, eigenmat* indices){
  const int n = indices->size[1] * indices->size[0];
  const int w = source->size[1], h_source = source->size[0], h_target = target->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int j = 0; j < w; j++) {
    for (int i = 0; i < n; i++) {
      target->data[j * h_target + i] = source->data[j * h_source + (int)indices->data[i]];
    }
  }
  return 0;
}

extern int swapCols(eigenmat* source, eigenmat* target, eigenmat* indices1, eigenmat* indices2){
  const int n = indices1->size[1] * indices1->size[0];
  const int h = source->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int i = 0; i < n; i++) {
    source_offset = (int)indices1->data[i] * h;
    target_offset = (int)indices2->data[i] * h;
    for (int j = 0; j < h; j++) {
      float temp = target->data[target_offset + j];
      target->data[target_offset + j] = source->data[source_offset + j];
      source->data[source_offset + j] = temp;
    }
  }
  return 0;
}

extern int swapRows(eigenmat* source, eigenmat* target, eigenmat* indices1, eigenmat* indices2){
  const int n = indices1->size[1] * indices1->size[0];
  const int w = source->size[1], h_source = source->size[0], h_target = target->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int j = 0; j < w; j++) {
    for (int i = 0; i < n; i++) {
      int source_index = j * h_source + (int)indices1->data[i];
      int target_index = j * h_target + (int)indices2->data[i];
      float temp = target->data[target_index];
      target->data[target_index] = source->data[source_index];
      source->data[source_index] = temp;
    }
  }
  return 0;
}

extern int setSelectedCols(eigenmat* source, eigenmat* target, eigenmat* indices){
  const int n = indices->size[1] * indices->size[0];
  const int h = source->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int i = 0; i < n; i++) {
    source_offset = i * h;
    target_offset = (int)indices->data[i] * h;
    for (int j = 0; j < h; j++) {
      target->data[target_offset + j] = source->data[source_offset + j];
    }
  }
  return 0;
}

extern int setSelectedRows(eigenmat* source, eigenmat* target, eigenmat* indices){
  const int n = indices->size[1] * indices->size[0];
  const int w = source->size[1], h_source = source->size[0], h_target = target->size[0];

  int target_offset, source_offset;
  #pragma omp parallel for
  for (int j = 0; j < w; j++) {
    for (int i = 0; i < n; i++) {
      target->data[j * h_target + (int)indices->data[i]] = source->data[j * h_source + i];
    }
  }
  return 0;
}

extern int generate_translations_big_var_off(eigenmat* source, eigenmat* target, eigenmat* off_x, eigenmat* off_y, int source_w, int target_w, int num_channels) {
  int num_images = source->size[1];
    int source_image_offset =source->size[0];
    int target_image_offset =target->size[0];
    int source_stride = source_w * num_channels;
    int target_stride = target_w * num_channels;

    int row_size = target_stride * sizeof(float);

    int zero_offset = (source_w-target_w)/2;

#pragma omp parallel for
    for (int image = 0; image < num_images; image++)
    {
        int oy = (int) off_y->data[image];
        int ox = (int) off_x->data[image];

        float *sourceP = source->data + image* source_image_offset + (zero_offset+oy)*source_stride + (zero_offset+ox)* num_channels;
        float *targetP = target->data + image* target_image_offset ;

        for (int j = 0; j < target_w; j++) {

            float *source_row = sourceP + j * source_stride ;
            float *target_row = targetP + j * target_stride;

            std::memcpy( target_row, source_row, row_size );

        }

    }

    return 0;
}

extern int blockify(eigenmat* source, eigenmat* target, int blocksize) {
  const int w = source->size[1], h = source->size[0];
  #pragma omp parallel for
  for (int i = 0; i < w; i++) {
    const int off = i * h;
    for (int j = 0; j < h; j++) {
      target->data[off + j] = source->data[off + (j / blocksize) * blocksize];
    }
  }
  return 0;
}
}


// Copyright Louis Dionne 2013-2017
// 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_HANA_TEST_AUTO_UNIQUE_HPP
#define BOOST_HANA_TEST_AUTO_UNIQUE_HPP

#include <boost/hana/assert.hpp>
#include <boost/hana/equal.hpp>
#include <boost/hana/unique.hpp>

#include "test_case.hpp"
#include <laws/base.hpp>
#include <support/equivalence_class.hpp>

TestCase test_unique{[]{
    namespace hana = boost::hana;
    using hana::test::ct_eq;

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE()),
        MAKE_TUPLE()
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{})),
        MAKE_TUPLE(ct_eq<0>{})
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<0>{})),
        MAKE_TUPLE(ct_eq<0>{})
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{})
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<0>{}, ct_eq<0>{})),
        MAKE_TUPLE(ct_eq<0>{})
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<0>{}, ct_eq<1>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{})
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<0>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<0>{})
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<1>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{})
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<2>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<2>{})
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(
            ct_eq<0>{}, ct_eq<0>{},
            ct_eq<1>{},
            ct_eq<2>{}, ct_eq<2>{}, ct_eq<2>{},
            ct_eq<3>{}, ct_eq<3>{}, ct_eq<3>{},
            ct_eq<0>{})),
        MAKE_TUPLE(ct_eq<0>{}, ct_eq<1>{}, ct_eq<2>{}, ct_eq<3>{}, ct_eq<0>{})
    ));
}};

TestCase test_unique_by{[]{
    namespace hana = boost::hana;
    using hana::test::ct_eq;

    struct undefined { };

    auto a = [](auto z) { return ::equivalence_class(ct_eq<999>{}, z); };
    auto b = [](auto z) { return ::equivalence_class(ct_eq<888>{}, z); };
    auto c = [](auto z) { return ::equivalence_class(ct_eq<777>{}, z); };

    auto pred = [](auto x, auto y) {
        return hana::equal(x.unwrap, y.unwrap);
    };

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(), undefined{}),
        MAKE_TUPLE()
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}))
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), b(ct_eq<0>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}))
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}))
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), b(ct_eq<0>{}), c(ct_eq<0>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}))
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), b(ct_eq<0>{}), c(ct_eq<1>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), c(ct_eq<1>{}))
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<0>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<0>{}))
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), b(ct_eq<1>{}), b(ct_eq<1>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), b(ct_eq<1>{}))
    ));
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<2>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<2>{}))
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique(MAKE_TUPLE(
            a(ct_eq<0>{}), b(ct_eq<0>{}),
            a(ct_eq<1>{}),
            a(ct_eq<2>{}), b(ct_eq<2>{}), c(ct_eq<2>{}),
            a(ct_eq<3>{}), b(ct_eq<3>{}), c(ct_eq<3>{}),
            a(ct_eq<0>{})), pred),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<2>{}), a(ct_eq<3>{}), a(ct_eq<0>{}))
    ));

    // unique.by
    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique.by(pred, MAKE_TUPLE(
            a(ct_eq<0>{}), b(ct_eq<0>{}),
            a(ct_eq<1>{}),
            a(ct_eq<2>{}), b(ct_eq<2>{}), c(ct_eq<2>{}),
            a(ct_eq<3>{}), b(ct_eq<3>{}), c(ct_eq<3>{}),
            a(ct_eq<0>{}))),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<2>{}), a(ct_eq<3>{}), a(ct_eq<0>{}))
    ));

    BOOST_HANA_CONSTANT_CHECK(hana::equal(
        hana::unique.by(pred)(MAKE_TUPLE(
            a(ct_eq<0>{}), b(ct_eq<0>{}),
            a(ct_eq<1>{}),
            a(ct_eq<2>{}), b(ct_eq<2>{}), c(ct_eq<2>{}),
            a(ct_eq<3>{}), b(ct_eq<3>{}), c(ct_eq<3>{}),
            a(ct_eq<0>{}))),
        MAKE_TUPLE(a(ct_eq<0>{}), a(ct_eq<1>{}), a(ct_eq<2>{}), a(ct_eq<3>{}), a(ct_eq<0>{}))
    ));
}};

#endif // !BOOST_HANA_TEST_AUTO_UNIQUE_HPP


/**
* @file       Tests.cpp
*
* @brief      Test routines for Zerocoin.
*
* @author     Ian Miers, Christina Garman and Matthew Green
* @date       June 2013
*
* @copyright  Copyright 2013 Ian Miers, Christina Garman and Matthew Green
* @license    This project is released under the MIT license.
**/
// Copyright (c) 2017-2018 The PIVX developers
// Copyright (c) 2018 The time developers

#include <boost/test/unit_test.hpp>
#include <string>
#include <iostream>
#include <fstream>
#include <cmath>
// #include <curses.h>
#include <exception>
#include "streams.h"
#include "libzerocoin/ParamGeneration.h"
#include "libzerocoin/Denominations.h"
#include "libzerocoin/Coin.h"
#include "libzerocoin/CoinSpend.h"
#include "libzerocoin/Accumulator.h"

using namespace std;
using namespace libzerocoin;

#define COLOR_STR_GREEN   "\033[32m"
#define COLOR_STR_NORMAL  "\033[0m"
#define COLOR_STR_RED     "\033[31m"

#define TESTS_COINS_TO_ACCUMULATE   10
#define NON_PRIME_TESTS				100

// Global test counters
uint32_t    gNumTests        = 0;
uint32_t    gSuccessfulTests = 0;

// Proof size
uint32_t    gProofSize			= 0;
uint32_t    gCoinSize			= 0;
uint32_t	gSerialNumberSize	= 0;

// Global coin array
PrivateCoin    *gCoins[TESTS_COINS_TO_ACCUMULATE];

// Global params
ZerocoinParams *g_Params;

//////////
// Utility routines
//////////

void
LogTestResult(string testName, bool (*testPtr)())
{
	string colorGreen(COLOR_STR_GREEN);
	string colorNormal(COLOR_STR_NORMAL);
	string colorRed(COLOR_STR_RED);

	cout << "Testing if " << testName << "..." << endl;

	bool testResult = testPtr();

	if (testResult == true) {
		cout << "\t" << colorGreen << "[PASS]"  << colorNormal << endl;
		gSuccessfulTests++;
	} else {
		cout << colorRed << "\t[FAIL]" << colorNormal << endl;
	}

	gNumTests++;
}

CBigNum
GetTestModulus()
{
	static CBigNum testModulus(0);

	// TODO: should use a hard-coded RSA modulus for testing
	if (!testModulus) {
		CBigNum p, q;

		// Note: we are NOT using safe primes for testing because
		// they take too long to generate. Don't do this in real
		// usage. See the paramgen utility for better code.
		p = CBigNum::generatePrime(1024, false);
		q = CBigNum::generatePrime(1024, false);
		testModulus = p * q;
	}

	return testModulus;
}

//////////
// Test routines
//////////

bool
Test_GenRSAModulus()
{
	CBigNum result = GetTestModulus();

	if (!result) {
		return false;
	}
	else {
		return true;
	}
}

bool
Test_CalcParamSizes()
{
	bool result = true;
#if 0

	uint32_t pLen, qLen;

	try {
		calculateGroupParamLengths(4000, 80, &pLen, &qLen);
		if (pLen < 1024 || qLen < 256) {
			result = false;
		}
		calculateGroupParamLengths(4000, 96, &pLen, &qLen);
		if (pLen < 2048 || qLen < 256) {
			result = false;
		}
		calculateGroupParamLengths(4000, 112, &pLen, &qLen);
		if (pLen < 3072 || qLen < 320) {
			result = false;
		}
		calculateGroupParamLengths(4000, 120, &pLen, &qLen);
		if (pLen < 3072 || qLen < 320) {
			result = false;
		}
		calculateGroupParamLengths(4000, 128, &pLen, &qLen);
		if (pLen < 3072 || qLen < 320) {
			result = false;
		}
	} catch (exception &e) {
		result = false;
	}
#endif

	return result;
}

bool
Test_GenerateGroupParams()
{
	uint32_t pLen = 1024, qLen = 256, count;
	IntegerGroupParams group;

	for (count = 0; count < 1; count++) {

		try {
			group = deriveIntegerGroupParams(calculateSeed(GetTestModulus(), "test", ZEROCOIN_DEFAULT_SECURITYLEVEL, "TEST GROUP"), pLen, qLen);
		} catch (std::runtime_error e) {
			cout << "Caught exception " << e.what() << endl;
			return false;
		}

		// Now perform some simple tests on the resulting parameters
		if ((uint32_t)group.groupOrder.bitSize() < qLen || (uint32_t)group.modulus.bitSize() < pLen) {
			return false;
		}

		CBigNum c = group.g.pow_mod(group.groupOrder, group.modulus);
		//cout << "g^q mod p = " << c << endl;
		if (!(c.isOne())) return false;

		// Try at multiple parameter sizes
		pLen = pLen * 1.5;
		qLen = qLen * 1.5;
	}

	return true;
}

bool
Test_ParamGen()
{
	bool result = true;

	try {
		// Instantiating testParams runs the parameter generation code
		ZerocoinParams testParams(GetTestModulus(),ZEROCOIN_DEFAULT_SECURITYLEVEL);
	} catch (runtime_error e) {
		cout << e.what() << endl;
		result = false;
	}

	return result;
}

bool
Test_Accumulator()
{
	// This test assumes a list of coins were generated during
	// the Test_MintCoin() test.
	if (gCoins[0] == NULL) {
		return false;
	}
	try {
		// Accumulate the coin list from first to last into one accumulator
            Accumulator accOne(&g_Params->accumulatorParams, CoinDenomination::ZQ_ONE);
            Accumulator accTwo(&g_Params->accumulatorParams,CoinDenomination::ZQ_ONE);
            Accumulator accThree(&g_Params->accumulatorParams,CoinDenomination::ZQ_ONE);
            Accumulator accFour(&g_Params->accumulatorParams,CoinDenomination::ZQ_ONE);
		AccumulatorWitness wThree(g_Params, accThree, gCoins[0]->getPublicCoin());

		for (uint32_t i = 0; i < TESTS_COINS_TO_ACCUMULATE; i++) {
			accOne += gCoins[i]->getPublicCoin();
			accTwo += gCoins[TESTS_COINS_TO_ACCUMULATE - (i+1)]->getPublicCoin();
			accThree += gCoins[i]->getPublicCoin();
			wThree += gCoins[i]->getPublicCoin();
			if(i != 0) {
				accFour += gCoins[i]->getPublicCoin();
			}
		}

		// Compare the accumulated results
		if (accOne.getValue() != accTwo.getValue() || accOne.getValue() != accThree.getValue()) {
			cout << "Accumulators don't match" << endl;
			return false;
		}

		if(accFour.getValue() != wThree.getValue()) {
			cout << "Witness math not working," << endl;
			return false;
		}

		// Verify that the witness is correct
		if (!wThree.VerifyWitness(accThree, gCoins[0]->getPublicCoin()) ) {
			cout << "Witness not valid" << endl;
			return false;
		}

		// Serialization test: see if we can serialize the accumulator
		CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
		ss << accOne;

		// Deserialize it into a new object
		Accumulator newAcc(g_Params, ss);

		// Compare the results
		if (accOne.getValue() != newAcc.getValue()) {
			return false;
		}

	} catch (runtime_error e) {
		return false;
	}

	return true;
}

bool
Test_EqualityPoK()
{
	// Run this test 10 times
	for (uint32_t i = 0; i < 10; i++) {
		try {
			// Generate a random integer "val"
			CBigNum val = CBigNum::randBignum(g_Params->coinCommitmentGroup.groupOrder);

			// Manufacture two commitments to "val", both
			// under different sets of parameters
			Commitment one(&g_Params->accumulatorParams.accumulatorPoKCommitmentGroup, val);

			Commitment two(&g_Params->serialNumberSoKCommitmentGroup, val);

			// Now generate a proof of knowledge that "one" and "two" are
			// both commitments to the same value
			CommitmentProofOfKnowledge pok(&g_Params->accumulatorParams.accumulatorPoKCommitmentGroup,
			                               &g_Params->serialNumberSoKCommitmentGroup,
			                               one, two);

			// Serialize the proof into a stream
			CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
			ss << pok;

			// Deserialize back into a PoK object
			CommitmentProofOfKnowledge newPok(&g_Params->accumulatorParams.accumulatorPoKCommitmentGroup,
			                                  &g_Params->serialNumberSoKCommitmentGroup,
			                                  ss);

			if (newPok.Verify(one.getCommitmentValue(), two.getCommitmentValue()) != true) {
				return false;
			}

			// Just for fun, deserialize the proof a second time
			CDataStream ss2(SER_NETWORK, PROTOCOL_VERSION);
			ss2 << pok;

			// This time tamper with it, then deserialize it back into a PoK
			ss2[15] = 0;
			CommitmentProofOfKnowledge newPok2(&g_Params->accumulatorParams.accumulatorPoKCommitmentGroup,
			                                   &g_Params->serialNumberSoKCommitmentGroup,
			                                   ss2);

			// If the tampered proof verifies, that's a failure!
			if (newPok2.Verify(one.getCommitmentValue(), two.getCommitmentValue()) == true) {
				return false;
			}

		} catch (runtime_error &e) {
			return false;
		}
	}

	return true;
}

bool
Test_MintCoin()
{
	gCoinSize = 0;

	try {
		// Generate a list of coins
		for (uint32_t i = 0; i < TESTS_COINS_TO_ACCUMULATE; i++) {
            gCoins[i] = new PrivateCoin(g_Params,libzerocoin::CoinDenomination::ZQ_ONE);

			PublicCoin pc = gCoins[i]->getPublicCoin();
			CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
			ss << pc;
			gCoinSize += ss.size();
		}

		gCoinSize /= TESTS_COINS_TO_ACCUMULATE;

	} catch (exception &e) {
		return false;
	}

	return true;
}

bool Test_InvalidCoin()
{
	CBigNum coinValue;

	try {
		// Pick a random non-prime CBigNum
		for (uint32_t i = 0; i < NON_PRIME_TESTS; i++) {
			coinValue = CBigNum::randBignum(g_Params->coinCommitmentGroup.modulus);
			coinValue = coinValue * 2;
			if (!coinValue.isPrime()) break;
		}

		PublicCoin pubCoin(g_Params);
		if (pubCoin.validate()) {
			// A blank coin should not be valid!
			return false;
		}

		PublicCoin pubCoin2(g_Params, coinValue, ZQ_ONE);
		if (pubCoin2.validate()) {
			// A non-prime coin should not be valid!
			return false;
		}

		PublicCoin pubCoin3 = pubCoin2;
		if (pubCoin2.validate()) {
			// A copy of a non-prime coin should not be valid!
			return false;
		}

		// Serialize and deserialize the coin
		CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
		ss << pubCoin;
		PublicCoin pubCoin4(g_Params, ss);
		if (pubCoin4.validate()) {
			// A deserialized copy of a non-prime coin should not be valid!
			return false;
		}

	} catch (runtime_error &e) {
		cout << "Caught exception: " << e.what() << endl;
		return false;
	}

	return true;
}

bool
Test_MintAndSpend()
{
	try {
		// This test assumes a list of coins were generated in Test_MintCoin()
		if (gCoins[0] == NULL)
		{
			// No coins: mint some.
			Test_MintCoin();
			if (gCoins[0] == NULL) {
				return false;
			}
		}

		// Accumulate the list of generated coins into a fresh accumulator.
		// The first one gets marked as accumulated for a witness, the
		// others just get accumulated normally.
        Accumulator acc(&g_Params->accumulatorParams,CoinDenomination::ZQ_ONE);
		AccumulatorWitness wAcc(g_Params, acc, gCoins[0]->getPublicCoin());

		for (uint32_t i = 0; i < TESTS_COINS_TO_ACCUMULATE; i++) {
			acc += gCoins[i]->getPublicCoin();
			wAcc +=gCoins[i]->getPublicCoin();
		}

		// Now spend the coin
		//SpendMetaData m(1,1);
		CDataStream cc(SER_NETWORK, PROTOCOL_VERSION);
		cc << *gCoins[0];
		PrivateCoin myCoin(g_Params,cc);

		CoinSpend spend(g_Params, g_Params, myCoin, acc, 0, wAcc, 0, SpendType::SPEND);
        spend.Verify(acc);

		// Serialize the proof and deserialize into newSpend
		CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
		ss << spend;
		gProofSize = ss.size();
		CoinSpend newSpend(g_Params, g_Params, ss);

		// See if we can verify the deserialized proof (return our result)
		bool ret =  newSpend.Verify(acc);

		// Extract the serial number
		CBigNum serialNumber = newSpend.getCoinSerialNumber();
		gSerialNumberSize = ceil((double)serialNumber.bitSize() / 8.0);

		return ret;
	} catch (runtime_error &e) {
		cout << e.what() << endl;
		return false;
	}

	return false;
}

void
Test_RunAllTests()
{
	// Make a new set of parameters from a random RSA modulus
	g_Params = new ZerocoinParams(GetTestModulus());

	gNumTests = gSuccessfulTests = gProofSize = 0;
	for (uint32_t i = 0; i < TESTS_COINS_TO_ACCUMULATE; i++) {
		gCoins[i] = NULL;
	}

	// Run through all of the Zerocoin tests
	LogTestResult("an RSA modulus can be generated", Test_GenRSAModulus);
	LogTestResult("parameter sizes are correct", Test_CalcParamSizes);
	LogTestResult("group/field parameters can be generated", Test_GenerateGroupParams);
	LogTestResult("parameter generation is correct", Test_ParamGen);
	LogTestResult("coins can be minted", Test_MintCoin);
	LogTestResult("invalid coins will be rejected", Test_InvalidCoin);
	LogTestResult("the accumulator works", Test_Accumulator);
	LogTestResult("the commitment equality PoK works", Test_EqualityPoK);
	LogTestResult("a minted coin can be spent", Test_MintAndSpend);

	cout << endl << "Average coin size is " << gCoinSize << " bytes." << endl;
	cout << "Serial number size is " << gSerialNumberSize << " bytes." << endl;
	cout << "Spend proof size is " << gProofSize << " bytes." << endl;

	// Summarize test results
	if (gSuccessfulTests < gNumTests) {
		cout << endl << "ERROR: SOME TESTS FAILED" << endl;
	}

	// Clear any generated coins
	for (uint32_t i = 0; i < TESTS_COINS_TO_ACCUMULATE; i++) {
		delete gCoins[i];
	}

	cout << endl << gSuccessfulTests << " out of " << gNumTests << " tests passed." << endl << endl;
	delete g_Params;
}

BOOST_AUTO_TEST_SUITE(libzerocoin)
BOOST_AUTO_TEST_CASE(libzerocoin_tests)
{
	cout << "libzerocoin v" << ZEROCOIN_VERSION_STRING << " test utility." << endl << endl;

	Test_RunAllTests();
}
BOOST_AUTO_TEST_SUITE_END()


//---------------------------------------------------------------------------//
// Copyright (c) 2013 Kyle Lutz <kyle.r.lutz@gmail.com>
//
// 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://boostorg.github.com/compute for more information.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE TestFunctionalUnpack
#include <boost/test/unit_test.hpp>

#include <boost/compute/system.hpp>
#include <boost/compute/algorithm/copy_n.hpp>
#include <boost/compute/algorithm/transform.hpp>
#include <boost/compute/container/vector.hpp>
#include <boost/compute/iterator/zip_iterator.hpp>
#include <boost/compute/functional/detail/unpack.hpp>

#include "check_macros.hpp"
#include "context_setup.hpp"

namespace compute = boost::compute;

BOOST_AUTO_TEST_CASE(plus_int)
{
    int data1[] = { 1, 3, 5, 7 };
    int data2[] = { 2, 4, 6, 8 };

    compute::vector<int> input1(4, context);
    compute::vector<int> input2(4, context);

    compute::copy_n(data1, 4, input1.begin(), queue);
    compute::copy_n(data2, 4, input2.begin(), queue);

    compute::vector<int> output(4, context);
    compute::transform(
        compute::make_zip_iterator(
            boost::make_tuple(input1.begin(), input2.begin())
        ),
        compute::make_zip_iterator(
            boost::make_tuple(input1.end(), input2.end())
        ),
        output.begin(),
        compute::detail::unpack(compute::plus<int>()),
        queue
    );
    CHECK_RANGE_EQUAL(int, 4, output, (3, 7, 11, 15));
}

BOOST_AUTO_TEST_CASE(fma_float)
{
    float data1[] = { 1, 3, 5, 7 };
    float data2[] = { 2, 4, 6, 8 };
    float data3[] = { 0, 9, 1, 2 };

    compute::vector<float> input1(4, context);
    compute::vector<float> input2(4, context);
    compute::vector<float> input3(4, context);

    compute::copy_n(data1, 4, input1.begin(), queue);
    compute::copy_n(data2, 4, input2.begin(), queue);
    compute::copy_n(data3, 4, input3.begin(), queue);

    compute::vector<int> output(4, context);
    compute::transform(
        compute::make_zip_iterator(
            boost::make_tuple(input1.begin(), input2.begin(), input3.begin())
        ),
        compute::make_zip_iterator(
            boost::make_tuple(input1.end(), input2.end(), input3.begin())
        ),
        output.begin(),
        compute::detail::unpack(compute::fma<float>()),
        queue
    );
}

BOOST_AUTO_TEST_CASE(subtract_int2)
{
    using compute::int2_;

    int data[] = { 4, 2, 5, 1, 6, 3, 7, 0 };
    compute::vector<int2_> input(4, context);
    compute::copy_n(reinterpret_cast<int2_ *>(data), 4, input.begin(), queue);

    compute::vector<int> output(4, context);
    compute::transform(
        input.begin(),
        input.end(),
        output.begin(),
        compute::detail::unpack(compute::minus<int>()),
        queue
    );
    CHECK_RANGE_EQUAL(int, 4, output, (2, 4, 3, 7));
}

BOOST_AUTO_TEST_SUITE_END()


//
// Created by Cristobal M on 10/7/17.
//
#include <iostream>
#include <vector>
#include <fstream>
#include "Rectangle.h"
#include "SplitHeuristic.h"

#include "RandomRectangleGenerator.h"
#include "IOControl.h"

#include "Experiments.h"
#include "FilenameGenerator.h"

#include <sys/resource.h>
#include <boost/filesystem/operations.hpp>

#include <boost/geometry.hpp>
#include <boost/foreach.hpp>

namespace bg = boost::geometry;
namespace bgi = boost::geometry::index;
typedef bg::model::point<float, 2, bg::cs::cartesian> point;
typedef bg::model::box<point> box;
typedef std::pair<box, unsigned> value;

void processInputBoostExperiment(std::string fname, bgi::rtree< value, bgi::quadratic<113> > rtree){
    std::ifstream infile(fname);
    int index = 0;
    while (infile) {
        std::string s;
        if (!getline(infile, s)) break;
        std::istringstream ss(s);
        std::vector <std::string> record;
        int count = 0;
        float minX, maxX, minY, maxY;
        minX = minY = std::numeric_limits<float>::infinity();
        maxX = maxY = -std::numeric_limits<float>::infinity();
        while (ss){
            std::string s2;
            if (!getline( ss, s2, ',')){
                break;
            }
            float coordinate = (float) ::atof(s2.c_str());
            if(count % 2 == 0){
                if(minX > coordinate){
                    minX = coordinate;
                }
                if(maxX < coordinate){
                    maxX = coordinate;
                }
            }
            else{
                if(minY > coordinate){
                    minY = coordinate;
                }
                if(maxY < coordinate){
                    maxY = coordinate;
                }
            }
            count++;
        }
        Rectangle mbr(minX, maxX, minY, maxY, index);
        box b(point(mbr.x1, mbr.y1), point(mbr.x2, mbr.y2));
        rtree.insert(std::make_pair(b, index));
        index++;
    }
    if (!infile.eof()) {
        std::cerr << "Error during input file processing\n";
    }

}
void writeToFile(std::string uuid, std::string file_name, std::string fileheader, std::vector<std::string> contents){
    std::string output_folder = "output_folder";
    if(!boost::filesystem::exists(output_folder)){
        boost::filesystem::create_directory(output_folder);
    }

    std::ofstream ofs;
    ofs.open(output_folder+"/"+uuid+file_name);
    ofs << fileheader;
    ofs << "\n";
    for(std::string &s : contents){
        ofs << s;
        ofs << "\n";
    }
    ofs.close();
}

void boostRTreeExperiment(){

    namespace bg = boost::geometry;
    namespace bgi = boost::geometry::index;
    typedef bg::model::point<float, 2, bg::cs::cartesian> point;
    typedef bg::model::box<point> box;
    typedef std::pair<box, unsigned> value;

    std::vector<std::string> buildTimeOutput;
    std::vector<std::string> searchTimeOutput;

    for (int i = 9; i <= 25; i++){
        bgi::rtree< value, bgi::quadratic<113> > rtree;
        std::string input = "generated2e" + std::to_string(i)  + ".txt";
        const int inputSize = 1 << i; // (int)std::pow(2, i);
        std::cout << "input size:" << inputSize << std::endl;
        RandomRectangleGenerator::generateDataFile(inputSize, input);
        vRect randomVRect = RandomRectangleGenerator::generateVRect(inputSize/10);
        for(int j = 0; j <= 1; j++) {
            Experiments::startExperiment("Build Rtree");
            processInputBoostExperiment(input, rtree);
            ExperimentData &experiment = Experiments::stopExperiment();

            std::string buildTimeStr, searchOutputStr;

            long buildTime =experiment.retrieveTimeElapsedNanoSeconds();
            buildTimeStr = std::to_string(inputSize)+","+
                           std::to_string(buildTime);

            //search
            long sumElapsedSearch = 0;

            std::vector<box> boxes;

            for(Rectangle &randomRectangle : randomVRect){
                box query_box(point(randomRectangle.x1, randomRectangle.y1), point(randomRectangle.x2, randomRectangle.y2));
                boxes.push_back(query_box);
            }

            for(box &randomBox : boxes){
                std::vector<value> result_s;
                Experiments::startExperiment("search");
                rtree.query(bgi::intersects(randomBox), std::back_inserter(result_s));
                ExperimentData &searchExperiment = Experiments::stopExperiment();
                sumElapsedSearch += searchExperiment.retrieveTimeElapsedNanoSeconds();
            }
            long averageDuration = sumElapsedSearch / randomVRect.size();

            searchOutputStr = std::to_string(inputSize)+","+
                              std::to_string(averageDuration);

            buildTimeOutput.push_back(buildTimeStr);
            searchTimeOutput.push_back(searchOutputStr);

            std::cout << "Build Time = ";
            std::cout << buildTime << std::endl;
            std::cout << "Search time = ";
            std::cout << averageDuration << std::endl << std::endl;

        }
        std::cout << std::endl << std::endl;
    }

    std::string anUuid = FilenameGenerator::makeUuid();
    writeToFile(anUuid, "experiment_build", "inputsize,heuristic,timeelapsed", buildTimeOutput);
    writeToFile(anUuid, "experiment_searchtime", "inputsize,heuristic,timeelapsed", searchTimeOutput);
}

void increaseStackSize(){
    const rlim_t kStackSize = 2000 * 1024 * 1024;   // min stack size = 100 MB
    struct rlimit rl;
    int result;

    result = getrlimit(RLIMIT_STACK, &rl);
    if (result == 0)
    {
        if (rl.rlim_cur < kStackSize)
        {
            rl.rlim_cur = kStackSize;
            result = setrlimit(RLIMIT_STACK, &rl);
            if (result != 0)
            {
                fprintf(stderr, "setrlimit returned result = %d\n", result);
            }
        }
    }

}

int main(int argc, char * argv[]){
    increaseStackSize();
    boostRTreeExperiment();
    return 0;
}


/**
 * @file testLinearContainerFactor.cpp
 *
 * @date Jul 6, 2012
 * @author Alex Cunningham
 */

#include <CppUnitLite/TestHarness.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/nonlinear/LinearContainerFactor.h>
#include <gtsam/linear/VectorValues.h>
#include <gtsam/linear/HessianFactor.h>
#include <gtsam/geometry/Pose2.h>
#include <gtsam/geometry/Point3.h>
#include <gtsam/base/TestableAssertions.h>
#include <boost/assign/std/vector.hpp>

using namespace std;
using namespace boost::assign;
using namespace gtsam;

const gtsam::noiseModel::Diagonal::shared_ptr diag_model2 = noiseModel::Diagonal::Sigmas((Vector(2) << 1.0, 1.0));
const double tol = 1e-5;

gtsam::Key  l1 = 101, l2 = 102, x1 = 1, x2 = 2;

Point2 landmark1(5.0, 1.5), landmark2(7.0, 1.5);
Pose2 poseA1(0.0, 0.0, 0.0), poseA2(2.0, 0.0, 0.0);

/* ************************************************************************* */
TEST( testLinearContainerFactor, generic_jacobian_factor ) {

  Matrix A1 = (Matrix(2, 2) <<
      2.74222, -0.0067457,
      0.0,  2.63624);
  Matrix A2 = (Matrix(2, 2) <<
      -0.0455167, -0.0443573,
      -0.0222154, -0.102489);
  Vector b = (Vector(2) << 0.0277052,
      -0.0533393);

  JacobianFactor expLinFactor(l1, A1, l2, A2, b, diag_model2);

  LinearContainerFactor actFactor(expLinFactor);
  EXPECT_LONGS_EQUAL(2, actFactor.size());
  EXPECT(actFactor.isJacobian());
  EXPECT(!actFactor.isHessian());

  // check keys
  FastVector<Key> expKeys; expKeys += l1, l2;
  EXPECT(assert_container_equality(expKeys, actFactor.keys()));

  Values values;
  values.insert(l1, landmark1);
  values.insert(l2, landmark2);
  values.insert(x1, poseA1);
  values.insert(x2, poseA2);

  // Check reconstruction
  GaussianFactor::shared_ptr actLinearizationA = actFactor.linearize(values);
  EXPECT(assert_equal(*expLinFactor.clone(), *actLinearizationA, tol));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, jacobian_factor_withlinpoints ) {

  Matrix A1 = (Matrix(2, 2) <<
      2.74222, -0.0067457,
      0.0,  2.63624);
  Matrix A2 = (Matrix(2, 2) <<
      -0.0455167, -0.0443573,
      -0.0222154, -0.102489);
  Vector b = (Vector(2) << 0.0277052,
      -0.0533393);

  JacobianFactor expLinFactor(l1, A1, l2, A2, b, diag_model2);

  Values values;
  values.insert(l1, landmark1);
  values.insert(l2, landmark2);
  values.insert(x1, poseA1);
  values.insert(x2, poseA2);

  LinearContainerFactor actFactor(expLinFactor, values);
  LinearContainerFactor actFactorNolin(expLinFactor);

  EXPECT(assert_equal(actFactor, actFactor, tol));
  EXPECT(assert_inequal(actFactor, actFactorNolin, tol));
  EXPECT(assert_inequal(actFactorNolin, actFactor, tol));

  // Check contents
  Values expLinPoint;
  expLinPoint.insert(l1, landmark1);
  expLinPoint.insert(l2, landmark2);
  CHECK(actFactor.linearizationPoint());
  EXPECT(actFactor.hasLinearizationPoint());
  EXPECT(assert_equal(expLinPoint, *actFactor.linearizationPoint()));

  // Check error evaluation
  Vector delta_l1 = (Vector(2) << 1.0, 2.0);
  Vector delta_l2 = (Vector(2) << 3.0, 4.0);

  VectorValues delta = values.zeroVectors();
  delta.at(l1) = delta_l1;
  delta.at(l2) = delta_l2;
  Values noisyValues = values.retract(delta);
  double expError = expLinFactor.error(delta);
  EXPECT_DOUBLES_EQUAL(expError, actFactor.error(noisyValues), tol);
  EXPECT_DOUBLES_EQUAL(expLinFactor.error(values.zeroVectors()), actFactor.error(values), tol);

  // Check linearization with corrections for updated linearization point
  GaussianFactor::shared_ptr actLinearizationB = actFactor.linearize(noisyValues);
  Vector bprime = b - A1 * delta_l1 - A2 * delta_l2;
  JacobianFactor expLinFactor2(l1, A1, l2, A2, bprime, diag_model2);
  EXPECT(assert_equal(*expLinFactor2.clone(), *actLinearizationB, tol));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, generic_hessian_factor ) {
  Matrix G11 = (Matrix(1, 1) << 1.0);
  Matrix G12 = (Matrix(1, 2) << 2.0, 4.0);
  Matrix G13 = (Matrix(1, 3) << 3.0, 6.0, 9.0);

  Matrix G22 = (Matrix(2, 2) << 3.0, 5.0,
                            0.0, 6.0);
  Matrix G23 = (Matrix(2, 3) << 4.0, 6.0, 8.0,
                            1.0, 2.0, 4.0);

  Matrix G33 = (Matrix(3, 3) << 1.0, 2.0, 3.0,
                            0.0, 5.0, 6.0,
                            0.0, 0.0, 9.0);

  Vector g1 = (Vector(1) << -7.0);
  Vector g2 = (Vector(2) << -8.0, -9.0);
  Vector g3 = (Vector(3) <<  1.0,  2.0,  3.0);

  double f = 10.0;

  HessianFactor initFactor(x1, x2, l1,
      G11, G12, G13, g1, G22, G23, g2, G33, g3, f);

  Values values;
  values.insert(l1, landmark1);
  values.insert(l2, landmark2);
  values.insert(x1, poseA1);
  values.insert(x2, poseA2);

  LinearContainerFactor actFactor(initFactor);
  EXPECT(!actFactor.isJacobian());
  EXPECT(actFactor.isHessian());

  GaussianFactor::shared_ptr actLinearization1 = actFactor.linearize(values);
  EXPECT(assert_equal(*initFactor.clone(), *actLinearization1, tol));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, hessian_factor_withlinpoints ) {
  // 2 variable example, one pose, one landmark (planar)
  // Initial ordering: x1, l1

  Matrix G11 = (Matrix(3, 3) <<
      1.0, 2.0, 3.0,
      0.0, 5.0, 6.0,
      0.0, 0.0, 9.0);
  Matrix G12 = (Matrix(3, 2) <<
      1.0, 2.0,
      3.0, 5.0,
      4.0, 6.0);
  Vector g1 = (Vector(3) << 1.0,  2.0,  3.0);

  Matrix G22 = (Matrix(2, 2) <<
        0.5, 0.2,
        0.0, 0.6);

  Vector g2 = (Vector(2) << -8.0, -9.0);

  double f = 10.0;

  // Construct full matrices
  Matrix G(5,5);
  G << G11, G12, Matrix::Zero(2,3), G22;

  HessianFactor initFactor(x1, l1, G11, G12, g1, G22, g2, f);

  Values linearizationPoint, expLinPoints;
  linearizationPoint.insert(l1, landmark1);
  linearizationPoint.insert(x1, poseA1);
  expLinPoints = linearizationPoint;
  linearizationPoint.insert(x2, poseA2);

  LinearContainerFactor actFactor(initFactor, linearizationPoint);
  EXPECT(!actFactor.isJacobian());
  EXPECT(actFactor.isHessian());

  EXPECT(actFactor.hasLinearizationPoint());
  Values actLinPoint = *actFactor.linearizationPoint();
  EXPECT(assert_equal(expLinPoints, actLinPoint));

  // Create delta
  Vector delta_l1 = (Vector(2) << 1.0, 2.0);
  Vector delta_x1 = (Vector(3) << 3.0, 4.0, 0.5);
  Vector delta_x2 = (Vector(3) << 6.0, 7.0, 0.3);

  // Check error calculation
  VectorValues delta = linearizationPoint.zeroVectors();
  delta.at(l1) = delta_l1;
  delta.at(x1) = delta_x1;
  delta.at(x2) = delta_x2;
  EXPECT(assert_equal((Vector(5) << 3.0, 4.0, 0.5, 1.0, 2.0), delta.vector(initFactor.keys())));
  Values noisyValues = linearizationPoint.retract(delta);

  double expError = initFactor.error(delta);
  EXPECT_DOUBLES_EQUAL(expError, actFactor.error(noisyValues), tol);
  EXPECT_DOUBLES_EQUAL(initFactor.error(linearizationPoint.zeroVectors()), actFactor.error(linearizationPoint), tol);

  // Compute updated versions
  Vector dv = (Vector(5) << 3.0, 4.0, 0.5, 1.0, 2.0);
  Vector g(5); g << g1, g2;
  Vector g_prime = g - G.selfadjointView<Eigen::Upper>() * dv;

  // Check linearization with corrections for updated linearization point
  Vector g1_prime = g_prime.head(3);
  Vector g2_prime = g_prime.tail(2);
  double f_prime = f + dv.transpose() * G.selfadjointView<Eigen::Upper>() * dv - 2.0 * dv.transpose() * g;
  HessianFactor expNewFactor(x1, l1, G11, G12, g1_prime, G22, g2_prime, f_prime);
  EXPECT(assert_equal(*expNewFactor.clone(), *actFactor.linearize(noisyValues), tol));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, creation ) {
  // Create a set of local keys (No robot label)
  Key  l1 = 11, l3 = 13, l5 = 15;

  // create a linear factor
  SharedDiagonal model = noiseModel::Unit::Create(2);
  JacobianFactor::shared_ptr linear_factor(new JacobianFactor(
      l3, eye(2,2), l5, 2.0 * eye(2,2), zero(2), model));

  // create a set of values - build with full set of values
  gtsam::Values full_values, exp_values;
  full_values.insert(l3, Point2(1.0, 2.0));
  full_values.insert(l5, Point2(4.0, 3.0));
  exp_values = full_values;
  full_values.insert(l1, Point2(3.0, 7.0));

  LinearContainerFactor actual(linear_factor, full_values);

  // Verify the keys
  FastVector<Key> expKeys;
  expKeys += l3, l5;
  EXPECT(assert_container_equality(expKeys, actual.keys()));

  // Verify subset of linearization points
  EXPECT(assert_equal(exp_values, actual.linearizationPoint(), tol));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, jacobian_relinearize )
{
  // Create a Between Factor from a Point3. This is actually a linear factor.
  gtsam::Key key1(1);
  gtsam::Key key2(2);
  gtsam::Values linpoint1;
  linpoint1.insert(key1, gtsam::Point3(-22.4,  +8.5,  +2.4));
  linpoint1.insert(key2, gtsam::Point3(-21.0,  +5.0, +21.0));

  gtsam::Point3 measured(1.0, -2.5, 17.8);
  gtsam::SharedNoiseModel model = gtsam::noiseModel::Isotropic::Sigma(3, 0.1);
  gtsam::BetweenFactor<gtsam::Point3> betweenFactor(key1, key2, measured, model);

  // Create a jacobian container factor at linpoint 1
  gtsam::JacobianFactor::shared_ptr jacobian(new gtsam::JacobianFactor(*betweenFactor.linearize(linpoint1)));
  gtsam::LinearContainerFactor jacobianContainer(jacobian, linpoint1);

  // Create a second linearization point
  gtsam::Values linpoint2;
  linpoint2.insert(key1, gtsam::Point3(+18.0, -0.25, +1.11));
  linpoint2.insert(key2, gtsam::Point3(-10.0, +11.2, +0.05));

  // Check the error at linpoint2 versus the original factor
  double expected_error = betweenFactor.error(linpoint2);
  double actual_error = jacobianContainer.error(linpoint2);
  EXPECT_DOUBLES_EQUAL(expected_error, actual_error, 1e-9 );

  // Re-linearize around the new point and check the factors
  gtsam::GaussianFactor::shared_ptr expected_factor = betweenFactor.linearize(linpoint2);
  gtsam::GaussianFactor::shared_ptr actual_factor   = jacobianContainer.linearize(linpoint2);
  CHECK(gtsam::assert_equal(*expected_factor, *actual_factor));
}

/* ************************************************************************* */
TEST( testLinearContainerFactor, hessian_relinearize )
{
  // Create a Between Factor from a Point3. This is actually a linear factor.
  gtsam::Key key1(1);
  gtsam::Key key2(2);
  gtsam::Values linpoint1;
  linpoint1.insert(key1, gtsam::Point3(-22.4,  +8.5,  +2.4));
  linpoint1.insert(key2, gtsam::Point3(-21.0,  +5.0, +21.0));

  gtsam::Point3 measured(1.0, -2.5, 17.8);
  gtsam::SharedNoiseModel model = gtsam::noiseModel::Isotropic::Sigma(3, 0.1);
  gtsam::BetweenFactor<gtsam::Point3> betweenFactor(key1, key2, measured, model);

  // Create a hessian container factor at linpoint 1
  gtsam::HessianFactor::shared_ptr hessian(new gtsam::HessianFactor(*betweenFactor.linearize(linpoint1)));
  gtsam::LinearContainerFactor hessianContainer(hessian, linpoint1);

  // Create a second linearization point
  gtsam::Values linpoint2;
  linpoint2.insert(key1, gtsam::Point3(+18.0, -0.25, +1.11));
  linpoint2.insert(key2, gtsam::Point3(-10.0, +11.2, +0.05));

  // Check the error at linpoint2 versus the original factor
  double expected_error = betweenFactor.error(linpoint2);
  double actual_error = hessianContainer.error(linpoint2);
  EXPECT_DOUBLES_EQUAL(expected_error, actual_error, 1e-9 );

  // Re-linearize around the new point and check the factors
  gtsam::GaussianFactor::shared_ptr expected_factor = gtsam::HessianFactor::shared_ptr(new gtsam::HessianFactor(*betweenFactor.linearize(linpoint2)));
  gtsam::GaussianFactor::shared_ptr actual_factor   = hessianContainer.linearize(linpoint2);
  CHECK(gtsam::assert_equal(*expected_factor, *actual_factor));
}

/* ************************************************************************* */
int main() { TestResult tr; return TestRegistry::runAllTests(tr); }
/* ************************************************************************* */


// Copyright (c) 2020-2021 FRC Team 3512. All Rights Reserved.

#pragma once

#include <array>
#include <tuple>
#include <utility>

#include <Eigen/Core>
#include <frc/logging/CSVLogFile.h>
#include <units/time.h>
#include <wpi/StringRef.h>
#include <wpi/Twine.h>

#include "logging/ControllerLabel.hpp"

namespace frc3512 {

/**
 * Generates CSV files for state-space controllers.
 */
template <int States, int Inputs, int Outputs>
class CSVControllerLogger {
public:
    /**
     * Constructs a CSVControllerLogger.
     *
     * @param controllerName Name of the controller log file.
     * @param stateLabels    Labels for states each consisting of its name and
     *                       unit.
     * @param inputLabels    Labels for inputs each consisting of its name and
     *                       unit.
     * @param outputLabels   Labels for outputs each consisting of its name and
     *                       unit.
     */
    CSVControllerLogger(
        wpi::StringRef controllerName,
        const std::array<ControllerLabel, States>& stateLabels,
        const std::array<ControllerLabel, Inputs>& inputLabels,
        const std::array<ControllerLabel, Outputs>& outputLabels)
        : m_stateLogger{(controllerName + " states").str(),
                        std::tuple_cat(MakeReferenceLabels(stateLabels),
                                       MakeStateEstimateLabels(stateLabels))},
          m_inputLogger{(controllerName + " inputs").str(),
                        MakeInputLabels(inputLabels)},
          m_outputLogger{(controllerName + " outputs").str(),
                         MakeOutputLabels(outputLabels)} {}

    /**
     * Move constructor.
     */
    CSVControllerLogger(CSVControllerLogger&&) = default;

    /**
     * Move assignment operator.
     */
    CSVControllerLogger& operator=(CSVControllerLogger&&) = default;

    /**
     * Logs the references, states, inputs, and outputs.
     *
     * @param time Timestamp for data.
     * @param r    Reference vector.
     * @param x    State vector.
     * @param u    Input vector.
     * @param y    Output vector.
     */
    void Log(units::second_t time, const Eigen::Matrix<double, States, 1>& r,
             const Eigen::Matrix<double, States, 1>& x,
             const Eigen::Matrix<double, Inputs, 1>& u,
             const Eigen::Matrix<double, Outputs, 1>& y) {
        LogVector2WithTime<States, States>(m_stateLogger, time, r, x);
        LogVectorWithTime<Inputs>(m_inputLogger, time, u);
        LogVectorWithTime<Outputs>(m_outputLogger, time, y);
    }

    /**
     * Logs the references, states, inputs, and outputs.
     *
     * @param r    Reference vector.
     * @param x    State vector.
     * @param u    Input vector.
     * @param y    Output vector.
     */
    void Log(const Eigen::Matrix<double, States, 1>& r,
             const Eigen::Matrix<double, States, 1>& x,
             const Eigen::Matrix<double, Inputs, 1>& u,
             const Eigen::Matrix<double, Outputs, 1>& y) {
        LogVector2<States, States>(m_stateLogger, r, x);
        LogVector<Inputs>(m_inputLogger, u);
        LogVector<Outputs>(m_outputLogger, y);
    }

private:
    frc::CSVLogFile m_stateLogger;
    frc::CSVLogFile m_inputLogger;
    frc::CSVLogFile m_outputLogger;

    /**
     * Returns tuple of reference labels from state names.
     */
    static auto MakeReferenceLabels(
        const std::array<ControllerLabel, States>& stateLabels) {
        return std::apply(
            [](auto... labels) {
                return std::make_tuple(
                    (labels.name + " reference (" + labels.unit + ")")...);
            },
            ArrayToTuple(stateLabels));
    }

    /**
     * Returns tuple of state estimate labels from state names.
     */
    static auto MakeStateEstimateLabels(
        const std::array<ControllerLabel, States>& stateLabels) {
        return std::apply(
            [](auto... labels) {
                return std::make_tuple(
                    (labels.name + " estimate (" + labels.unit + ")")...);
            },
            ArrayToTuple(stateLabels));
    }

    /**
     * Returns tuple of input labels from input names.
     */
    static auto MakeInputLabels(
        const std::array<ControllerLabel, Inputs>& inputLabels) {
        return std::apply(
            [](auto... labels) {
                return std::make_tuple(
                    (labels.name + " input (" + labels.unit + ")")...);
            },
            ArrayToTuple(inputLabels));
    }

    /**
     * Returns tuple of output labels from output names.
     */
    static auto MakeOutputLabels(
        const std::array<ControllerLabel, Outputs>& outputLabels) {
        return std::apply(
            [](auto... labels) {
                return std::make_tuple(
                    (labels.name + " measurement (" + labels.unit + ")")...);
            },
            ArrayToTuple(outputLabels));
    }

    template <typename Array, size_t... I>
    static auto ArrayToTupleImpl(const Array& a, std::index_sequence<I...>) {
        return std::make_tuple(a[I]...);
    }

    template <typename T, size_t N,
              typename Indices = std::make_index_sequence<N>>
    static auto ArrayToTuple(const std::array<T, N>& a) {
        return ArrayToTupleImpl(a, Indices{});
    }

    template <typename T, size_t... I>
    void LogVectorImpl(frc::CSVLogFile& logger, const T& vec,
                       std::index_sequence<I...>) {
        return logger.Log(vec(I)...);
    }

    template <typename T, size_t... I>
    void LogVectorWithTimeImpl(frc::CSVLogFile& logger, units::second_t time,
                               const T& vec, std::index_sequence<I...>) {
        return logger.Log(time, vec(I)...);
    }

    template <typename T, size_t... I1, typename U, size_t... I2>
    void LogVector2Impl(frc::CSVLogFile& logger, const T& vec1,
                        std::index_sequence<I1...>, const U& vec2,
                        std::index_sequence<I2...>) {
        return logger.Log(vec1(I1)..., vec2(I2)...);
    }

    template <typename T, size_t... I1, typename U, size_t... I2>
    void LogVector2WithTimeImpl(frc::CSVLogFile& logger, units::second_t time,
                                const T& vec1, std::index_sequence<I1...>,
                                const U& vec2, std::index_sequence<I2...>) {
        return logger.Log(time, vec1(I1)..., vec2(I2)...);
    }

    template <int Rows, typename Indices = std::make_index_sequence<Rows>>
    void LogVector(frc::CSVLogFile& logger,
                   const Eigen::Matrix<double, Rows, 1>& vec) {
        return LogVectorImpl(logger, vec, Indices{});
    }

    template <int Rows, typename Indices = std::make_index_sequence<Rows>>
    void LogVectorWithTime(frc::CSVLogFile& logger, units::second_t time,
                           const Eigen::Matrix<double, Rows, 1>& vec) {
        return LogVectorWithTimeImpl(logger, time, vec, Indices{});
    }

    template <int Rows1, int Rows2,
              typename Indices1 = std::make_index_sequence<Rows1>,
              typename Indices2 = std::make_index_sequence<Rows2>>
    void LogVector2(frc::CSVLogFile& logger,
                    const Eigen::Matrix<double, Rows1, 1>& vec1,
                    const Eigen::Matrix<double, Rows2, 1>& vec2) {
        return LogVector2Impl(logger, vec1, Indices1{}, vec2, Indices2{});
    }

    template <int Rows1, int Rows2,
              typename Indices1 = std::make_index_sequence<Rows1>,
              typename Indices2 = std::make_index_sequence<Rows2>>
    void LogVector2WithTime(frc::CSVLogFile& logger, units::second_t time,
                            const Eigen::Matrix<double, Rows1, 1>& vec1,
                            const Eigen::Matrix<double, Rows2, 1>& vec2) {
        return LogVector2WithTimeImpl(logger, time, vec1, Indices1{}, vec2,
                                      Indices2{});
    }
};

}  // namespace frc3512