#include <stan/math/rev.hpp>
#include <gtest/gtest.h>
#include <test/unit/math/rev/util.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <vector>

template <typename T_x>
std::vector<T_x> fill_vec(Eigen::Matrix<T_x, -1, 1> inp) {
  std::vector<T_x> ret_vec;
  ret_vec.reserve(inp.rows());
  for (int i = 0; i < inp.rows(); ++i)
    ret_vec.push_back(inp(i));
  return ret_vec;
}

template <typename T>
Eigen::Matrix<T, -1, -1> create_mat(Eigen::VectorXd inp, T alpha, T len,
                                    T jitter) {
  std::vector<double> test_inp = fill_vec(inp);
  Eigen::Matrix<T, -1, -1> test_mat_dense
      = stan::math::cov_exp_quad(test_inp, alpha, len);
  for (int i = 0; i < inp.rows(); ++i)
    test_mat_dense(i, i) = test_mat_dense(i, i) + jitter;
  return test_mat_dense;
}

struct gp_chol {
  Eigen::VectorXd inp, mean, y;
  gp_chol(Eigen::VectorXd inp_, Eigen::VectorXd mean_, Eigen::VectorXd y_)
      : inp(inp_), mean(mean_), y(y_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    Eigen::Matrix<T, -1, -1> x_c = create_mat(inp, x[0], x[1], x[2]);
    Eigen::Matrix<T, -1, -1> L = stan::math::cholesky_decompose(x_c);
    T lp = stan::math::multi_normal_cholesky_lpdf(y, mean, L);
    return lp;
  }
};

struct chol_functor {
  int i, j, K;
  chol_functor(int i_, int j_, int K_) : i(i_), j(j_), K(K_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    using stan::math::cholesky_decompose;
    using stan::math::cov_matrix_constrain;
    T lp(0.0);
    Eigen::Matrix<T, -1, -1> x_c = cov_matrix_constrain(x, K, lp);
    Eigen::Matrix<T, -1, -1> L = cholesky_decompose(x_c);
    lp += L(i, j);
    return lp;
  }
};

struct chol_functor_mult_scal {
  int K;
  Eigen::VectorXd vec;
  chol_functor_mult_scal(int K_, Eigen::VectorXd vec_) : K(K_), vec(vec_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    using stan::math::cholesky_decompose;
    using stan::math::cov_matrix_constrain;
    using stan::math::multiply;
    using stan::math::transpose;
    T lp(0.0);
    Eigen::Matrix<T, -1, -1> x_c = cov_matrix_constrain(x, K, lp);
    Eigen::Matrix<T, -1, -1> L = cholesky_decompose(x_c);
    lp += multiply(transpose(vec), multiply(L, vec));
    return lp;
  }
};

struct chol_functor_2 {
  int K;
  explicit chol_functor_2(int K_) : K(K_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    using stan::math::cholesky_decompose;
    using stan::math::cov_matrix_constrain;
    using stan::math::multi_normal_cholesky_log;
    T lp(0.0);
    Eigen::Matrix<T, -1, -1> x_c = cov_matrix_constrain(x, K, lp);
    Eigen::Matrix<T, -1, -1> L = cholesky_decompose(x_c);
    Eigen::Matrix<double, -1, 1> vec(K);
    Eigen::Matrix<double, -1, 1> mu(K);
    vec.setZero();
    mu.setOnes();
    lp += multi_normal_cholesky_log(vec, mu, L);
    return lp;
  }
};

struct chol_functor_simple {
  int i, j, K;
  chol_functor_simple(int i_, int j_, int K_) : i(i_), j(j_), K(K_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    using stan::math::cholesky_decompose;
    Eigen::Matrix<T, -1, -1> x_c(K, K);
    int pos = 0;
    for (int n = 0; n < K; ++n)
      for (int m = 0; m < K; ++m) {
        x_c(m, n) = x(pos++);
        x_c(n, m) = x_c(m, n);
      }
    Eigen::Matrix<T, -1, -1> L = cholesky_decompose(x_c);
    return L(i, j);
  }
};

struct chol_functor_simple_vec {
  int K;
  Eigen::VectorXd vec;
  chol_functor_simple_vec(int K_, Eigen::VectorXd vec_) : K(K_), vec(vec_) {}
  template <typename T>
  T operator()(Eigen::Matrix<T, -1, 1> x) const {
    using stan::math::cholesky_decompose;
    using stan::math::multiply;
    using stan::math::transpose;
    Eigen::Matrix<T, -1, -1> x_c(K, K);
    int pos = 0;
    for (int n = 0; n < K; ++n)
      for (int m = 0; m < K; ++m) {
        x_c(m, n) = x(pos++);
        x_c(n, m) = x_c(m, n);
      }
    Eigen::Matrix<T, -1, -1> L = cholesky_decompose(x_c);
    T lp = multiply(transpose(vec), multiply(L, vec));
    return lp;
  }
};

void test_gradients(int size, double prec) {
  std::vector<std::vector<chol_functor> > functors;
  std::vector<std::vector<Eigen::Matrix<double, -1, 1> > > grads_ad;
  std::vector<std::vector<Eigen::Matrix<double, -1, 1> > > grads_fd;
  Eigen::Matrix<double, -1, -1> evals_ad(size, size);
  Eigen::Matrix<double, -1, -1> evals_fd(size, size);
  functors.resize(size);
  grads_ad.resize(size);
  grads_fd.resize(size);

  for (int i = 0; i < size; ++i)
    for (int j = 0; j < size; ++j) {
      functors[i].push_back(chol_functor(i, j, size));
      grads_fd[i].push_back(Eigen::Matrix<double, -1, 1>(size));
      grads_ad[i].push_back(Eigen::Matrix<double, -1, 1>(size));
    }

  int numels = size + size * (size - 1) / 2;
  Eigen::Matrix<double, -1, 1> x(numels);
  for (int i = 0; i < numels; ++i)
    x(i) = i % 10 / 100.0;

  for (size_t i = 0; i < static_cast<size_t>(size); ++i) {
    for (size_t j = 0; j < static_cast<size_t>(size); ++j) {
      stan::math::gradient(functors[i][j], x, evals_ad(i, j), grads_ad[i][j]);
      stan::math::finite_diff_gradient(functors[i][j], x, evals_fd(i, j),
                                       grads_fd[i][j]);

      for (int k = 0; k < numels; ++k)
        EXPECT_NEAR(grads_fd[i][j](k), grads_ad[i][j](k), prec);
      EXPECT_FLOAT_EQ(evals_fd(i, j), evals_ad(i, j));
    }
  }
}

void test_gradients_simple(int size, double prec) {
  std::vector<std::vector<chol_functor_simple> > functors;
  std::vector<std::vector<Eigen::Matrix<double, -1, 1> > > grads_ad;
  std::vector<std::vector<Eigen::Matrix<double, -1, 1> > > grads_fd;
  Eigen::Matrix<double, -1, -1> evals_ad(size, size);
  Eigen::Matrix<double, -1, -1> evals_fd(size, size);
  functors.resize(size);
  grads_ad.resize(size);
  grads_fd.resize(size);

  for (int i = 0; i < size; ++i)
    for (int j = 0; j < size; ++j) {
      functors[i].push_back(chol_functor_simple(i, j, size));
      grads_fd[i].push_back(Eigen::Matrix<double, -1, 1>(size));
      grads_ad[i].push_back(Eigen::Matrix<double, -1, 1>(size));
    }

  stan::math::welford_covar_estimator estimator(size);

  boost::random::mt19937 rng;
  for (int i = 0; i < 1000; ++i) {
    Eigen::VectorXd q(size);
    for (int j = 0; j < size; ++j)
      q(j) = stan::math::normal_rng(0.0, 1.0, rng);
    estimator.add_sample(q);
  }

  Eigen::MatrixXd covar(size, size);
  estimator.sample_covariance(covar);

  Eigen::Matrix<double, -1, 1> x(size * size);
  int pos = 0;
  for (int j = 0; j < size; ++j)
    for (int i = 0; i < size; ++i)
      x(pos++) = covar(i, j);

  for (size_t j = 0; j < static_cast<size_t>(size); ++j) {
    for (size_t i = j; i < static_cast<size_t>(size); ++i) {
      stan::math::gradient(functors[i][j], x, evals_ad(i, j), grads_ad[i][j]);
      stan::math::finite_diff_gradient(functors[i][j], x, evals_fd(i, j),
                                       grads_fd[i][j]);

      for (int k = 0; k < size; ++k)
        EXPECT_NEAR(grads_fd[i][j](k), grads_ad[i][j](k), prec);
      EXPECT_FLOAT_EQ(evals_fd(i, j), evals_ad(i, j));
    }
  }
}

void test_gp_grad(int mat_size, double prec) {
  using Eigen::MatrixXd;
  using Eigen::RowVectorXd;
  using Eigen::VectorXd;
  using stan::math::var;

  Eigen::VectorXd test_vec(mat_size);
  Eigen::VectorXd mean_vec(mat_size);
  Eigen::VectorXd draw_vec(mat_size);
  Eigen::VectorXd test_vals(3);
  test_vals[0] = 1;
  test_vals[1] = 1.5;
  test_vals[2] = 10;

  boost::random::mt19937 rng(2);

  for (int i = 0; i < mat_size; ++i) {
    test_vec(i) = stan::math::normal_rng(0.0, 0.1, rng);
    mean_vec(i) = 0;
    draw_vec(i) = stan::math::normal_rng(0.0, 0.1, rng);
  }

  Eigen::MatrixXd cov_mat
      = create_mat(test_vec, test_vals[0], test_vals[1], test_vals[2]);
  Eigen::MatrixXd chol_cov = cov_mat.llt().matrixL();
  Eigen::VectorXd y_vec = chol_cov * draw_vec;

  gp_chol gp_fun(test_vec, mean_vec, y_vec);
  double val_ad;
  Eigen::VectorXd grad_ad;
  stan::math::gradient(gp_fun, test_vals, val_ad, grad_ad);

  VectorXd grad_fd;
  double val_fd;

  stan::math::finite_diff_gradient(gp_fun, test_vals, val_fd, grad_fd);
  EXPECT_NEAR(val_fd, val_ad, 1e-10);
  for (int i = 0; i < grad_ad.size(); ++i) {
    EXPECT_NEAR(grad_fd(i), grad_ad(i), prec);
  }
}

void test_chol_mult(int mat_size, double prec) {
  using Eigen::MatrixXd;
  using Eigen::RowVectorXd;
  using Eigen::VectorXd;
  using stan::math::var;

  int vec_size = mat_size * (mat_size + 1) / 2;
  Eigen::VectorXd test_vec(mat_size);
  Eigen::VectorXd test_vals(vec_size);

  boost::random::mt19937 rng(2);

  for (int i = 0; i < test_vals.size(); ++i) {
    if (i < test_vec.size()) {
      test_vec(i) = stan::math::normal_rng(0.0, 0.1, rng);
    }
    test_vals(i) = i % 10 / 100.0;
  }

  chol_functor_mult_scal mult_fun(mat_size, test_vec);
  double val_ad;
  Eigen::VectorXd grad_ad;
  stan::math::gradient(mult_fun, test_vals, val_ad, grad_ad);

  VectorXd grad_fd;
  double val_fd;

  stan::math::finite_diff_gradient(mult_fun, test_vals, val_fd, grad_fd);
  EXPECT_NEAR(val_fd, val_ad, 1e-10);
  for (int i = 0; i < grad_ad.size(); ++i) {
    EXPECT_NEAR(grad_fd(i), grad_ad(i), prec);
  }
}

void test_simple_vec_mult(int size, double prec) {
  Eigen::VectorXd test_vec(size);
  boost::random::mt19937 rng(2);

  for (int i = 0; i < test_vec.size(); ++i)
    test_vec(i) = stan::math::normal_rng(0.0, 0.1, rng);

  chol_functor_simple_vec f(size, test_vec);

  stan::math::welford_covar_estimator estimator(size);

  for (int i = 0; i < 1000; ++i) {
    Eigen::VectorXd q(size);
    for (int j = 0; j < size; ++j)
      q(j) = stan::math::normal_rng(0.0, 1.0, rng);
    estimator.add_sample(q);
  }

  Eigen::MatrixXd covar(size, size);
  estimator.sample_covariance(covar);

  Eigen::Matrix<double, -1, 1> x(size * size);
  int pos = 0;
  for (int j = 0; j < size; ++j)
    for (int i = 0; i < size; ++i)
      x(pos++) = covar(i, j);

  double eval_ad;
  Eigen::VectorXd grad_ad;
  stan::math::gradient(f, x, eval_ad, grad_ad);
  double eval_fd;
  Eigen::VectorXd grad_fd;
  stan::math::finite_diff_gradient(f, x, eval_fd, grad_fd);

  EXPECT_FLOAT_EQ(eval_fd, eval_ad);
  for (int k = 0; k < grad_fd.size(); ++k)
    EXPECT_NEAR(grad_fd(k), grad_ad(k), prec) << " for k=" << k;
}

double test_gradient(int size, double prec) {
  chol_functor_2 functown(size);
  Eigen::Matrix<double, -1, 1> grads_ad;
  Eigen::Matrix<double, -1, 1> grads_fd;
  double evals_ad;
  double evals_fd;

  int numels = size + size * (size - 1) / 2;
  Eigen::Matrix<double, -1, 1> x(numels);
  for (int i = 0; i < numels; ++i)
    x(i) = i / 100.0;

  stan::math::gradient(functown, x, evals_ad, grads_ad);
  stan::math::finite_diff_gradient(functown, x, evals_fd, grads_fd);

  for (int k = 0; k < numels; ++k)
    EXPECT_NEAR(grads_fd(k), grads_ad(k), prec) << " for k=" << k;
  EXPECT_FLOAT_EQ(evals_fd, evals_ad);
  return grads_ad.sum();
}

TEST(AgradRevMatrix, mat_cholesky) {
  using stan::math::cholesky_decompose;
  using stan::math::matrix_v;
  using stan::math::singular_values;
  using stan::math::transpose;

  // symmetric
  matrix_v X(2, 2);
  stan::math::var a = 3.0;
  stan::math::var b = -1.0;
  stan::math::var c = -1.0;
  stan::math::var d = 1.0;
  X << a, b, c, d;

  matrix_v L = cholesky_decompose(X);

  matrix_v LL_trans = multiply(L, transpose(L));
  EXPECT_FLOAT_EQ(a.val(), LL_trans(0, 0).val());
  EXPECT_FLOAT_EQ(b.val(), LL_trans(0, 1).val());
  EXPECT_FLOAT_EQ(c.val(), LL_trans(1, 0).val());
  EXPECT_FLOAT_EQ(d.val(), LL_trans(1, 1).val());

  EXPECT_NO_THROW(singular_values(X));
}

TEST(AgradRevMatrix, exception_mat_cholesky) {
  stan::math::matrix_v m;

  // not positive definite
  m.resize(2, 2);
  m << 1.0, 2.0, 2.0, 3.0;
  EXPECT_THROW(stan::math::cholesky_decompose(m), std::domain_error);

  // zero size
  m.resize(0, 0);
  EXPECT_NO_THROW(stan::math::cholesky_decompose(m));

  // not square
  m.resize(2, 3);
  EXPECT_THROW(stan::math::cholesky_decompose(m), std::invalid_argument);

  // not symmetric
  m.resize(2, 2);
  m << 1.0, 2.0, 3.0, 4.0;
  EXPECT_THROW(stan::math::cholesky_decompose(m), std::domain_error);
}

TEST(AgradRevMatrix, exception_varmat_cholesky) {
  stan::math::matrix_d m;

  // not positive definite
  m.resize(2, 2);
  m << 1.0, 2.0, 2.0, 3.0;
  stan::math::var_value<stan::math::matrix_d> mv1(m);
  EXPECT_THROW(stan::math::cholesky_decompose(mv1), std::domain_error);

  // zero size
  m.resize(0, 0);
  stan::math::var_value<stan::math::matrix_d> mv2(m);
  EXPECT_NO_THROW(stan::math::cholesky_decompose(mv2));

  // not square
  m.resize(2, 3);
  stan::math::var_value<stan::math::matrix_d> mv3(m);
  EXPECT_THROW(stan::math::cholesky_decompose(mv3), std::invalid_argument);

  // not symmetric
  m.resize(2, 2);
  m << 1.0, 2.0, 3.0, 4.0;
  stan::math::var_value<stan::math::matrix_d> mv4(m);
  EXPECT_THROW(stan::math::cholesky_decompose(mv4), std::domain_error);
}

TEST(AgradRevMatrix, mat_cholesky_1st_deriv_small) {
  test_gradients(9, 1e-10);
  test_gradients_simple(10, 1e-10);
  test_gradient(15, 1e-10);
  test_gp_grad(20, 1e-10);
}

TEST(AgradRevMatrix, check_varis_on_stack_small) {
  stan::math::matrix_v X(2, 2);
  X << 3, -1, -1, 1;

  test::check_varis_on_stack(stan::math::cholesky_decompose(X));
}

TEST(AgradRevMatrix, mat_cholesky_1st_deriv_large_gradients) {
  test_gradient(36, 1e-08);
  test_gp_grad(100, 1e-08);
  test_gp_grad(1000, 1e-08);
  test_chol_mult(37, 1e-08);
  test_simple_vec_mult(45, 1e-08);
}

TEST(AgradRevMatrix, cholesky_replicated_input) {
  using stan::math::var;

  auto f = [](int size, const auto& y) {
    auto m = stan::math::diag_matrix(stan::math::rep_vector(y, size));
    auto L = stan::math::cholesky_decompose(m);
    return stan::math::sum(L);
  };

  double ydbl = 1.5;
  double dx = 1e-5;
  var y = ydbl;
  int size = 4;
  var s = f(size, y);
  s.grad();

  double fd_ref = (f(size, ydbl + dx) - f(size, ydbl - dx)) / (2.0 * dx);
  EXPECT_FLOAT_EQ(y.adj(), fd_ref);

  stan::math::set_zero_all_adjoints();
  size = 40;
  s = f(size, y);
  s.grad();

  fd_ref = (f(size, ydbl + dx) - f(size, ydbl - dx)) / (2.0 * dx);
  EXPECT_FLOAT_EQ(y.adj(), fd_ref);
}


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

///////////////////////////////////////////////////////////////////////////////////////////////////
#include <gtest/gtest.h>
#include "api/CPP/memory.hpp"
#include <api/CPP/input_layout.hpp>
#include "api/CPP/lstm.hpp"
#include <api/CPP/split.hpp>
#include <api/CPP/crop.hpp>
#include <api/CPP/concatenation.hpp>
#include <api/CPP/topology.hpp>
#include <api/CPP/tensor.hpp>
#include <api/CPP/network.hpp>
#include <api/CPP/engine.hpp>
#include "test_utils/test_utils.h"
#include <api/CPP/data.hpp>
#include "instrumentation.h"
#include <boost/filesystem.hpp>

#include <sstream>
#include <iomanip>

using namespace cldnn;
using namespace tests;

#define FERROR 1E-4

namespace {
    float sigmoid(float x) {
        return 1.f / (1.f + (float)std::exp((float)(-x)));
    }
}

struct offset_order {
    size_t it, ot, ft, zt;
    offset_order(size_t scale, const cldnn_lstm_offset_order& t = cldnn_lstm_offset_order_iofz) {
        static const std::map<cldnn_lstm_offset_order, std::vector<size_t>> offset_map{
            { cldnn_lstm_offset_order_iofz,{ 0, 1, 2, 3 } },
            { cldnn_lstm_offset_order_ifoz,{ 0, 2, 1, 3 } }
        };
        std::vector<size_t> v = offset_map.at(t);
        it = v[0] * scale;
        ot = v[1] * scale;
        ft = v[2] * scale;
        zt = v[3] * scale;
    }
};
cldnn_lstm_offset_order default_offset_type = cldnn_lstm_offset_order_iofz;

// [ARIEL] TODO: use move semantics when required

template<typename T>
T clip(T val, T threshold) {
    if (threshold > 0) {
        if (val > threshold) return threshold;
        if (val < -threshold) return -threshold;
    }
    return val;
}

template <typename T>
VVVVF<T> lstm_elt_reference(VVVVF<T>& tempGEMM, VVVVF<T>& cell, bool hasCell = true, float clip_threshold = 0, bool input_forget = false) {
    size_t hidden_size = tempGEMM[0][0][0].size() / 4;
    size_t batch_size = tempGEMM.size();
    VVVVF<T> tempOut(batch_size, VVVF<T>(2, VVF<T>(1, VF<T>(hidden_size))));
    offset_order off(hidden_size, default_offset_type);

    for (size_t b = 0; b < batch_size; ++b) {
        T *it = &tempGEMM[b][0][0][off.it];
        T *ot = &tempGEMM[b][0][0][off.ot];
        T *ft = &tempGEMM[b][0][0][off.ft];
        T *zt = &tempGEMM[b][0][0][off.zt];
        for (size_t h = 0; h < hidden_size; ++h) {
            T val = sigmoid(clip(it[h], clip_threshold)) * std::tanh((float)clip(zt[h], clip_threshold));
            if (input_forget) {
                val *= (1 - ft[h]);
            }
            if (hasCell) {
                val += cell[b][0][0][h] * sigmoid(clip(ft[h], clip_threshold));
            }
            tempOut[b][0][0][h] = std::tanh((float)val) * sigmoid(ot[h]);
            tempOut[b][1][0][h] = val;
        }
    }
    return tempOut;
}

template <typename T>
VVVVF<T> lstm_gemm_reference(VVVVF<T>& input, VVVVF<T>& weights, VVVVF<T>& recurrent, VVVVF<T>& bias, VVVVF<T>& hidden,
    bool hasBias = true, bool hasHidden = true) {
    size_t input_size = input[0][0][0].size();
    size_t hidden_size = hidden[0][0][0].size();
    size_t batch_size = input.size();

    // Temporary output from GEMM operations [f, i, o, z]
    VVVVF<T> tempGEMM(batch_size, VVVF<T>(1, VVF<T>(1, VF<T>(4 * hidden_size))));
    for (size_t b = 0; b < batch_size; ++b) {
        for (size_t y = 0; y < 4 * hidden_size; ++y) {
            T res = 0;
            for (size_t x = 0; x < input_size; ++x) {
                res += (T)weights[0][0][y][x] * (T)input[b][0][0][x];
            }
            if (hasHidden) {
                for (size_t x = 0; x < hidden_size; ++x) {
                    res += (T)recurrent[0][0][y][x] * (T)hidden[b][0][0][x];
                }
            }
            if (hasBias) {
                res += (T)bias[0][0][0][y];
            }
            tempGEMM[b][0][0][y] = res;
        }
    }
    return tempGEMM;
}

template<typename T>
void print(const std::string& s, VVVVF<T>& input) {
    printf("%s -------------\n", s.c_str());
    printf("Size = [%d, %d, %d, %d]\n", (int)input.size(), (int)input[0].size(), (int)input[0][0].size(), (int)input[0][0][0].size());
    for (size_t b = 0; b < input.size(); ++b) {
        for (size_t f = 0; f < input[0].size(); ++f) {
            for (size_t y = 0; y < input[0][0].size(); ++y) {
                for (size_t x = 0; x < input[0][0][0].size(); ++x) {
                    printf("%f ", input[b][f][y][x]);
                }
                printf("\n");
            }
        }
    }
    printf("---------------------------------------\n");
}

template<typename T>
VVVVF<T> lstm_split_reference(VVVVF<T>& input, size_t idx, size_t bufferId) {
    VVVVF<T> tempOut;
    switch (idx) {
    case 0:
        tempOut = VVVVF<T>(input.size(), VVVF<T>(input[0].size(), VVF<T>(1, VF<T>(input[0][0][0].size()))));
        for (size_t i = 0; i < input.size(); i++)
            tempOut[i][0] = input[i][bufferId];
        break;
    case 1:
        tempOut = VVVVF<T>(input.size(), VVVF<T>(1, VVF<T>(1, VF<T>(input[0][0][0].size()))));
        //tempOut[0][0] = input[0][bufferId];
        for (size_t i = 0; i < input.size(); i++)
            tempOut[i][0] = input[i][bufferId];
        break;
    case 2:
        tempOut = VVVVF<T>(1, VVVF<T>(1, VVF<T>(1, VF<T>(input[0][0][0].size()))));
        tempOut[0][0][0] = input[0][0][bufferId];
        break;
    }
    return tempOut;
}

template <typename T>
void lstm_reference(VVVVF<T>& input, VVVVF<T>& hidden, VVVVF<T>& cell, VVVVF<T>& weights, VVVVF<T>& recurrent, VVVVF<T>& bias,
    VVVVF<T>& output, VVVVF<T>& last_hidden, VVVVF<T>& last_cell,
    bool hasBias = true, bool hasInitialHidden = true, bool hasInitialCell = true,
    float clip_threshold = 0, bool input_forget = false) {

    size_t sequence_len = input[0].size();
    size_t dir_len = weights[0].size();
    size_t batch = input.size();
    for (size_t dir = 0; dir < dir_len; ++dir) {
        for (size_t seq = 0; seq < sequence_len; ++seq) {
            VVVVF<T> splitInput = lstm_split_reference(input, 1, seq);
            VVVVF<T> tempGEMM = lstm_gemm_reference(splitInput, weights, recurrent, bias, hidden, hasBias, hasInitialHidden);
            VVVVF<T> tempOutput = lstm_elt_reference(tempGEMM, cell, hasInitialCell, clip_threshold, input_forget);
            for (size_t i = 0; i < batch; i++)
                output[i][seq] = tempOutput[i][0]; // hidden, output[dir,seq] = tempOutput[0,dir,batch,hidden]
            hidden = lstm_split_reference(tempOutput, 0, 0);
            cell = lstm_split_reference(tempOutput, 0, 1);
            hasInitialHidden = true;
            hasInitialCell = true;
        }
    }

    last_hidden = hidden;
    last_cell = cell;
}

template<typename T>
void generic_lstm_gemm_gpu_test(int sequence_len, int direction, int batch_size, int input_size, int hidden_size,
    bool hasBias = true, bool hasHidden = true) {
    int min_random = -2, max_random = 2;

    VVVVF<T> ref_input = generate_random_4d<T>(batch_size, sequence_len, 1, input_size, min_random, max_random);
    VVVVF<T> ref_weights = generate_random_4d<T>(1, direction, 4 * hidden_size, input_size, min_random, max_random);
    VVVVF<T> ref_recurrent = generate_random_4d<T>(1, direction, 4 * hidden_size, hidden_size, min_random, max_random);
    VVVVF<T> ref_bias = generate_random_4d<T>(1, 1, direction, 4 * hidden_size, min_random, max_random);
    VVVVF<T> ref_hidden = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VF<T> ref_input_vec = flatten_4d<T>(cldnn::format::bfyx, ref_input);
    VF<T> ref_weights_vec = flatten_4d<T>(cldnn::format::bfyx, ref_weights);
    VF<T> ref_recurrent_vec = flatten_4d<T>(cldnn::format::bfyx, ref_recurrent);
    VF<T> ref_bias_vec = flatten_4d<T>(cldnn::format::bfyx, ref_bias);
    VF<T> ref_hidden_vec = flatten_4d<T>(cldnn::format::bfyx, ref_hidden);

    VVVVF<T> ref_output = lstm_gemm_reference(ref_input, ref_weights, ref_recurrent, ref_bias, ref_hidden, hasBias, hasHidden);

    engine engine;
    memory input = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,   sequence_len,  input_size,      1 } });
    memory weights = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,            direction,     input_size,      4 * hidden_size } });
    memory recurrent = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,            direction,     hidden_size,     4 * hidden_size } });
    memory biases = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,            1,             4 * hidden_size, direction } });
    memory hidden = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,   direction,     hidden_size,     1 } });

    set_values(input, ref_input_vec);
    set_values(weights, ref_weights_vec);
    set_values(recurrent, ref_recurrent_vec);
    set_values(biases, ref_bias_vec);
    set_values(hidden, ref_hidden_vec);

    topology topology;
    topology.add(input_layout("input", input.get_layout()));
    topology.add(data("weights", weights));
    topology.add(data("recurrent", recurrent));
    if (hasBias) {
        topology.add(data("biases", biases));
    }
    if (hasHidden) {
        topology.add(input_layout("hidden", hidden.get_layout()));
    }

    topology.add(lstm_gemm("lstm_gemm", "input", "weights", "recurrent", hasBias ? "biases" : "", hasHidden ? "hidden" : ""));

    network network(engine, topology);
    network.set_input_data("input", input);
    if (hasHidden) {
        network.set_input_data("hidden", hidden);
    }

    auto outputs = network.execute();
    EXPECT_EQ(outputs.size(), size_t(1));

    auto output = outputs.begin()->second.get_memory();
    auto output_ptr = output.pointer<T>();
    int i = 0;
    for (int b = 0; b < batch_size; ++b) {
        for (int x = 0; x < 4 * hidden_size; ++x)
            EXPECT_EQ(ref_output[b][0][0][x], output_ptr[i++]);
    }
}

template<typename T>
void generic_lstm_elt_gpu_test(int sequence_len, int direction, int batch_size, int input_size, int hidden_size, bool hasCell = true,
    float clip_threshold = 0.f, bool input_forget = false) {
    // tempGEMM  = [        1, direction,           batch, 4 * hidden_size ] input
    // cell      = [        1, direction,           batch,     hidden_size ] optional
    // output    = [        2, direction,           batch,     hidden_size ] output concat[hidden, cell]
    int min_random = -2, max_random = 2;

    VVVVF<T> ref_tempGEMM = generate_random_4d<T>(batch_size, direction, 1, 4 * hidden_size, min_random, max_random);
    VVVVF<T> ref_cell = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VF<T> ref_tempGEMM_vec = flatten_4d<T>(cldnn::format::bfyx, ref_tempGEMM);
    VF<T> ref_cell_vec = flatten_4d<T>(cldnn::format::bfyx, ref_cell);

    VVVVF<T> ref_output = lstm_elt_reference(ref_tempGEMM, ref_cell, hasCell, clip_threshold, input_forget);

    engine engine;
    memory tempGEMM = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,    direction, 4 * hidden_size, 1 } });
    memory cell = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,    direction,     hidden_size, 1 } });
    set_values(tempGEMM, ref_tempGEMM_vec);
    set_values(cell, ref_cell_vec);

    topology topology;
    topology.add(input_layout("tempGEMM", tempGEMM.get_layout()));
    if (hasCell) {
        topology.add(input_layout("cell", cell.get_layout()));
    }
    topology.add(lstm_elt("lstm_elt", "tempGEMM", hasCell ? "cell" : "", clip_threshold, input_forget));

    network network(engine, topology);
    network.set_input_data("tempGEMM", tempGEMM);
    if (hasCell) {
        network.set_input_data("cell", cell);
    }

    auto outputs = network.execute();
    EXPECT_EQ(outputs.size(), size_t(1));

    auto output = outputs.begin()->second.get_memory();
    auto output_ptr = output.pointer<T>();
    for (int b = 0; b < batch_size; ++b) {
        for (int j = 0; j < 2; ++j) {
            for (int x = 0; x < hidden_size; ++x)
            {
                auto idx = b * 2 * hidden_size + j * hidden_size + x;
                EXPECT_NEAR(ref_output[b][j][0][x], output_ptr[idx], FERROR);
            }
        }
    }
}

std::string getIdString(size_t i) {
    std::stringstream ss;
    ss << std::setw(5) << std::setfill('0') << i;
    return ss.str();
}

// --------------- Manually constructed LSTM ----------------------------------------
// This function manually generates an lstm node sequence by conbining lstm_gemm and lstm_elt nodes
// it requires that the output of the lstm_elt node is croped to obtain the corresponding hidden and cell outputs
void generate_lstm_topology(topology& t, memory& input, memory& hidden, memory& cell,
    memory& weights, memory& recurrent, memory& biases, int sequence_len,
    bool hasBias = true, bool hasInitialHidden = true, bool hasInitialCell = true) {
    auto hidden_size = hidden.get_layout().size;
    t.add(input_layout("input", input.get_layout()));
    std::vector<std::pair<primitive_id, tensor>> input_ids_offsets;
    std::vector<primitive_id> output_ids_offsets;
    for (int i = 0; i < sequence_len; ++i)
        input_ids_offsets.push_back({ getIdString(i),{ 0, i, 0, 0 } });
    t.add(split("inputSplit", "input", input_ids_offsets));
    t.add(data("weights", weights));
    t.add(data("recurrent", recurrent));

    std::string biasStr = "";
    std::string hiddenStr = "";
    std::string cellStr = "";
    if (hasBias)
    {
        t.add(data("biases", biases));
        biasStr = "biases";
    }
    if (hasInitialHidden)
    {
        t.add(input_layout("hidden", hidden.get_layout()));
        hiddenStr = "hidden";
    }
    if (hasInitialCell)
    {
        t.add(input_layout("cell", cell.get_layout()));
        cellStr = "cell";
    }
    for (int i = 0; i < sequence_len; ++i) {
        std::string lstm_gemm_id = "lstm_gemm" + getIdString(i);
        std::string lstm_elt_id = "lstm_elt" + getIdString(i);
        std::string crop_id = "crop" + getIdString(i);

        t.add(lstm_gemm(lstm_gemm_id, "inputSplit:" + getIdString(i), "weights", "recurrent", biasStr, hiddenStr));
        t.add(lstm_elt(lstm_elt_id, lstm_gemm_id, cellStr));

        hiddenStr = crop_id + ":hidden";
        t.add(crop(hiddenStr, lstm_elt_id, hidden_size, tensor{ 0,0,0,0 }));
        if (i < sequence_len - 1) {
            cellStr = crop_id + ":cell";
            t.add(crop(cellStr, lstm_elt_id, hidden_size, tensor{ 0,1,0,0 }));
        }
        output_ids_offsets.push_back(hiddenStr);
    }
    t.add(concatenation("concatenation", output_ids_offsets, concatenation::along_f));
}

template<typename T>
void generic_lstm_custom_gpu_test(int sequence_len, int direction, int batch_size, int input_size, int hidden_size,
    bool hasBias = true, bool hasInitialHidden = true, bool hasInitialCell = true) {
    std::cout << "Input Size = " << input_size << " Hidden Size = " << hidden_size << " Sequence Len = " << sequence_len << " Batch Size = " << batch_size << std::endl;
    int min_random = -2, max_random = 2;
    VVVVF<T> ref_input = generate_random_4d<T>(batch_size, sequence_len, 1, input_size, min_random, max_random);
    VVVVF<T> ref_weights = generate_random_4d<T>(1, direction, 4 * hidden_size, input_size, min_random, max_random);
    VVVVF<T> ref_recurrent = generate_random_4d<T>(1, direction, 4 * hidden_size, hidden_size, min_random, max_random);
    VVVVF<T> ref_bias = generate_random_4d<T>(1, 1, direction, 4 * hidden_size, min_random, max_random);
    VVVVF<T> ref_hidden = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VVVVF<T> ref_cell = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VVVVF<T> ref_output(batch_size, VVVF<T>(sequence_len, VVF<T>(direction, VF<T>(hidden_size))));
    VVVVF<T> last_hidden(batch_size, VVVF<T>(direction, VVF<T>(1, VF<T>(hidden_size))));
    VVVVF<T> last_cell(batch_size, VVVF<T>(direction, VVF<T>(1, VF<T>(hidden_size))));

    VF<T> ref_input_vec = flatten_4d<T>(cldnn::format::bfyx, ref_input);
    VF<T> ref_weights_vec = flatten_4d<T>(cldnn::format::bfyx, ref_weights);
    VF<T> ref_recurrent_vec = flatten_4d<T>(cldnn::format::bfyx, ref_recurrent);
    VF<T> ref_bias_vec = flatten_4d<T>(cldnn::format::bfyx, ref_bias);
    VF<T> ref_hidden_vec = flatten_4d<T>(cldnn::format::bfyx, ref_hidden);
    VF<T> ref_cell_vec = flatten_4d<T>(cldnn::format::bfyx, ref_cell);
    lstm_reference(ref_input, ref_hidden, ref_cell, ref_weights, ref_recurrent, ref_bias, ref_output, last_hidden, last_cell,
        hasBias, hasInitialHidden, hasInitialCell);

    engine engine;
    memory input = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size, sequence_len,  input_size,       1 } });
    memory weights = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,          direction,     input_size,       4 * hidden_size } });
    memory recurrent = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,          direction,     hidden_size,      4 * hidden_size } });
    memory biases = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,          1,             4 * hidden_size,  direction } });
    memory hidden = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size, direction,     hidden_size,      1 } });
    memory cell = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size, direction,     hidden_size,      1 } });
    set_values(input, ref_input_vec);
    set_values(weights, ref_weights_vec);
    set_values(recurrent, ref_recurrent_vec);
    set_values(biases, ref_bias_vec);
    set_values(hidden, ref_hidden_vec);
    set_values(cell, ref_cell_vec);

    topology topology;
    generate_lstm_topology(topology, input, hidden, cell, weights, recurrent, biases, sequence_len,
        hasBias, hasInitialHidden, hasInitialCell);

    network network(engine, topology);
    network.set_input_data("input", input);
    if (hasInitialHidden) network.set_input_data("hidden", hidden);
    if (hasInitialCell) network.set_input_data("cell", cell);
    auto outputs = network.execute();

    ASSERT_EQ(outputs.size(), size_t(1));
    size_t output_size = outputs.begin()->second.get_memory().size() / sizeof(T);
    ASSERT_EQ(output_size, size_t(hidden_size * sequence_len * batch_size * direction));

    auto output = outputs.begin()->second.get_memory();
    auto output_ptr = output.pointer<T>();
    int i = 0;
    for (int b = 0; b < batch_size; ++b) {
        for (int s = 0; s < sequence_len; ++s) {
            for (int x = 0; x < hidden_size; ++x) {
                for (int d = 0; d < direction; ++d) {
                    ASSERT_NEAR(ref_output[b][s][d][x], output_ptr[i++], FERROR);
                }
            }
        }
    }
}

// -------------------------------------------------------

template<typename T>
void generic_lstm_gpu_test(int sequence_len, int direction, int batch_size, int input_size, int hidden_size,
    bool hasBias = true, bool hasInitialHidden = true, bool hasInitialCell = true,
    float clip_threshold = 0, bool input_forget = false) {
    std::cout << "Input Size = " << input_size << " Hidden Size = " << hidden_size << " Sequence Len = " << sequence_len << " Batch Size = " << batch_size << std::endl;
    int min_random = -2, max_random = 2;
    VVVVF<T> ref_input = generate_random_4d<T>(batch_size, sequence_len, 1, input_size, min_random, max_random);
    VVVVF<T> ref_weights = generate_random_4d<T>(1, direction, 4 * hidden_size, input_size, min_random, max_random);
    VVVVF<T> ref_recurrent = generate_random_4d<T>(1, direction, 4 * hidden_size, hidden_size, min_random, max_random);
    VVVVF<T> ref_bias = generate_random_4d<T>(1, 1, direction, 4 * hidden_size, min_random, max_random);
    VVVVF<T> ref_hidden = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VVVVF<T> ref_cell = generate_random_4d<T>(batch_size, direction, 1, hidden_size, min_random, max_random);
    VVVVF<T> ref_output(batch_size, VVVF<T>(sequence_len, VVF<T>(direction, VF<T>(hidden_size))));
    VVVVF<T> last_hidden(batch_size, VVVF<T>(direction, VVF<T>(1, VF<T>(hidden_size))));
    VVVVF<T> last_cell(batch_size, VVVF<T>(direction, VVF<T>(1, VF<T>(hidden_size))));

    VF<T> ref_input_vec = flatten_4d<T>(cldnn::format::bfyx, ref_input);
    VF<T> ref_weights_vec = flatten_4d<T>(cldnn::format::bfyx, ref_weights);
    VF<T> ref_recurrent_vec = flatten_4d<T>(cldnn::format::bfyx, ref_recurrent);
    VF<T> ref_bias_vec = flatten_4d<T>(cldnn::format::bfyx, ref_bias);
    VF<T> ref_hidden_vec = flatten_4d<T>(cldnn::format::bfyx, ref_hidden);
    VF<T> ref_cell_vec = flatten_4d<T>(cldnn::format::bfyx, ref_cell);
    lstm_reference(ref_input, ref_hidden, ref_cell, ref_weights, ref_recurrent, ref_bias, ref_output, last_hidden, last_cell,
        hasBias, hasInitialHidden, hasInitialCell, clip_threshold, input_forget);

    engine engine;

    memory input = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,    sequence_len,   input_size,      1 } });
    memory weights = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,             direction,      input_size,      4 * hidden_size } });
    memory recurrent = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,             direction,      hidden_size,     4 * hidden_size } });
    memory biases = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ 1,             1,              4 * hidden_size, direction } });
    memory hidden = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,    direction,      hidden_size,     1 } });
    memory cell = memory::allocate(engine, { type_to_data_type<T>::value, format::bfyx,{ batch_size,    direction,      hidden_size,     1 } });

    set_values(input, ref_input_vec);
    set_values(weights, ref_weights_vec);
    set_values(recurrent, ref_recurrent_vec);
    if (hasBias) set_values(biases, ref_bias_vec);
    if (hasInitialHidden) set_values(hidden, ref_hidden_vec);
    if (hasInitialCell) set_values(cell, ref_cell_vec);

    topology topology;
    std::vector<std::pair<primitive_id, tensor>> input_ids_offsets;
    std::vector<primitive_id> lstm_inputs;
    std::vector<primitive_id> output_ids_offsets;

    topology.add(input_layout("input", input.get_layout()));
    for (int i = 0; i < sequence_len; ++i) {
        input_ids_offsets.push_back({ getIdString(i),{ 0, i, 0, 0 } });
        lstm_inputs.push_back("inputSplit:" + getIdString(i));
    }
    topology.add(split("inputSplit", "input", input_ids_offsets));
    topology.add(data("weights", weights));
    topology.add(data("recurrent", recurrent));
    if (hasBias) topology.add(data("biases", biases));
    if (hasInitialHidden) topology.add(input_layout("hidden", hidden.get_layout()));
    if (hasInitialCell) topology.add(input_layout("cell", cell.get_layout()));
    topology.add(lstm("lstm", lstm_inputs, "weights", "recurrent",
        hasBias ? "biases" : "", hasInitialHidden ? "hidden" : "", hasInitialCell ? "cell" : "", "", clip_threshold, input_forget, {}, {}, default_offset_type));

    network network(engine, topology);
    network.set_input_data("input", input);
    if (hasInitialHidden) network.set_input_data("hidden", hidden);
    if (hasInitialCell) network.set_input_data("cell", cell);
    auto outputs = network.execute();

    ASSERT_EQ(outputs.size(), size_t(1));
    size_t output_size = outputs.begin()->second.get_memory().size() / sizeof(T);
    ASSERT_EQ(output_size, size_t(hidden_size * sequence_len * batch_size * direction));

    auto output = outputs.begin()->second.get_memory();
    auto output_ptr = output.pointer<T>();
    int i = 0;
    for (int b = 0; b < batch_size; ++b) {
        for (int s = 0; s < sequence_len; ++s) {
            for (int x = 0; x < hidden_size; ++x) {
                for (int d = 0; d < direction; ++d) {
                    ASSERT_NEAR(ref_output[b][s][d][x], output_ptr[i++], FERROR);
                }
            }
        }
    }
}

TEST(lstm_gemm_gpu, generic_lstm_gemm_test_f32) {
    generic_lstm_gemm_gpu_test<float>(1, 1, 3, 6, 2, true, true);
}

TEST(lstm_gemm_gpu, generic_lstm_gemm_no_bias_f32) {
    generic_lstm_gemm_gpu_test<float>(1, 1, 3, 6, 2, false, true);
}

TEST(lstm_gemm_gpu, generic_lstm_gemm_no_hidden_f32) {
    generic_lstm_gemm_gpu_test<float>(1, 1, 3, 6, 2, true, false);
}

TEST(lstm_gemm_gpu, generic_lstm_gemm_no_hidden_bias_f32) {
    generic_lstm_gemm_gpu_test<float>(1, 1, 3, 6, 2, false, false);
}

TEST(lstm_elt_gpu, generic_lstm_elt_test_clip_f32) {
    generic_lstm_elt_gpu_test<float>(1, 1, 4, 6, 3, true, 0.3f);
}

TEST(lstm_elt_gpu, generic_lstm_elt_test_input_forget_f32) {
    generic_lstm_elt_gpu_test<float>(1, 1, 4, 6, 3, true, 0.f, 1);
}

TEST(lstm_elt_gpu, generic_lstm_elt_test_clip_input_forget_f32) {
    generic_lstm_elt_gpu_test<float>(1, 1, 4, 6, 3, true, 0.5f, 1);
}

TEST(lstm_elt_gpu, generic_lstm_elt_test_f32) {
    generic_lstm_elt_gpu_test<float>(1, 1, 4, 6, 3, true);
}

TEST(lstm_elt_gpu, generic_lstm_elt_no_cell_f32) {
    generic_lstm_elt_gpu_test<float>(1, 1, 4, 6, 3, false);
}

TEST(lstm_custom_gpu, generic_lstm_custom_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, true, true, true);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_biasf32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, false, true, true);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_hidden_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, true, false, true);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_bias_hidden_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, false, false, true);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_cell_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, true, true, false);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_bias_cell_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, false, true, false);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_hidden_cell_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, true, false, false);
}

TEST(lstm_custom_gpu, generic_lstm_custom_no_bias_hidden_cell_f32) {
    generic_lstm_custom_gpu_test<float>(3, 1, 3, 3, 2, false, false, false);
}

TEST(lstm_gpu, generic_lstm_f32) {
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, true, true, true);
}

TEST(lstm_gpu, generic_lstm_no_bias_f32) {
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, false, true, true);
}

TEST(lstm_gpu, generic_lstm_no_hidden_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, true, false, true);
}

TEST(lstm_gpu, generic_lstm_no_bias_hidden_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, false, false, true);
}

TEST(lstm_gpu, generic_lstm_no_cell_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, true, true, false);
}

TEST(lstm_gpu, generic_lstm_no_bias_cell_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, false, true, false);
}

TEST(lstm_gpu, generic_lstm_no_hidden_cell_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, true, false, false);
}

TEST(lstm_gpu, generic_lstm_no_bias_hidden_cell_f32) {
    generic_lstm_gpu_test<float>(3, 1, 5, 4, 3, false, false, false);
}

TEST(lstm_gpu, generic_lstm_clip_f32) {
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, true, true, true, 0.3f, 0);
}

TEST(lstm_gpu, generic_lstm_input_forget_f32) {
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, true, true, true, 0.f, 1);
}

TEST(lstm_gpu, generic_lstm_clip_input_forget_f32) {
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, true, true, true, 0.3f, 1);
}

TEST(lstm_gpu, generic_lstm_offset_order_ifoz_f32) {
    default_offset_type = cldnn_lstm_offset_order_ifoz;
    generic_lstm_gpu_test<float>(3, 1, 3, 3, 2, true, true, true);
    default_offset_type = cldnn_lstm_offset_order_iofz;
}

// TODO: Add tests for the following:
// optional concatenate output
// optional last hidden
// optional last cell
// optional activation list


//=======================================================================
// Copyright 2001 Jeremy G. Siek, Andrew Lumsdaine, Lie-Quan Lee,
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
#include <boost/config.hpp>
#include <fstream>
#include <iostream>
#include <vector>
#include <iomanip>
#include <boost/property_map/property_map.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/graphviz.hpp>
#include <boost/graph/johnson_all_pairs_shortest.hpp>

int
main()
{
  using namespace boost;
  typedef adjacency_list<vecS, vecS, directedS, no_property,
    property< edge_weight_t, int, property< edge_weight2_t, int > > > Graph;
  const int V = 6;
  typedef std::pair < int, int >Edge;
  Edge edge_array[] =
    { Edge(0, 1), Edge(0, 2), Edge(0, 3), Edge(0, 4), Edge(0, 5),
    Edge(1, 2), Edge(1, 5), Edge(1, 3), Edge(2, 4), Edge(2, 5),
    Edge(3, 2), Edge(4, 3), Edge(4, 1), Edge(5, 4)
  };
  const std::size_t E = sizeof(edge_array) / sizeof(Edge);
#if defined(BOOST_MSVC) && BOOST_MSVC <= 1300
  // VC++ can't handle the iterator constructor
  Graph g(V);
  for (std::size_t j = 0; j < E; ++j)
    add_edge(edge_array[j].first, edge_array[j].second, g);
#else
  Graph g(edge_array, edge_array + E, V);
#endif

  property_map < Graph, edge_weight_t >::type w = get(edge_weight, g);
  int weights[] = { 0, 0, 0, 0, 0, 3, -4, 8, 1, 7, 4, -5, 2, 6 };
  int *wp = weights;

  graph_traits < Graph >::edge_iterator e, e_end;
  for (boost::tie(e, e_end) = edges(g); e != e_end; ++e)
    w[*e] = *wp++;

  std::vector < int >d(V, (std::numeric_limits < int >::max)());
  int D[V][V];
  johnson_all_pairs_shortest_paths(g, D, distance_map(&d[0]));

  std::cout << "       ";
  for (int k = 0; k < V; ++k)
    std::cout << std::setw(5) << k;
  std::cout << std::endl;
  for (int i = 0; i < V; ++i) {
    std::cout << std::setw(3) << i << " -> ";
    for (int j = 0; j < V; ++j) {
      if (D[i][j] == (std::numeric_limits<int>::max)())
        std::cout << std::setw(5) << "inf";
      else
        std::cout << std::setw(5) << D[i][j];
    }
    std::cout << std::endl;
  }

  std::ofstream fout("figs/johnson-eg.dot");
  fout << "digraph A {\n"
    << "  rankdir=LR\n"
    << "size=\"5,3\"\n"
    << "ratio=\"fill\"\n"
    << "edge[style=\"bold\"]\n" << "node[shape=\"circle\"]\n";

  graph_traits < Graph >::edge_iterator ei, ei_end;
  for (boost::tie(ei, ei_end) = edges(g); ei != ei_end; ++ei)
    fout << source(*ei, g) << " -> " << target(*ei, g)
      << "[label=" << get(edge_weight, g)[*ei] << "]\n";

  fout << "}\n";
  return 0;
}


#pragma once

#include <limits>

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <boost/type_traits/is_integral.hpp>
#include <boost/type_traits/is_floating_point.hpp>
#include <boost/type_traits/is_same.hpp>

#include <boost/iterator/counting_iterator.hpp>
#include <boost/iterator/transform_iterator.hpp>

#include "fmmtl/numeric/Vec.hpp"
#include "fmmtl/numeric/Complex.hpp"

namespace fmmtl {

static boost::random::mt19937 default_generator;

using boost::enable_if;
using boost::is_integral;
using boost::is_floating_point;

template <typename T, class Enable = void>
struct random;

template <typename T>
struct random<T, typename enable_if<is_integral<T> >::type> {
  typedef T result_type;
  result_type operator()(T a, T b) const { return get(a,b); }
  result_type operator()()         const { return get();    }

  static result_type get(T a, T b) {
    boost::random::uniform_int_distribution<T> dist(a, b);
    return dist(default_generator);
  }
  static result_type get() {
    return get(T(0), std::numeric_limits<T>::max());
  }
};

template <typename T>
struct random<T, typename enable_if<is_floating_point<T> >::type> {
  typedef T result_type;
  result_type operator()(T a, T b) const { return get(a,b); }
  result_type operator()()         const { return get();    }

  static result_type get(T a, T b) {
    boost::random::uniform_real_distribution<T> dist(a, b);
    return dist(default_generator);
  }
  static result_type get() {
    return get(T(0), T(1));
  }
};

template <typename T>
struct random<complex<T> > {
  typedef complex<T> result_type;
  result_type operator()(T a, T b) const { return get(a,b); }
  result_type operator()()         const { return get();    }

  static result_type get(T a, T b) {
    return complex<T>(random<T>::get(a,b), random<T>::get(a,b));
  }
  static result_type get() {
    return get(T(0), T(1));
  }
};

template <typename T>
struct random<std::complex<T> > {
  typedef std::complex<T> result_type;
  result_type operator()(T a, T b) const { return get(a,b); }
  result_type operator()()         const { return get();    }

  static result_type get(T a, T b) {
    return std::complex<T>(random<T>::get(a,b), random<T>::get(a,b));
  }
  static result_type get() {
    return get(T(0), T(1));
  }
};

template <std::size_t N, typename T>
struct random<Vec<N,T> > {
  typedef Vec<N,T> result_type;
  result_type operator()(T a, T b) const { return get(a,b); }
  result_type operator()()         const { return get();    }

  static result_type get(T a, T b) {
    Vec<N,T> v;
    for (std::size_t i = 0; i != N; ++i)
      v[i] = fmmtl::random<T>::get(a, b);
    return v;
  }
  static result_type get() {
    return get(T(0), T(1));
  }
};

class random_n {
  template <typename T>
  struct generator {
    T operator()(const std::size_t&) const { return random<T>::get(); }
  };
  std::size_t N;

 public:
  random_n(const std::size_t& _N) : N(_N) {}

  template <typename Container>
  operator Container() const {
    typedef typename Container::value_type value_type;
    return Container(
        boost::make_transform_iterator(
            boost::make_counting_iterator(std::size_t(0)),
            generator<value_type>()),
        boost::make_transform_iterator(
            boost::make_counting_iterator(std::size_t(N)),
            generator<value_type>()));
  }
};

} // end namespace fmmtl


#include <iostream>
#include <vector>
#include <Eigen/Dense>
#include "point.h"
#include "node.h"
#include "decision_tree.h"
#include <fstream>

void test_decision_tree()
{
  std::cout << "test_decision_tree\n";

  using Label = int;
  std::vector<Point<Label>> datapoints {{-2, 1.5, 0},
                                       {1, 1, 0},
                                       {0, 1.5, 0},
                                       {0.5, -0.5, 1},
                                       {1, -1, 1},
                                       {-1, -2, 1},
                                       {1, 2, 0}};

  // stream out
  std::ofstream out("data/data.txt");
  for(std::size_t i = 0; i < datapoints.size(); ++i)
  {
    out << datapoints[i].x << " " << datapoints[i].y << " "
        << datapoints[i].label << "\n";
  }
  out.close();

  Node<Label> root(datapoints);
  Decision_tree<Label> tree(&root);
  tree.fit();
}

void test_decision_tree2()
{
  std::cout << "test_decision_tree\n";

  using Label = int;
  std::vector<Point<Label>> datapoints {{0, 1, 0},
                                        {1, 0, 1}};

  // stream out
  std::ofstream out("data/data.txt");
  for(std::size_t i = 0; i < datapoints.size(); ++i)
  {
    out << datapoints[i].x << " " << datapoints[i].y << " "
        << datapoints[i].label << "\n";
  }
  out.close();

  Node<Label> root(datapoints);
  Decision_tree<Label> tree(&root);
  tree.fit();
  // todo: assert line b1 = 1, b0 = 0
}

void test_decision_tree_outlier()
{
  std::cout << "test_decision_tree_outlier\n";

  using Label = int;
  std::vector<Point<Label>> datapoints {{-2, 1.5, 0}, //0
                                       {1, 1, 0},  //1
                                       {0, 1.5, 0}, //2
                                       {0.5, -0.5, 1}, //3
                                       {1, -1, 1}, //4
                                       {-1, -2, 1}, //5
                                       {1, 2, 0}, //6
                                       {-1, 2, 1}}; // outlier

  // stream out
  std::ofstream out("data/data.txt");
  for(std::size_t i = 0; i < datapoints.size(); ++i)
  {
    out << datapoints[i].x << " " << datapoints[i].y << " "
        << datapoints[i].label << "\n";
  }
  out.close();

  Node<Label> root(datapoints);
  Decision_tree<Label> tree(&root);
  tree.fit();
}

int main()
{

  test_decision_tree();
  test_decision_tree2();
  test_decision_tree_outlier();
  return 0;
}


#ifndef mesh_reader
#define mesh_reader

#include <fstream>
#include <sstream>
#include <iostream>
#include <string>
#include <math.h>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>

using namespace boost::numeric::ublas;

namespace mesh {

    matrix<double> read_vertices() {
        /* index # is vertix id, contains x and y
        coordinate respectively of vertix, and
        a flag, 1 if boundary point 0 if not */
        std::ifstream infile("../triangle/circle.1.node");
        std::string line;
        double a, b, c, d, line_nb = 0;
        std::getline(infile, line);
        std::istringstream iss(line);
        iss >> a >> b >> c;
        int nb_nodes = a;
        matrix<double> vertices(nb_nodes, 3);

        while (std::getline(infile, line)) {
            std::istringstream iss(line);
            if (iss >> a >> b >> c >> d) {
                vertices(line_nb, 0) = b;
                vertices(line_nb, 1) = c;
                vertices(line_nb, 2) = d;
            }
            else {
                // std::cout << "ERROR: in .1.node file" << std::endl;
                break;
            }
            line_nb += 1;
        }

        return vertices;
    }

    matrix<double> read_triangles(matrix<double> &vertices) {
        /* index # is triangle id, contains
        the 3 vertix id's of the corner nodes  and
        the area of the triangle, calculated using
        the vertices matrix */
        std::ifstream infile("../triangle/circle.1.ele");
        std::string line;
        int a, b, c, d, line_nb = 0;
        std::getline(infile, line);
        std::istringstream iss(line);
        iss >> a >> b >> c;
        int nb_triangles = a;
        matrix<double> triangles(nb_triangles, 4);

        while (std::getline(infile, line)) {
            std::istringstream iss(line);
            if (iss >> a >> b >> c >> d) {
                triangles(line_nb, 0) = b-1;
                triangles(line_nb, 1) = c-1;
                triangles(line_nb, 2) = d-1;
                triangles(line_nb, 3) = 0.5*fabs(vertices(b-1, 0)*(vertices(c-1, 1) - vertices(d-1, 1)) +
                    vertices(c-1, 0)*(vertices(d-1, 1) - vertices(b-1, 1)) +
                    vertices(d-1, 0)*(vertices(b-1, 1) - vertices(c-1, 1)));
            }
            else {
                // std::cout << "ERROR: in .1.ele file" << std::endl;
                break;
            }
            line_nb += 1;
        }

        return triangles;
    }

    matrix<int> read_boundaries(matrix<double> &vertices) {
        /* index # is triangle id, contains
        the 3 vertix id's of the corner nodes  and
        the area of the triangle, calculated using
        the vertices matrix */
        std::ifstream infile("../triangle/circle.1.poly");
        std::string line;
        int a, b, c, d, count = 0;
        std::getline(infile, line);
        std::getline(infile, line);
        std::istringstream iss(line);
        iss >> a >> b;
        int nb_boundaries = a;
        matrix<int> boundaries(nb_boundaries, 2);

        while (std::getline(infile, line)) {
            std::istringstream iss(line);
            if (iss >> a >> b >> c >> d) {
                if (vertices(b-1, 0) != 0 || vertices(c-1, 0) != 0) {
                  boundaries(count, 0) = b-1;
                  boundaries(count, 1) = c-1;
                  count += 1;
                }
            }
            else {
                // std::cout << "ERROR: in .1.poly file" << std::endl;
                break;
            }
        }

        matrix<int> boundaries_red(count, 2);
        boundaries_red = project(boundaries, range(0, count), range(0, 2));
        return boundaries_red;
    }

}

#endif


// Copyright (c) 2016 Graphcore Ltd. All rights reserved.
// Simple test case for test log of softmax
//
#define BOOST_TEST_MODULE NonLinearityTest
#include <boost/test/unit_test.hpp>
#include <iostream>
#include <limits>
#include <poplar/Engine.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/Util.hpp>
#include <poplin/codelets.hpp>
#include <popnn/LogSoftmax.hpp>
#include <popnn/codelets.hpp>
#include <popops/codelets.hpp>
#include <poputil/TileMapping.hpp>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace popnn;
using namespace poplibs_test;
using namespace poplibs_test::util;
using namespace poplibs_support;

namespace utf = boost::unit_test;
namespace fpc = boost::test_tools::fpc;

#define TOL 0.1
#define FLOAT_ATOL 1e-20
#define HALF_ATOL 1e-7

void validateLogSoftmax(unsigned batchSize, unsigned numChannels) {
  auto device = createTestDevice(TEST_TARGET);
  auto &target = device.getTarget();
  Graph graph(target);
  popnn::addCodelets(graph);
  popops::addCodelets(graph);
  poplin::addCodelets(graph);

  auto actF = graph.addVariable(FLOAT, {batchSize, numChannels}, "actF");
  auto actH = graph.addVariable(HALF, {batchSize, numChannels}, "actH");

  // arbitrary mappings
  mapTensorLinearly(graph, actF);
  mapTensorLinearly(graph, actH);

  graph.createHostWrite("inF", actF);
  graph.createHostWrite("inH", actH);
  graph.createHostRead("outF", actF);
  graph.createHostRead("outH", actH);

  std::vector<std::pair<std::string, char *>> tmap;
  Sequence uploadProg, downloadProg;

  auto rawHActF = allocateHostMemoryForTensor(actF, "actF", graph, uploadProg,
                                              downloadProg, tmap);
  auto rawHActH = allocateHostMemoryForTensor(actH, "actH", graph, uploadProg,
                                              downloadProg, tmap);

  boost::multi_array<double, 2> hActIn(boost::extents[batchSize][numChannels]),
      hOutRef(boost::extents[batchSize][numChannels]),
      hActOutF(boost::extents[batchSize][numChannels]),
      hActOutH(boost::extents[batchSize][numChannels]);

  // Reference computation
  for (unsigned b = 0; b < batchSize; ++b) {
    double maxBatch = std::numeric_limits<double>::lowest();
    for (unsigned c = 0; c < numChannels; ++c) {
      double sample = (1.0 - 2 * (c & 1)) * (1 + b) * 0.01 * c;
      hActIn[b][c] = sample;
      maxBatch = std::max(maxBatch, sample);
    }
    // compute sum of exponent
    double sum = 0.0;
    for (unsigned c = 0; c < numChannels; ++c) {
      sum += std::exp(hActIn[b][c] - maxBatch);
    }

    for (unsigned c = 0; c < numChannels; ++c) {
      hOutRef[b][c] = hActIn[b][c] - maxBatch - std::log(sum);
    }
  }

  // To test 1D
  if (batchSize == 1) {
    actF = actF.squeeze({0});
    actH = actH.squeeze({0});
  }

  // build and run the target code: non-inplace followed by in-place
  auto prog = Sequence();
  auto outF = popnn::logSoftmax(graph, actF, prog);
  auto outH = popnn::logSoftmax(graph, actH, prog);
  popnn::logSoftmaxInPlace(graph, actF, prog);
  popnn::logSoftmaxInPlace(graph, actH, prog);

  auto rawHOutF = allocateHostMemoryForTensor(outF, "outF", graph, uploadProg,
                                              downloadProg, tmap);
  auto rawHOutH = allocateHostMemoryForTensor(outH, "outH", graph, uploadProg,
                                              downloadProg, tmap);

  copy(target, hActIn, FLOAT, rawHActF.get());
  copy(target, hActIn, HALF, rawHActH.get());

  Engine fwdEng(graph, Sequence{uploadProg, prog, downloadProg});
  attachStreams(fwdEng, tmap);
  device.bind([&](const Device &d) { fwdEng.loadAndRun(d); });

  // inplace variant
  copy(target, FLOAT, rawHActF.get(), hActOutF);
  copy(target, HALF, rawHActH.get(), hActOutH);
  BOOST_TEST(checkIsClose("actOutF", hActOutF, hOutRef, TOL, FLOAT_ATOL));
  BOOST_TEST(checkIsClose("actOutH", hActOutH, hOutRef, TOL, HALF_ATOL));

  // non-inplace variant
  copy(target, FLOAT, rawHOutF.get(), hActOutF);
  copy(target, HALF, rawHOutH.get(), hActOutH);
  BOOST_TEST(checkIsClose("actOutF", hActOutF, hOutRef, TOL, FLOAT_ATOL));
  BOOST_TEST(checkIsClose("actOutH", hActOutH, hOutRef, TOL, HALF_ATOL));
}

BOOST_AUTO_TEST_CASE(logSoftmax_1D) { validateLogSoftmax(1, 100); }

BOOST_AUTO_TEST_CASE(logSoftmax_2D) { validateLogSoftmax(4, 100); }


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

#ifndef MATRIX_HPP
#define MATRIX_HPP

#include <Eigen/Dense>
#include <vector>

namespace coati {

using Matrix64f = Eigen::Matrix<float, 64, 64>;
using float_t = float;

template <class T>
class Matrix {
   public:
    Matrix() = default;
    Matrix(std::size_t rows, std::size_t cols, T value = static_cast<T>(0))
        : rows_(rows), cols_(cols), data_(rows*cols, value) { }
    Matrix(std::size_t rows, std::size_t cols, Matrix64f& eigen_m)
        : rows_(rows), cols_(cols) {
        Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> m(
            eigen_m);
        data_.resize(rows_ * cols_);
        for(std::size_t i = 0; i < rows_ * cols_; i++) {
            data_[i] = m(i);
        }
    }
    // copy constructor
    Matrix(const Matrix &) = default;
    // move constructor
    Matrix(Matrix &&) = default;
    // assignment operator
    Matrix & operator=(const Matrix &) = default;
    // move assignment operator
    Matrix & operator=(Matrix &&) = default;
    // destructor
    ~Matrix() = default;

    T operator()(std::size_t row, std::size_t col) const {
        assert(row < rows_ && col < cols_);
        return data_[row * cols_ + col];
    }
    T& operator()(std::size_t row, std::size_t col) {
        assert(row < rows_ && col < cols_);
        return data_[row * cols_ + col];
    }
    bool operator==(const Matrix& mat) const {
        if((rows_ != mat.rows_) || (cols_ != mat.cols_)) {
            return false;
        }

        for(std::size_t i = 0; i < rows_ * cols_; i++) {
            if(data_[i] != mat.data_[i]) {
                return false;
            }
        }
        return true;
    }

    void resize(std::size_t rows, std::size_t cols,
                T value = static_cast<T>(0)) {
        rows_ = rows;
        cols_ = cols;
        data_.resize(rows * cols);
        data_.assign(data_.size(), value);
    }

    std::size_t rows() const { return rows_; }
    std::size_t cols() const { return cols_; }

   private:
    std::size_t rows_{0}, cols_{0};
    std::vector<T> data_;
};  // class matrix

////////////////////////////////////////////////////////////////////////////////

// template <class T>
template <class T>
class Tensor {
   public:
    Tensor(std::size_t dims, std::size_t rows, std::size_t cols, T value = 0.0f)
        : dims_(dims), rows_(rows), cols_(cols) {
        data_.resize(dims_ * rows_ * cols_);
        data_.assign(data_.size(), value);
    }
    // copy constructor
    Tensor(const Tensor& tens)
        : dims_(tens.dims_), rows_(tens.rows_), cols_(tens.cols_) {
        memcpy(&data_[0], &tens.data_[0], dims_ * rows_ * cols_ * sizeof(T));
    }
    // move constructor
    Tensor(Tensor&& tens) noexcept
        : dims_(tens.dims_),
          rows_(tens.rows_),
          cols_(tens.cols_),
          data_(std::move(tens.data_)) {}
    // destructor
    ~Tensor() = default;
    // assignment operator
    Tensor& operator=(const Tensor& tens) {
        if(&tens != this) {
            memcpy(&data_[0], &tens.data_[0],
                   dims_ * rows_ * cols_ * sizeof(T));
        }
        return *this;
    }
    // move assignment operator
    Tensor& operator=(Tensor&& tens) noexcept {
        dims_ = tens.dims_;
        rows_ = tens.rows_;
        cols_ = tens.cols_;
        data_ = std::move(tens.data_);
        return *this;
    }
    T operator()(std::size_t dims, std::size_t row, std::size_t col) const {
        return data_[dims * rows_ * cols_ + row * cols_ + col];
    }
    T& operator()(std::size_t dims, std::size_t row, std::size_t col) {
        return data_[dims * rows_ * cols_ + row * cols_ + col];
    }
    bool operator==(const Tensor& tens) const;

   private:
    std::size_t dims_, rows_, cols_;
    std::vector<T> data_;
};  // class Tensor

using Matrixf = Matrix<float_t>;
using Matrixi = Matrix<int>;
using Tensorf = Tensor<float_t>;

}  // namespace coati
#endif


// Boost.Geometry

// Copyright (c) 2018, Oracle and/or its affiliates.

// Contributed and/or modified by Vissarion Fysikopoulos, on behalf of Oracle

// Licensed under the Boost Software License version 1.0.
// http://www.boost.org/users/license.html

#ifndef BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_LINE_INTERPOLATE_HPP
#define BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_LINE_INTERPOLATE_HPP

#include <boost/geometry/core/assert.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/coordinate_type.hpp>
#include <boost/geometry/core/radian_access.hpp>
#include <boost/geometry/srs/spheroid.hpp>
#include <boost/geometry/strategies/line_interpolate.hpp>
#include <boost/geometry/strategies/geographic/parameters.hpp>
#include <boost/geometry/util/select_calculation_type.hpp>

namespace boost { namespace geometry
{

namespace strategy { namespace line_interpolate
{

/*!
\brief Interpolate point on a geographic segment.
\ingroup strategies
\tparam FormulaPolicy The geodesic formulas used internally.
\tparam Spheroid The spheroid model.
\tparam CalculationType \tparam_calculation

\qbk{
[heading See also]
\* [link geometry.reference.algorithms.line_interpolate.line_interpolate_4_with_strategy line_interpolate (with strategy)]
\* [link geometry.reference.srs.srs_spheroid srs::spheroid]
}
 */
template
<
    typename FormulaPolicy = strategy::andoyer,
    typename Spheroid = srs::spheroid<double>,
    typename CalculationType = void
>
class geographic
{
public:
    geographic()
        : m_spheroid()
    {}

    explicit geographic(Spheroid const& spheroid)
        : m_spheroid(spheroid)
    {}

    // point-point strategy getters
    struct distance_pp_strategy
    {
        typedef distance::geographic<FormulaPolicy, Spheroid, CalculationType> type;
    };

    inline typename distance_pp_strategy::type get_distance_pp_strategy() const
    {
        typedef typename distance_pp_strategy::type distance_type;
        return distance_type(m_spheroid);
    }

    template <typename Point, typename Fraction, typename Distance>
    inline void apply(Point const& p0,
                      Point const& p1,
                      Fraction const& fraction, //fraction of segment
                      Point & p,
                      Distance const& distance) const
    {
        typedef typename select_calculation_type_alt
            <
                CalculationType,
                Point
            >::type calc_t;

        typedef typename FormulaPolicy::template inverse
                <calc_t, false, true, false, false, false> inverse_t;

        calc_t azimuth = inverse_t::apply(get_as_radian<0>(p0), get_as_radian<1>(p0),
                                          get_as_radian<0>(p1), get_as_radian<1>(p1),
                                          m_spheroid).azimuth;

        typedef typename FormulaPolicy::template direct
                <calc_t, true, false, false, false> direct_t;

        typename direct_t::result_type
        dir_r = direct_t::apply(get_as_radian<0>(p0), get_as_radian<1>(p0),
                                distance * fraction, azimuth,
                                m_spheroid);

        set_from_radian<0>(p, dir_r.lon2);
        set_from_radian<1>(p, dir_r.lat2);
    }

private:
    Spheroid m_spheroid;
};

#ifndef DOXYGEN_NO_STRATEGY_SPECIALIZATIONS
namespace services
{

template <>
struct default_strategy<geographic_tag>
{
    typedef strategy::line_interpolate::geographic<> type;
};

} // namespace services
#endif // DOXYGEN_NO_STRATEGY_SPECIALIZATIONS

}} // namespace strategy::line_interpolate

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_LINE_INTERPOLATE_HPP


/*
 * BlueNoise.hpp
 *
 *  Created on: Feb 6, 2012
 *      Author: david
 */

#ifndef BLUENOISE_HPP_
#define BLUENOISE_HPP_
//----------------------------------------------------------------------------//
#include <Slimage/Slimage.hpp>
#include <Danvil/Tools/FunctionCache.h>
#include <Eigen/Dense>
#include <vector>
#include <algorithm>
#include <cmath>
//----------------------------------------------------------------------------//
namespace pds {
//----------------------------------------------------------------------------//

namespace fattal
{
	// need to change some other functions too!!!
	constexpr unsigned int D = 2;

	struct Point {
		float x, y;
		float weight;
		float scale;
	};

	constexpr float KernelRange = 2.5f;

	constexpr float cMaxRefinementScale = 10.0f;

	constexpr float cPi = 3.141592654f;

	/** phi(x) = exp(-pi*x*x) */
	inline
	float KernelFunctorImpl(float d) {
		return std::exp(-cPi*d*d);
	}

	/**
	 * Warning: only defined for y <= 1!
	 */
	inline float KernelFunctorInverse(float y) {
		return std::sqrt(- std::log(y) / cPi);
	}

	inline
	float KernelFunctor(float d) {
		static Danvil::FunctionCache<float,1> cache(0.0f, KernelRange, &KernelFunctorImpl);
		return cache(std::abs(d));
	}

	inline
	float KernelFunctorSquareImpl(float d) {
		return std::exp(-cPi*d);
	}

	inline
	float KernelFunctorSquare(float d) {
		static Danvil::FunctionCache<float,1> cache(0.0f, KernelRange*KernelRange, &KernelFunctorSquareImpl);
		return cache(d);
	}

	inline
	float ZeroBorderAccess(const Eigen::MatrixXf& density, int x, int y) {
		if(0 <= x && x < int(density.rows()) && 0 <= y && y < int(density.cols())) {
			return density(x, y);
		}
		else {
			return 0.0f;
		}
	}

	inline
	float ZeroBorderAccess(const Eigen::MatrixXf& density, float x, float y) {
		return ZeroBorderAccess(density, (int)std::round(x), (int)std::round(y));
	}

	inline
	float KernelScaleFunction(float roh, float weight) {
//		return std::pow(roh / weight, -1.0f / float(D));
		return 1.0f / std::sqrt(roh / weight);
	}

	inline
	float ScalePowerD(float s) {
		//return std::pow(s, -float(D));
		return 1.0f / (s*s);
	}

	float EnergyApproximation(const std::vector<Point>& pnts, float x, float y);

	float Energy(const std::vector<Point>& pnts, const Eigen::MatrixXf& density);

	float EnergyDerivative(const std::vector<Point>& pnts, const Eigen::MatrixXf& density, unsigned int i, float& result_dE_x, float& result_dE_y);

	std::vector<Point> PlacePoints(const Eigen::MatrixXf& density, unsigned int p);

	void Refine(std::vector<Point>& points, const Eigen::MatrixXf& density, unsigned int iterations);

	std::vector<Point> Split(const std::vector<Point>& points, const Eigen::MatrixXf& density, bool& result_added);

	std::vector<Point> Compute(const Eigen::MatrixXf& density);

	struct Color {
		unsigned char r,g,b;
	};

	void PlotPoints(const std::vector<Point>& points, const slimage::Image1ub& img, unsigned char grey=0, bool plot_1px=true);

	void PlotPoints(const std::vector<Point>& points, const slimage::Image3ub& img, const slimage::Pixel3ub& color=slimage::Pixel3ub{{0,0,0}}, bool plot_1px=true);

}

//----------------------------------------------------------------------------//
}
//----------------------------------------------------------------------------//
#endif


/**
 * @file tests/gan_test.cpp
 * @author Kris Singh
 * @author Shikhar Jaiswal
 *
 * Tests the GAN network.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#include <mlpack/core.hpp>

#include <mlpack/methods/ann/init_rules/gaussian_init.hpp>
#include <mlpack/methods/ann/loss_functions/sigmoid_cross_entropy_error.hpp>
#include <mlpack/methods/ann/gan/gan.hpp>
#include <mlpack/methods/ann/layer/layer.hpp>
#include <mlpack/methods/softmax_regression/softmax_regression.hpp>

#include <ensmallen.hpp>

#include <boost/test/unit_test.hpp>
#include "test_tools.hpp"
#include "serialization.hpp"

using namespace mlpack;
using namespace mlpack::ann;
using namespace mlpack::math;
using namespace mlpack::regression;
using namespace std::placeholders;

BOOST_AUTO_TEST_SUITE(GANNetworkTest);

/*
 * Load pre trained network values
 * for generating distribution that
 * is close to N(4, 0.5)
 */
BOOST_AUTO_TEST_CASE(GANTest)
{
  size_t generatorHiddenLayerSize = 8;
  size_t discriminatorHiddenLayerSize = 8;
  size_t generatorOutputSize = 1;
  size_t discriminatorOutputSize = 1;
  size_t discriminatorPreTrain = 0;
  size_t batchSize = 8;
  size_t noiseDim = 1;
  size_t generatorUpdateStep = 1;
  size_t numSamples = 10000;
  double multiplier = 1;

  arma::mat trainData(1, 10000);
  trainData.imbue( [&]() { return arma::as_scalar(RandNormal(4, 0.5));});
  trainData = arma::sort(trainData);

  // Create the Discriminator network.
  FFN<SigmoidCrossEntropyError<> > discriminator;
  discriminator.Add<Linear<> > (
      generatorOutputSize, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorOutputSize);

  // Create the Generator network.
  FFN<SigmoidCrossEntropyError<> > generator;
  generator.Add<Linear<> >(noiseDim, generatorHiddenLayerSize);
  generator.Add<SoftPlusLayer<> >();
  generator.Add<Linear<> >(generatorHiddenLayerSize, generatorOutputSize);

  // Create GAN.
  GaussianInitialization gaussian(0, 0.1);
  std::function<double ()> noiseFunction = [](){ return math::Random(-8, 8) +
      math::RandNormal(0, 1) * 0.01;};
  GAN<FFN<SigmoidCrossEntropyError<> >,
      GaussianInitialization,
      std::function<double()> >
  gan(generator, discriminator, gaussian, noiseFunction, noiseDim, batchSize,
      generatorUpdateStep, discriminatorPreTrain, multiplier);
  gan.ResetData(trainData);

  Log::Info << "Loading Parameters" << std::endl;
  arma::mat parameters, generatorParameters;
  parameters.load("preTrainedGAN.arm");
  gan.Parameters() = parameters;

  // Generate samples.
  Log::Info << "Sampling..." << std::endl;
  arma::mat noise(noiseDim, batchSize);

  size_t dim = std::sqrt(trainData.n_rows);
  arma::mat generatedData(2 * dim, dim * numSamples);

  for (size_t i = 0; i < numSamples; i++)
  {
    arma::mat samples;
    noise.imbue( [&]() { return noiseFunction(); } );

    gan.Generator().Forward(noise, samples);
    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(0, i * dim, dim - 1, i * dim + dim - 1) = samples;

    samples = trainData.col(math::RandInt(0, trainData.n_cols));
    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(dim,
        i * dim, 2 * dim - 1, i * dim + dim - 1) = samples;
  }

  double generatedMean = arma::as_scalar(arma::mean(
      generatedData.rows(0, dim - 1), 1));
  double originalMean = arma::as_scalar(arma::mean(
      generatedData.rows(dim, 2 * dim - 1), 1));
  double generatedStd = arma::as_scalar(arma::stddev(
      generatedData.rows(0, dim - 1), 0, 1));
  double originalStd = arma::as_scalar(arma::stddev(
      generatedData.rows(dim, 2 * dim - 1), 0, 1));

  BOOST_REQUIRE_LE(generatedMean - originalMean, 0.2);
  BOOST_REQUIRE_LE(generatedStd - originalStd, 0.2);
}

/*
 * Tests the GAN implementation of the O'Reilly Test on the MNIST dataset.
 * It's not viable to train on bigger parameters due to time constraints.
 * Please refer mlpack/models repository for the tutorial.
 */
BOOST_AUTO_TEST_CASE(GANMNISTTest)
{
  size_t dNumKernels = 32;
  size_t discriminatorPreTrain = 5;
  size_t batchSize = 5;
  size_t noiseDim = 100;
  size_t generatorUpdateStep = 1;
  size_t numSamples = 10;
  double stepSize = 0.0003;
  double eps = 1e-8;
  size_t numEpoches = 1;
  double tolerance = 1e-5;
  int datasetMaxCols = 10;
  bool shuffle = true;
  double multiplier = 10;

  Log::Info << std::boolalpha
      << " batchSize = " << batchSize << std::endl
      << " generatorUpdateStep = " << generatorUpdateStep << std::endl
      << " noiseDim = " << noiseDim << std::endl
      << " numSamples = " << numSamples << std::endl
      << " stepSize = " << stepSize << std::endl
      << " numEpoches = " << numEpoches << std::endl
      << " tolerance = " << tolerance << std::endl
      << " shuffle = " << shuffle << std::endl;

  arma::mat trainData;
  trainData.load("mnist_first250_training_4s_and_9s.arm");
  Log::Info << arma::size(trainData) << std::endl;

  trainData = trainData.cols(0, datasetMaxCols - 1);

  size_t numIterations = trainData.n_cols * numEpoches;
  numIterations /= batchSize;

  Log::Info << "Dataset loaded (" << trainData.n_rows << ", "
            << trainData.n_cols << ")" << std::endl;
  Log::Info << trainData.n_rows << "--------" << trainData.n_cols << std::endl;

  // Create the Discriminator network.
  FFN<SigmoidCrossEntropyError<> > discriminator;
  discriminator.Add<Convolution<> >(1, dNumKernels, 5, 5, 1, 1, 2, 2, 28, 28);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<MeanPooling<> >(2, 2, 2, 2);
  discriminator.Add<Convolution<> >(dNumKernels, 2 * dNumKernels, 5, 5, 1, 1,
      2, 2, 14, 14);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<MeanPooling<> >(2, 2, 2, 2);
  discriminator.Add<Linear<> >(7 * 7 * 2 * dNumKernels, 1024);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> >(1024, 1);

  // Create the Generator network.
  FFN<SigmoidCrossEntropyError<> > generator;
  generator.Add<Linear<> >(noiseDim, 3136);
  generator.Add<BatchNorm<> >(3136);
  generator.Add<ReLULayer<> >();
  generator.Add<Convolution<> >(1, noiseDim / 2, 3, 3, 2, 2, 1, 1, 56, 56);
  generator.Add<BatchNorm<> >(39200);
  generator.Add<ReLULayer<> >();
  generator.Add<BilinearInterpolation<> >(28, 28, 56, 56, noiseDim / 2);
  generator.Add<Convolution<> >(noiseDim / 2, noiseDim / 4, 3, 3, 2, 2, 1, 1,
      56, 56);
  generator.Add<BatchNorm<> >(19600);
  generator.Add<ReLULayer<> >();
  generator.Add<BilinearInterpolation<> >(28, 28, 56, 56, noiseDim / 4);
  generator.Add<Convolution<> >(noiseDim / 4, 1, 3, 3, 2, 2, 1, 1, 56, 56);
  generator.Add<TanHLayer<> >();

  // Create GAN.
  GaussianInitialization gaussian(0, 1);
  ens::Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
      tolerance, shuffle);
  std::function<double()> noiseFunction = [] () {
      return math::RandNormal(0, 1);};
  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
      std::function<double()> > gan(generator, discriminator,
      gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
      discriminatorPreTrain, multiplier);

  Log::Info << "Training..." << std::endl;
  std::stringstream stream;
  double objVal = gan.Train(trainData, optimizer, ens::ProgressBar(70, stream));
  BOOST_REQUIRE_GT(stream.str().length(), 0);
  BOOST_REQUIRE_EQUAL(std::isfinite(objVal), true);

  // Generate samples.
  Log::Info << "Sampling..." << std::endl;
  arma::mat noise(noiseDim, batchSize);
  size_t dim = std::sqrt(trainData.n_rows);
  arma::mat generatedData(2 * dim, dim * numSamples);

  for (size_t i = 0; i < numSamples; i++)
  {
    arma::mat samples;
    noise.imbue( [&]() { return noiseFunction(); } );

    gan.Generator().Forward(noise, samples);
    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(0, i * dim, dim - 1, i * dim + dim - 1) = samples;

    samples = trainData.col(math::RandInt(0, trainData.n_cols));
    samples.reshape(dim, dim);
    samples = samples.t();

    generatedData.submat(dim,
        i * dim, 2 * dim - 1, i * dim + dim - 1) = samples;
  }

  Log::Info << "Output generated!" << std::endl;

  // Check that Serialization is working correctly.
  arma::mat orgPredictions;
  gan.Predict(noise, orgPredictions);

  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
      std::function<double()> > ganText(generator, discriminator,
      gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
      discriminatorPreTrain, multiplier);

  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
      std::function<double()> > ganXml(generator, discriminator,
      gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
      discriminatorPreTrain, multiplier);

  GAN<FFN<SigmoidCrossEntropyError<> >, GaussianInitialization,
      std::function<double()> > ganBinary(generator, discriminator,
      gaussian, noiseFunction, noiseDim, batchSize, generatorUpdateStep,
      discriminatorPreTrain, multiplier);

  SerializeObjectAll(gan, ganXml, ganText, ganBinary);

  arma::mat predictions, xmlPredictions, textPredictions, binaryPredictions;
  gan.Predict(noise, predictions);
  ganXml.Predict(noise, xmlPredictions);
  ganText.Predict(noise, textPredictions);
  ganBinary.Predict(noise, binaryPredictions);

  CheckMatrices(orgPredictions, predictions);
  CheckMatrices(orgPredictions, xmlPredictions);
  CheckMatrices(orgPredictions, textPredictions);
  CheckMatrices(orgPredictions, binaryPredictions);
}

/*
 * Create GAN network and test for memory sharing
 * between discriminator and gan predictors.
 */
BOOST_AUTO_TEST_CASE(GANMemorySharingTest)
{
  size_t generatorHiddenLayerSize = 8;
  size_t discriminatorHiddenLayerSize = 8;
  size_t generatorOutputSize = 1;
  size_t discriminatorOutputSize = 1;
  size_t discriminatorPreTrain = 0;
  size_t batchSize = 8;
  size_t noiseDim = 1;
  size_t generatorUpdateStep = 1;
  double multiplier = 1;
  double eps = 1e-8;
  double stepSize = 0.0003;
  size_t numIterations = 8;
  double tolerance = 1e-5;
  bool shuffle = true;

  arma::mat trainData(1, 10000);
  trainData.imbue( [&]() { return arma::as_scalar(RandNormal(4, 0.5));});
  trainData = arma::sort(trainData);

  // Create the Discriminator network.
  FFN<SigmoidCrossEntropyError<> > discriminator;
  discriminator.Add<Linear<> > (
      generatorOutputSize, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorHiddenLayerSize * 2);
  discriminator.Add<ReLULayer<> >();
  discriminator.Add<Linear<> > (
      discriminatorHiddenLayerSize * 2, discriminatorOutputSize);

  // Create the Generator network.
  FFN<SigmoidCrossEntropyError<> > generator;
  generator.Add<Linear<> >(noiseDim, generatorHiddenLayerSize);
  generator.Add<SoftPlusLayer<> >();
  generator.Add<Linear<> >(generatorHiddenLayerSize, generatorOutputSize);

  // Create GAN.
  GaussianInitialization gaussian(0, 0.1);
  ens::Adam optimizer(stepSize, batchSize, 0.9, 0.999, eps, numIterations,
      tolerance, shuffle);
  std::function<double ()> noiseFunction = [](){ return math::Random(-8, 8) +
      math::RandNormal(0, 1) * 0.01;};
  GAN<FFN<SigmoidCrossEntropyError<> >,
      GaussianInitialization,
      std::function<double()> >
  gan(generator, discriminator, gaussian, noiseFunction,
      noiseDim, batchSize, generatorUpdateStep, discriminatorPreTrain,
      multiplier);

  gan.Train(trainData, optimizer);

  CheckMatrices(gan.Predictors().head_cols(trainData.n_cols), trainData);
  CheckMatrices(gan.Predictors(), gan.Discriminator().Predictors());
  gan.Shuffle();
  CheckMatrices(gan.Predictors(), gan.Discriminator().Predictors());
  CheckMatricesNotEqual(gan.Predictors().head_cols(trainData.n_cols),
      trainData);
}

BOOST_AUTO_TEST_SUITE_END();


#include <mex.h>
#include <math.h>
#include <iostream>
#include <vector>

#include <igl/matlab/MexStream.h>
#include <igl/matlab/parse_rhs.h>
#include <igl/matlab/prepare_lhs.h>
#include <igl/matlab/validate_arg.h>

#include <igl/PI.h>
#include <igl/forward_kinematics.h>
#include <igl/directed_edge_parents.h>

#include <Eigen/Core>

using namespace Eigen;
using namespace std;

void euler_to_quat(const Eigen::Vector3d& euler, Eigen::Quaterniond& q) {

  q = Eigen::AngleAxisd(euler(2), Eigen::Vector3d::UnitZ())
      * Eigen::AngleAxisd(euler(1), Eigen::Vector3d::UnitY())
      * Eigen::AngleAxisd(euler(0), Eigen::Vector3d::UnitX());
}

void euler_to_quat(const Eigen::Vector3d& euler, const Eigen::Affine3d& a, Eigen::Quaterniond& q) {
  q = Eigen::AngleAxisd(euler(2), a.rotation().col(2))
      * Eigen::AngleAxisd(euler(1), a.rotation().col(1))
      * Eigen::AngleAxisd(euler(0), a.rotation().col(0));
}

void read_bone_anim(std::string anim_file,
    const Eigen::MatrixXd& C,
    const Eigen::MatrixXi& BE,
    const Eigen::VectorXi& P,
    std::vector<Eigen::MatrixXd>& T_list)
{
  FILE* file;
  file= fopen(anim_file.c_str(), "rb");

  double degree2radian = igl::PI/180.0;

  int num_bone, num_frame;
  double val = 0;

  typedef std::vector<Eigen::Quaterniond,
      Eigen::aligned_allocator<Eigen::Quaterniond> > RotationList;

  fscanf(file, "%d %d\n", &num_bone, &num_frame);
  T_list.resize(num_frame);
  RotationList rot_list(num_bone);
  std::vector<Eigen::Vector3d> tran_list(num_bone);
  std::vector<Eigen::Affine3d> rest_list(num_bone);

  Eigen::Vector3d root_rest_tran;
  Eigen::Affine3d root_affine;
  for (int j = 0; j < 3; j++) {
    fscanf(file, "%lf", &val);
    root_rest_tran(j) = val;
  }
  Eigen::Vector3d euler;
  for (int j = 0; j < 3; j++) {
    fscanf(file, "%lf", &val);
    euler(j) = val;
  }
  euler = euler.array() * degree2radian;
  Eigen::Quaterniond q;
  euler_to_quat(euler, q);
  root_affine = Eigen::Affine3d::Identity();
  root_affine.rotate(q);

  for (int i = 0; i < num_bone; i++) {
    Eigen::Vector3d euler;
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      euler(j) = val;
    }
    euler = euler.array() * degree2radian;
    Eigen::Quaterniond q;
    euler_to_quat(euler, q);
    Eigen::Affine3d a = Eigen::Affine3d::Identity();
    a.rotate(q);
    rest_list[i] = a;
  }
  //Eigen::Affine3d root_affine = rest_list[0];
  Eigen::Vector3d root_tran, root_rot;
  for (int k = 0; k < num_frame; k++) {
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      root_rot(j) = val;
    }
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      root_tran(j) = val;
    }
    root_rot = root_rot.array() * degree2radian;
    Eigen::Quaterniond root_q;
    euler_to_quat(root_rot, root_affine, root_q);

    for (int i = 0; i < num_bone; i++) {
      Eigen::Vector3d euler;
      for (int j = 0; j < 3; j++) {
        fscanf(file, "%lf", &val);
        euler(j) = val;
      }
      euler = euler.array() * degree2radian;
      Eigen::Quaterniond q;
      euler_to_quat(euler, rest_list[i], q);

      if(P(i)==-1){
        int root_cnt = 0;
        for(int ii=0; ii<num_bone; ii++)
          if(P(ii)==-1)
            root_cnt++;
        if(root_cnt>1)
          q = q*root_q;
        tran_list[i] = root_tran - root_rest_tran;
      }
      else
        tran_list[i] = Eigen::Vector3d::Zero();
      rot_list[i] = q;
    }
    RotationList vQ;
    std::vector<Eigen::Vector3d> vT;
    igl::forward_kinematics(C, BE, P, rot_list, tran_list, vQ, vT);

    Eigen::MatrixXd T(num_bone * 4, 3);
    for (int i = 0; i < num_bone; i++) {
      Eigen::Affine3d a = Eigen::Affine3d::Identity();
      a.translate(vT[i]);
      a.rotate(vQ[i]);
      T.block(i * 4, 0, 4, 3) = a.matrix().transpose().block(0, 0, 4, 3);
    }
    T_list[k] = T;
  }
  fclose(file);
}

void read_bone_anim(std::string anim_file,
    const Eigen::MatrixXd& C,
    const Eigen::MatrixXi& BE,
    const Eigen::VectorXi& P,
    const Eigen::VectorXd& center,
    const double& scale,
    std::vector<Eigen::MatrixXd>& T_list)
{
  FILE* file;
  file= fopen(anim_file.c_str(), "rb");

  double degree2radian = igl::PI/180.0;

  int num_bone, num_frame;
  double val = 0;

  typedef std::vector<Eigen::Quaterniond,
      Eigen::aligned_allocator<Eigen::Quaterniond> > RotationList;

  fscanf(file, "%d %d\n", &num_bone, &num_frame);
  T_list.resize(num_frame);
  RotationList rot_list(num_bone);
  std::vector<Eigen::Vector3d> tran_list(num_bone);
  std::vector<Eigen::Affine3d> rest_list(num_bone);

  Eigen::Vector3d root_rest_tran;
  Eigen::Affine3d root_affine;
  for (int j = 0; j < 3; j++) {
    fscanf(file, "%lf", &val);
    root_rest_tran(j) = val;
  }
  root_rest_tran = (root_rest_tran-center)/scale;

  Eigen::Vector3d euler;
  for (int j = 0; j < 3; j++) {
    fscanf(file, "%lf", &val);
    euler(j) = val;
  }
  euler = euler.array() * degree2radian;
  Eigen::Quaterniond q;
  euler_to_quat(euler, q);
  root_affine = Eigen::Affine3d::Identity();
  root_affine.rotate(q);

  for (int i = 0; i < num_bone; i++) {
    Eigen::Vector3d euler;
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      euler(j) = val;
    }
    euler = euler.array() * degree2radian;
    Eigen::Quaterniond q;
    euler_to_quat(euler, q);
    Eigen::Affine3d a = Eigen::Affine3d::Identity();
    a.rotate(q);
    rest_list[i] = a;
  }
  //Eigen::Affine3d root_affine = rest_list[0];
  Eigen::Vector3d root_tran, root_rot;
  for (int k = 0; k < num_frame; k++) {
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      root_rot(j) = val;
    }
    for (int j = 0; j < 3; j++) {
      fscanf(file, "%lf", &val);
      root_tran(j) = val;
    }
    root_tran = (root_tran - center)/scale;
    root_rot = root_rot.array() * degree2radian;
    Eigen::Quaterniond root_q;
    euler_to_quat(root_rot, root_affine, root_q);

    for (int i = 0; i < num_bone; i++) {
      Eigen::Vector3d euler;
      for (int j = 0; j < 3; j++) {
        fscanf(file, "%lf", &val);
        euler(j) = val;
      }
      euler = euler.array() * degree2radian;
      Eigen::Quaterniond q;
      euler_to_quat(euler, rest_list[i], q);

      if(P(i)==-1){
        int root_cnt = 0;
        for(int ii=0; ii<num_bone; ii++)
          if(P(ii)==-1)
            root_cnt++;
        if(root_cnt>1)
          q = q*root_q;
        tran_list[i] = root_tran - root_rest_tran;
      }
      else
        tran_list[i] = Eigen::Vector3d::Zero();
      rot_list[i] = q;
    }
    RotationList vQ;
    std::vector<Eigen::Vector3d> vT;
    igl::forward_kinematics(C, BE, P, rot_list, tran_list, vQ, vT);

    Eigen::MatrixXd T(num_bone * 4, 3);
    for (int i = 0; i < num_bone; i++) {
      Eigen::Affine3d a = Eigen::Affine3d::Identity();
      a.translate(vT[i]);
      a.rotate(vQ[i]);
      T.block(i * 4, 0, 4, 3) = a.matrix().transpose().block(0, 0, 4, 3);
    }
    T_list[k] = T;

  }
  fclose(file);
}

void mexFunction(
  int          nlhs,
  mxArray      *plhs[],
  int          nrhs,
  const mxArray *prhs[])
{
    Eigen::MatrixXd C;
    Eigen::MatrixXi BE;
    Eigen::VectorXi P;
    Eigen::VectorXd center;
    double scale;
    char* anim_file = mxArrayToString(prhs[0]);

    igl::matlab::parse_rhs_double(prhs+1, C);
    igl::matlab::parse_rhs_index(prhs+2, BE);

    igl::matlab::parse_rhs_double(prhs+3, center);
    scale = (double) *mxGetPr(prhs[4]);

    igl::directed_edge_parents(BE, P);

    std::vector<Eigen::MatrixXd> T_list;
    read_bone_anim(anim_file, C,BE,P,center,scale, T_list);
    //read_bone_anim(anim_file, C,BE,P, T_list);

    plhs[0] = mxCreateCellMatrix(T_list.size(),1);

    mxArray *x;
    for(int i=0; i<T_list.size(); i++){
        const int m = T_list[i].rows();
        const int n = T_list[i].cols();
        x = mxCreateDoubleMatrix(m,n, mxREAL);
        Eigen::Map< Eigen::Matrix<double,Eigen::Dynamic,Eigen::Dynamic> > map(mxGetPr(x),m,n);
        map = T_list[i].template cast<double>();
        mxSetCell(plhs[0], i, x);
    }
    return;
}


// Boost.Geometry (aka GGL, Generic Geometry Library)
// This file is manually converted from PROJ4

// Copyright (c) 2008-2012 Barend Gehrels, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to 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)

// This file is converted from PROJ4, http://trac.osgeo.org/proj
// PROJ4 is originally written by Gerald Evenden (then of the USGS)
// PROJ4 is maintained by Frank Warmerdam
// PROJ4 is converted to Geometry Library by Barend Gehrels (Geodan, Amsterdam)

// Original copyright notice:

// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

#ifndef BOOST_GEOMETRY_PROJECTIONS_PJ_MLFN_HPP
#define BOOST_GEOMETRY_PROJECTIONS_PJ_MLFN_HPP

#include <boost/geometry/util/math.hpp>

namespace boost { namespace geometry { namespace projections {

namespace detail {

/* meridinal distance for ellipsoid and inverse
**    8th degree - accurate to < 1e-5 meters when used in conjuction
**        with typical major axis values.
**    Inverse determines phi to EPS (1e-11) radians, about 1e-6 seconds.
*/
static const double C00 = 1.;
static const double C02 = .25;
static const double C04 = .046875;
static const double C06 = .01953125;
static const double C08 = .01068115234375;
static const double C22 = .75;
static const double C44 = .46875;
static const double C46 = .01302083333333333333;
static const double C48 = .00712076822916666666;
static const double C66 = .36458333333333333333;
static const double C68 = .00569661458333333333;
static const double C88 = .3076171875;
static const double EPS = 1e-11;
static const int MAX_ITER = 10;
static const int EN_SIZE = 5;

inline bool pj_enfn(double es, double* en)
{
    double t; //, *en;

    //if (en = (double *)pj_malloc(EN_SIZE * sizeof(double)))
    {
        en[0] = C00 - es * (C02 + es * (C04 + es * (C06 + es * C08)));
        en[1] = es * (C22 - es * (C04 + es * (C06 + es * C08)));
        en[2] = (t = es * es) * (C44 - es * (C46 + es * C48));
        en[3] = (t *= es) * (C66 - es * C68);
        en[4] = t * es * C88;
    }
    // return en;
    return true;
}

inline double pj_mlfn(double phi, double sphi, double cphi, const double *en)
{
    cphi *= sphi;
    sphi *= sphi;
    return(en[0] * phi - cphi * (en[1] + sphi*(en[2]
        + sphi*(en[3] + sphi*en[4]))));
}

inline double pj_inv_mlfn(double arg, double es, const double *en)
{
    double s, t, phi, k = 1./(1.-es);
    int i;

    phi = arg;
    for (i = MAX_ITER; i ; --i) { /* rarely goes over 2 iterations */
        s = sin(phi);
        t = 1. - es * s * s;
        phi -= t = (pj_mlfn(phi, s, cos(phi), en) - arg) * (t * sqrt(t)) * k;
        if (geometry::math::abs(t) < EPS)
            return phi;
    }
    throw proj_exception(-17);
    return phi;
}

} // namespace detail
}}} // namespace boost::geometry::projections

#endif