/*
 * wave_equation_reference2.cpp
 *
 *  Created on: 27.05.2019
 *      Author: thies
 */

#include <base/AdaptiveMesh.h>
#include <base/ConstantMesh.h>
#include <base/DiscretizedFunction.h>
#include <base/SpaceTimeMesh.h>
#include <base/Util.h>
#include <deal.II/base/exceptions.h>
#include <deal.II/base/function.h>
#include <deal.II/base/logstream.h>
#include <deal.II/base/numbers.h>
#include <deal.II/base/point.h>
#include <deal.II/base/quadrature.h>
#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/timer.h>
#include <deal.II/fe/fe_q.h>
#include <deal.II/grid/grid_generator.h>
#include <deal.II/grid/grid_tools.h>
#include <deal.II/grid/tria.h>
#include <forward/L2RightHandSide.h>
#include <forward/WaveEquation.h>
#include <gtest/gtest.h>
#include <norms/L2H1.h>
#include <norms/L2L2.h>
#include <stddef.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <vector>

namespace {

using namespace dealii;
using namespace wavepi::forward;
using namespace wavepi::base;
using namespace wavepi;

/*****
NOTE: This module is used for automatic figure generation
and should therefore only be changed together with the thesis!
*****/

template <int dim>
class TestQ : public LightFunction<dim> {
 public:
  virtual ~TestQ() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const { return 0.05*t+0.05*p[0]*p[1]; }
};

template <int dim>
class TestNu : public LightFunction<dim> {
 public:
  virtual ~TestNu() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const { return 0.05*abs(cos(p[1]*t)); }
};

template <int dim>
class TestC : public LightFunction<dim> {
 public:
  virtual ~TestC() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const { return sqrt(1 + t) * (1 + p[0]); }
};

template <int dim>
class TestRho : public LightFunction<dim> {
 public:
  virtual ~TestRho() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const { return (1 + t) * (1 + p[1]); }
};

template <int dim>
class TestU : public Function<dim> {
 public:
  virtual ~TestU() = default;

  virtual double value(const Point<dim> &p, const unsigned int component __attribute__((unused))) const {
    double tmp = 1;

    for (size_t i = 0; i < dim; i++)
      tmp *= sin(p[i]);

    return cos(this->get_time()) * tmp;
  }

  virtual double evaluate(const Point<dim> &p, double t) const {
    double tmp = 1;

    for (size_t i = 0; i < dim; i++)
      tmp *= sin(p[i]);

    return cos(t) * tmp;
  }

  virtual Tensor<1, dim, double> gradient(const Point<dim> &p,
                                          const unsigned int component __attribute__((unused))) const {
    Tensor<1, dim, double> res;

    for (size_t j = 0; j < dim; j++) {
      double tmp = 1;

      for (size_t i = 0; i < dim; i++)
        if (i != j) tmp *= sin(p[i]);

      tmp *= cos(p[j]);
      res[j] = cos(this->get_time()) * tmp;
    }

    return res;
  }
};

template <int dim>
class TestV : public LightFunction<dim> {
 public:
  virtual ~TestV() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const {
    double tmp = 1;

    for (size_t i = 0; i < dim; i++)
      tmp *= sin(p[i]);

    return -sin(t) * tmp;
  }
};

template <int dim>
class TestF : public LightFunction<dim> {
 public:
  virtual ~TestF() = default;

  virtual double evaluate(const Point<dim> &p, const double t) const {
    double tmp2 = 1;
    for (size_t i = 0; i < dim; i++)
      tmp2 *= sin(p[i]);

    // 1/rho (u'/c^2)'
    double tmp =
        tmp2 * (sin(t) - (1 + t) * cos(t)) / ((1 + t) * (1 + t) * (1 + t) * (1 + p[0]) * (1 + p[0]) * (1 + p[1]));

    // div (nabla u / rho)

    double fac = cos(t) / ((1 + t) * (1 + p[1]));
    tmp -= -dim * fac * tmp2;

    double tmp3 = cos(p[1]);
    for (size_t i = 0; i < dim; i++)
      if (i != 1) tmp3 *= sin(p[i]);

    tmp -= -1 * fac * tmp3 / (1 + p[1]);

    // q and nu
    tmp+= q.evaluate(p, t) * u.evaluate(p,t) + nu.evaluate(p,t)*v.evaluate(p,t);

    return tmp;
  }

  TestQ<dim> q;
  TestNu<dim> nu;
  TestU<dim> u;
  TestV<dim> v;
};

template <int dim>
void run_reference_test2(std::shared_ptr<SpaceTimeMesh<dim>> mesh, int refines, bool expect = true, bool save = false,
                         std::shared_ptr<std::ofstream> log = nullptr, int precondition = -1) {
  deallog << std::endl << "----------  n_dofs(0): " << mesh->get_dof_handler(0)->n_dofs();
  deallog << ", n_steps: " << mesh->get_times().size() << "  ----------" << std::endl;

  WaveEquation<dim> wave_eq(mesh);
  wave_eq.set_precondition_max_age(precondition);

  wave_eq.set_param_rho(std::make_shared<TestRho<dim>>());
  wave_eq.set_param_c(std::make_shared<TestC<dim>>());
  wave_eq.set_param_q(std::make_shared<TestQ<dim>>());
  wave_eq.set_param_nu(std::make_shared<TestNu<dim>>());

  auto u_cont = std::make_shared<TestU<dim>>();
  auto v_cont = std::make_shared<TestV<dim>>();

  wave_eq.set_initial_values_u(u_cont);
  wave_eq.set_initial_values_v(v_cont);

  Timer timer;
  timer.start();
  DiscretizedFunction<dim> solu = wave_eq.run(std::make_shared<L2RightHandSide<dim>>(std::make_shared<TestF<dim>>()));
  timer.stop();

  DiscretizedFunction<dim> solv = solu.derivative();
  solu.throw_away_derivative();

  DiscretizedFunction<dim> u_disc(mesh, *u_cont);
  DiscretizedFunction<dim> tmp_u(solu);
  tmp_u -= u_disc;

  DiscretizedFunction<dim> v_disc(mesh, *v_cont);
  DiscretizedFunction<dim> tmp_v(solv);
  tmp_v -= v_disc;

  u_disc.set_norm(std::make_shared<norms::L2L2<dim>>());
  double err_u_L2L2 = norms::L2L2<dim>::absolute_error(solu, *u_cont) / u_disc.norm();

  u_disc.set_norm(std::make_shared<norms::L2H1<dim>>());
  double err_u_L2H1 = norms::L2H1<dim>::absolute_error(solu, *u_cont) / u_disc.norm();

  v_disc.set_norm(std::make_shared<norms::L2L2<dim>>());
  double err_v_L2L2 = norms::L2L2<dim>::absolute_error(solv, *v_cont) / v_disc.norm();

  tmp_u.set_norm(std::make_shared<norms::L2L2<dim>>());
  u_disc.set_norm(std::make_shared<norms::L2L2<dim>>());
  double err_u_L2L2_disc = tmp_u.norm() / u_disc.norm();

  tmp_u.set_norm(std::make_shared<norms::L2H1<dim>>());
  u_disc.set_norm(std::make_shared<norms::L2H1<dim>>());
  double err_u_L2H1_disc = tmp_u.norm() / u_disc.norm();

  tmp_v.set_norm(std::make_shared<norms::L2L2<dim>>());
  v_disc.set_norm(std::make_shared<norms::L2L2<dim>>());
  double err_v_L2L2_disc = tmp_v.norm() / v_disc.norm();

  if (expect) {
    EXPECT_LT(err_u_L2L2, 1e-1);
    EXPECT_LT(err_v_L2L2, 1e-1);
  }

  if (save) {
    DiscretizedFunction<dim> u_disc2(mesh, *u_cont);

    DiscretizedFunction<dim> tmp(solu);
    tmp -= u_disc2;

    u_disc2.write_pvd("./", "refu", "uref");
    solu.write_pvd("./", "solu", "u");
    tmp.write_pvd("./", "diff", "udiff");
  }

  deallog << std::scientific << " L2L2 rerr of u = " << err_u_L2L2 << " (disc: " << err_u_L2L2_disc
          << "), L2L2 rerr of v = " << err_v_L2L2 << " (disc: " << err_v_L2L2_disc
          << "), L2H1 rerr of u = " << err_u_L2H1 << " (disc: " << err_u_L2H1_disc << "), cpu = " << std::fixed
          << std::setprecision(2) << timer.cpu_time() << "s" << std::endl;

  double dt = mesh->get_time(1) - mesh->get_time(0);
  double h  = dealii::GridTools::maximal_cell_diameter(*mesh->get_triangulation(0));

  if (log)
    *log << std::scientific << mesh->length() << " " << dt << " " << refines << " " << h << " " << err_u_L2L2 << " "
         << err_v_L2L2 << " " << err_u_L2H1 << std::endl;
}

template <int dim>
void run_reference_test2_constant(int fe_order, int quad_order, int refines, int steps, bool expect = true,
                                  bool save = false, std::shared_ptr<std::ofstream> log = nullptr,
                                  int precondition = -1) {
  auto triangulation = std::make_shared<Triangulation<dim>>();
  GridGenerator::hyper_cube(*triangulation, 0.0, numbers::PI);
  Util::set_all_boundary_ids(*triangulation, 0);
  triangulation->refine_global(refines);

  double t_end   = 3.0;
  double t_start = 0.0, dt = t_end / (steps - 1);
  std::vector<double> times;

  for (size_t i = 0; t_start + i * dt <= t_end; i++)
    times.push_back(t_start + i * dt);

  FE_Q<dim> fe(fe_order);
  Quadrature<dim> quad = QGauss<dim>(quad_order);  // exact in poly degree 2n-1 (needed: fe_dim^3)

  std::shared_ptr<SpaceTimeMesh<dim>> mesh = std::make_shared<ConstantMesh<dim>>(times, fe, quad, triangulation);

  return run_reference_test2<dim>(mesh, refines, expect, save, log, precondition);
}

TEST(WaveEquation, ReferenceTestParameters2DFE1) {
  auto file_time = std::make_shared<std::ofstream>("./ReferenceTestParameters2DFE1_time.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_time) << "could not open file for output";

  for (int steps = 4; steps <= 128; steps = (int)(steps * 1.41))
    run_reference_test2_constant<2>(1, 5, 9, steps, steps >= 64, false, file_time);
  file_time->close();

  auto file_space = std::make_shared<std::ofstream>("./ReferenceTestParameters2DFE1_space.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_space) << "could not open file for output";

  for (int refine = 1; refine <= 9; refine++)
    run_reference_test2_constant<2>(1, 5, refine, 256, refine >= 4, false, file_space);
  file_space->close();
}

TEST(WaveEquation, ReferenceTestParameters2DFE2) {
  auto file_space = std::make_shared<std::ofstream>("./ReferenceTestParameters2DFE2_space.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_space) << "could not open file for output";

  for (int refine = 1; refine <= 5; refine++)
    run_reference_test2_constant<2>(2, 5, refine, 1024, refine >= 4, false, file_space);
  file_space->close();
}

TEST(WaveEquation, ReferenceTestParameters3DFE1) {
  auto file_time = std::make_shared<std::ofstream>("./ReferenceTestParameters3DFE1_time.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_time) << "could not open file for output";

  for (int steps = 4; steps <= 16; steps = (int)(steps * 1.41))
    run_reference_test2_constant<3>(1, 5, 6, steps, steps >= 64, false, file_time);
  file_time->close();

  auto file_space = std::make_shared<std::ofstream>("./ReferenceTestParameters3DFE1_space.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_space) << "could not open file for output";

  for (int refine = 1; refine <= 6; refine++)
    run_reference_test2_constant<3>(1, 5, refine, 256, refine >= 4, false, file_space);
  file_space->close();
}

TEST(WaveEquation, ReferenceTestParameters3DFE2) {
  auto file_space = std::make_shared<std::ofstream>("./ReferenceTestParameters3DFE2_space.dat", std::ios_base::trunc);
  ASSERT_TRUE(*file_space) << "could not open file for output";

  for (int refine = 1; refine <= 4; refine++)
    run_reference_test2_constant<3>(2, 5, refine, 128, refine >= 4, false, file_space);
  file_space->close();
}

TEST(WaveEquation, PreconditionTest2DFE1) {
  for (int i = 0; i <= 9; i++) {
    int max_age;
    if (i == 0)
      max_age = -1;
    else if (i == 1)
      max_age = 0;
    else
      max_age = 1 << (i - 2);

    run_reference_test2_constant<2>(1, 5, 6, 128, false, false, nullptr, max_age);
    deallog << " precon = " << max_age << std::endl;
  }
}

TEST(WaveEquation, PreconditionTest2DFE2) {
  for (int i = 0; i <= 9; i++) {
    int max_age;
    if (i == 0)
      max_age = -1;
    else if (i == 1)
      max_age = 0;
    else
      max_age = 1 << (i - 2);

    run_reference_test2_constant<2>(2, 5, 6, 128, false, false, nullptr, max_age);
    deallog << " precon = " << max_age << std::endl;
  }
}
}  // namespace


//  (C) Copyright Eric Niebler 2005.
//  Use, modification and distribution are 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)

// Test case for extended_p_square_quantile.hpp

#include <iostream>
#include <boost/random.hpp>
#include <boost/test/unit_test.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/accumulators/numeric/functional/vector.hpp>
#include <boost/accumulators/numeric/functional/complex.hpp>
#include <boost/accumulators/numeric/functional/valarray.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/extended_p_square_quantile.hpp>

using namespace boost;
using namespace unit_test;
using namespace boost::accumulators;

///////////////////////////////////////////////////////////////////////////////
// test_stat
//
void test_stat()
{
    typedef accumulator_set<double, stats<tag::extended_p_square_quantile> > accumulator_t;
    typedef accumulator_set<double, stats<tag::weighted_extended_p_square_quantile>, double > accumulator_t_weighted;
    typedef accumulator_set<double, stats<tag::extended_p_square_quantile(quadratic)> > accumulator_t_quadratic;
    typedef accumulator_set<double, stats<tag::weighted_extended_p_square_quantile(quadratic)>, double > accumulator_t_weighted_quadratic;

    // tolerance
    double epsilon = 1;

    // a random number generator
    boost::lagged_fibonacci607 rng;

    std::vector<double> probs;

    probs.push_back(0.990);
    probs.push_back(0.991);
    probs.push_back(0.992);
    probs.push_back(0.993);
    probs.push_back(0.994);
    probs.push_back(0.995);
    probs.push_back(0.996);
    probs.push_back(0.997);
    probs.push_back(0.998);
    probs.push_back(0.999);

    accumulator_t acc(extended_p_square_probabilities = probs);
    accumulator_t_weighted acc_weighted(extended_p_square_probabilities = probs);
    accumulator_t_quadratic acc2(extended_p_square_probabilities = probs);
    accumulator_t_weighted_quadratic acc_weighted2(extended_p_square_probabilities = probs);

    for (int i=0; i<10000; ++i)
    {
        double sample = rng();
        acc(sample);
        acc2(sample);
        acc_weighted(sample, weight = 1.);
        acc_weighted2(sample, weight = 1.);
    }

    for (std::size_t i = 0; i < probs.size() - 1; ++i)
    {
        BOOST_CHECK_CLOSE(
            quantile(acc, quantile_probability = 0.99025 + i*0.001)
          , 0.99025 + i*0.001
          , epsilon
        );
        BOOST_CHECK_CLOSE(
            quantile(acc2, quantile_probability = 0.99025 + i*0.001)
          , 0.99025 + i*0.001
          , epsilon
        );
        BOOST_CHECK_CLOSE(
            quantile(acc_weighted, quantile_probability = 0.99025 + i*0.001)
          , 0.99025 + i*0.001
          , epsilon
        );
        BOOST_CHECK_CLOSE(
            quantile(acc_weighted2, quantile_probability = 0.99025 + i*0.001)
          , 0.99025 + i*0.001
          , epsilon
        );
    }
}

///////////////////////////////////////////////////////////////////////////////
// init_unit_test_suite
//
test_suite* init_unit_test_suite( int argc, char* argv[] )
{
    test_suite *test = BOOST_TEST_SUITE("extended_p_square_quantile test");

    test->add(BOOST_TEST_CASE(&test_stat));

    return test;
}


#pragma once

/*
 * Model a non-convex optimization problem by defining Cost and Constraint objects
 * which know how to generate a convex approximation
 *
 *
 */

#include <vector>
#include <boost/shared_ptr.hpp>
#include <trajopt_sco/sco_fwd.hpp>
#include <trajopt_sco/solver_interface.hpp>

namespace sco {

using std::vector;

/**
Stores convex terms in a objective
For non-quadratic terms like hinge(x) and abs(x), it needs to add auxilliary variables and linear constraints to the model
Note: When this object is deleted, the constraints and variables it added to the model are removed
 */
class ConvexObjective {
public:
  ConvexObjective(Model* model) : model_(model) {}
  void addAffExpr(const AffExpr&);
  void addQuadExpr(const QuadExpr&);
  void addHinge(const AffExpr&, double coeff);
  void addAbs(const AffExpr&, double coeff);
  void addHinges(const AffExprVector&);
  void addL1Norm(const AffExprVector&);
  void addL2Norm(const AffExprVector&);
  void addMax(const AffExprVector&);

  bool inModel() {
    return model_ != NULL;
  }
  void addConstraintsToModel();
  void removeFromModel();
  double value(const vector<double>& x);

  ~ConvexObjective();

  Model* model_;
  QuadExpr quad_;
  vector<Var> vars_;
  vector<AffExpr> eqs_;
  vector<AffExpr> ineqs_;
  vector<Cnt> cnts_;
private:
  ConvexObjective()  {}
  ConvexObjective(ConvexObjective&)  {}
};

/**
Stores convex inequality constraints and affine equality constraints.
Actually only affine inequality constraints are currently implemented.
*/
class ConvexConstraints {
public:
  ConvexConstraints(Model* model) : model_(model) {}
  /** Expression that should == 0 */
  void addEqCnt(const AffExpr&);
  /** Expression that should <= 0 */
  void addIneqCnt(const AffExpr&);
  void setModel(Model* model) {
    assert(!inModel());
    model_ = model;
  }
  bool inModel() {
    return model_ != NULL;
  }
  void addConstraintsToModel();
  void removeFromModel();

  vector<double> violations(const vector<double>& x);
  double violation(const vector<double>& x);

  ~ConvexConstraints();
  vector<AffExpr> eqs_;
  vector<AffExpr> ineqs_;
private:
   Model* model_;
   vector<Cnt> cnts_;
   ConvexConstraints() : model_(NULL) {}
   ConvexConstraints(ConvexConstraints&) {}
};

/**
Non-convex cost function, which knows how to calculate its convex approximation (convexify() method)
*/
class Cost {
public:
  /** Evaluate at solution vector x*/
  virtual double value(const vector<double>&) = 0;
  /** Convexify at solution vector x*/
  virtual ConvexObjectivePtr convex(const vector<double>& x, Model* model) = 0;
  /** Get problem variables associated with this cost */
  virtual VarVector getVars() {return VarVector();}

  string name() {return name_;}
  void setName(const string& name) {name_=name;}
  Cost() : name_("unnamed") {}
  Cost(const string& name) : name_(name) {}
  virtual ~Cost() {}
protected:
  string name_;
};

/**
Non-convex vector-valued constraint function, which knows how to calculate its convex approximation
*/
class Constraint {
public:

  /** inequality vs equality */
  virtual ConstraintType type() = 0;
  /** Evaluate at solution vector x*/
  virtual vector<double> value(const vector<double>& x) = 0;
  /** Convexify at solution vector x*/
  virtual ConvexConstraintsPtr convex(const vector<double>& x, Model* model) = 0;
  /** Calculate constraint violations (positive part for inequality constraint, absolute value for inequality constraint)*/
  vector<double> violations(const vector<double>& x);
  /** Sum of violations */
  double violation(const vector<double>& x);
  /** Get problem variables associated with this constraint */
  virtual VarVector getVars() {return VarVector();}

  string name() {return name_;}
  void setName(const string& name) {name_=name;}
  Constraint() : name_("unnamed") {}
  Constraint(const string& name) : name_(name) {}
  virtual ~Constraint() {}

protected:
  string name_;
};

class EqConstraint : public Constraint{
public:
  ConstraintType type() {return EQ;}
};

class IneqConstraint : public Constraint {
public:
  ConstraintType type() {return INEQ;}
};

/**
Non-convex optimization problem
*/
class OptProb {
public:
  OptProb();
  /** create variables with bounds [-INFINITY, INFINITY]  */
  VarVector createVariables(const vector<string>& names);
  /** create variables with bounds [lb[i], ub[i] */
  VarVector createVariables(const vector<string>& names, const vector<double>& lb, const vector<double>& ub);
  /** set the lower bounds of all the variables */
  void setLowerBounds(const vector<double>& lb);
  /** set the upper bounds of all the variables */
  void setUpperBounds(const vector<double>& ub);
  /** set lower bounds of some of the variables */
  void setLowerBounds(const vector<double>& lb, const vector<Var>& vars);
  /** set upper bounds of some of the variables */
  void setUpperBounds(const vector<double>& ub, const vector<Var>& vars);
  /** Note: in the current implementation, this function just adds the constraint to the
   * model. So if you're not careful, you might end up with an infeasible problem. */
  void addLinearConstraint(const AffExpr&, ConstraintType type);
  /** Add nonlinear cost function */
  void addCost(CostPtr);
  /** Add nonlinear constraint function */
  void addConstraint(ConstraintPtr);
  void addEqConstraint(ConstraintPtr);
  void addIneqConstraint(ConstraintPtr);
  virtual ~OptProb() {}
  /** Find closest point to solution vector x that satisfies linear inequality constraints */
  vector<double> getCentralFeasiblePoint(const vector<double>& x);
  vector<double> getClosestFeasiblePoint(const vector<double>& x);

  vector<ConstraintPtr> getConstraints() const;
  vector<CostPtr>& getCosts() {return costs_;}
  vector<ConstraintPtr>& getIneqConstraints() {return ineqcnts_;}
  vector<ConstraintPtr>& getEqConstraints() {return eqcnts_;}
  DblVec& getLowerBounds() {return lower_bounds_;}
  DblVec& getUpperBounds() {return upper_bounds_;}
  ModelPtr getModel() {return model_;}
  vector<Var>& getVars() {return vars_;}
  int getNumCosts() {return costs_.size();}
  int getNumConstraints() {return eqcnts_.size() + ineqcnts_.size();}
  int getNumVars() {return vars_.size();}

protected:
  ModelPtr model_;
  vector<Var> vars_;
  vector<double> lower_bounds_;
  vector<double> upper_bounds_;
  vector<CostPtr> costs_;
  vector<ConstraintPtr> eqcnts_;
  vector<ConstraintPtr> ineqcnts_;

  OptProb(OptProb&);
};

template <typename VecType>
inline void setVec(DblVec& x, const VarVector& vars, const VecType& vals) {
  assert(vars.size() == vals.size());
  for (int i = 0; i < vars.size(); ++i) {
    x[vars[i].var_rep->index] = vals[i];
  }
}
template <typename OutVecType>
inline OutVecType getVec1(const vector<double>& x, const VarVector& vars) {
  OutVecType out(vars.size());
  for (unsigned i=0; i < vars.size(); ++i) out[i] = x[vars[i].var_rep->index];
  return out;
}

}


/**
 * MIT License
 *
 * Copyright (c) 2018 Prabhsimran Singh
 *
 * 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.
 */
#pragma once

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include <iostream>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <Eigen/Dense>
using namespace Eigen;

#include "functions.hpp"
#include "layers.hpp"
#include "lstm/cell.hpp"
#include "util.hpp"

namespace nn {

std::string SOS_TOKEN = "~";
std::string EOS_TOKEN = "#";

class LSTMNetwork {
  private:
    size_t n_layers;
    int batch_size;
    int embedding_dim;
    int n_tokens;
    int hidden_size;

    Embedding embedding;
    std::vector<LSTMCell> layers;
    Dense out;

    std::vector<std::vector<LSTMState>> states;

    friend class util::Trainer;

  public:
    explicit LSTMNetwork(const size_t &, const int &, const int &, const int &, const int &);

    MatrixXf operator()(const MatrixXf &);

    MatrixXf forward(const MatrixXf &);

    void backward(const MatrixXf &);
};

LSTMNetwork::LSTMNetwork(const size_t &n_layers, const int &hidden_size, const int &n_tokens, const int &embedding_dim, const int &batch_size)
    : embedding(Embedding(n_tokens, embedding_dim)), out(Dense(hidden_size, n_tokens)) {

    this->n_layers = n_layers;
    this->hidden_size = hidden_size;
    this->n_tokens = n_tokens;
    this->embedding_dim = embedding_dim;
    this->batch_size = batch_size;

    // initial layer (embedding -> hidden)
    layers.push_back(LSTMCell(hidden_size, embedding_dim, batch_size));
    for (size_t i = 1; i < n_layers; i++) {
        // rest of the layers (hidden -> hidden)
        layers.push_back(LSTMCell(hidden_size, hidden_size, batch_size));
    }
}

MatrixXf LSTMNetwork::operator()(const MatrixXf &inputs) {
    return forward(inputs);
}

MatrixXf LSTMNetwork::forward(const MatrixXf &inputs) {
    std::vector<LSTMState> t_states;
    MatrixXf output = embedding(inputs);
    for (size_t i = 0; i < layers.size(); i++) {
        output = layers[i](output);
        t_states.push_back(layers[i].state);
    }
    states.push_back(t_states);
    MatrixXf logits = out(output);
    return F::log_softmax(logits);
}

// void LSTMNetwork::backward(const MatrixXf &loss_grad) {

//     states.clear();
// }
} // namespace nn


#include <stan/math/rev/mat.hpp>
#include <gtest/gtest.h>
#include <iostream>
#include <sstream>
#include <vector>
#include <boost/numeric/odeint.hpp>
#include <test/unit/math/rev/mat/functor/util_cvodes.hpp>
#include <test/unit/math/prim/arr/functor/harmonic_oscillator.hpp>
#include <test/unit/math/prim/arr/functor/lorenz.hpp>

template <typename F, typename T_y0, typename T_theta>
void sho_value_test(F harm_osc,
                    std::vector<double>& y0,
                    double t0,
                    std::vector<double>& ts,
                    std::vector<double>& theta,
                    std::vector<double>& x,
                    std::vector<int>& x_int) {

  using stan::math::var;
  using stan::math::promote_scalar;

  std::vector<std::vector<var> > ode_res_vd
    = stan::math::integrate_ode_bdf(harm_osc, promote_scalar<T_y0>(y0), t0,
                                    ts, promote_scalar<T_theta>(theta), x,
                                    x_int);

  EXPECT_NEAR(0.995029, ode_res_vd[0][0].val(), 1e-5);
  EXPECT_NEAR(-0.0990884, ode_res_vd[0][1].val(), 1e-5);

  EXPECT_NEAR(-0.421907, ode_res_vd[99][0].val(), 1e-5);
  EXPECT_NEAR(0.246407, ode_res_vd[99][1].val(), 1e-5);
}

void sho_finite_diff_test(double t0) {
  using stan::math::var;
  harm_osc_ode_fun harm_osc;

  std::vector<double> theta;
  theta.push_back(0.15);

  std::vector<double> y0;
  y0.push_back(1.0);
  y0.push_back(0.0);

  std::vector<double> ts;
  for (int i = 0; i < 100; i++)
    ts.push_back(t0 + 0.1 * (i + 1));

  std::vector<double> x;
  std::vector<int> x_int;

  test_ode_cvode(harm_osc, t0, ts, y0, theta, x, x_int, 1e-8, 1e-4);

  sho_value_test<harm_osc_ode_fun, double, var>
    (harm_osc, y0, t0, ts, theta, x, x_int);
  sho_value_test<harm_osc_ode_fun, var, double>
    (harm_osc, y0, t0, ts, theta, x, x_int);
  sho_value_test<harm_osc_ode_fun, var, var>
    (harm_osc, y0, t0, ts, theta, x, x_int);
}

void sho_data_finite_diff_test(double t0) {
  using stan::math::var;
  harm_osc_ode_data_fun harm_osc;

  std::vector<double> theta;
  theta.push_back(0.15);

  std::vector<double> y0;
  y0.push_back(1.0);
  y0.push_back(0.0);

  std::vector<double> ts;
  for (int i = 0; i < 100; i++)
    ts.push_back(t0 + 0.1 * (i + 1));

  std::vector<double> x(3,1);
  std::vector<int> x_int(2,0);

  test_ode_cvode(harm_osc, t0, ts, y0, theta, x, x_int, 1e-8, 1e-4);

  sho_value_test<harm_osc_ode_data_fun,double,var>(harm_osc, y0, t0, ts,
                                                   theta, x, x_int);
  sho_value_test<harm_osc_ode_data_fun,var,double>(harm_osc, y0, t0, ts,
                                                   theta, x, x_int);
  sho_value_test<harm_osc_ode_data_fun,var,var>(harm_osc, y0, t0, ts,
                                                theta, x, x_int);
}

template <typename T_y0, typename T_theta, typename F>
void sho_error_test(F harm_osc,
                    std::vector<double>& y0,
                    double t0,
                    std::vector<double>& ts,
                    std::vector<double>& theta,
                    std::vector<double>& x,
                    std::vector<int>& x_int,
                    std::string error_msg) {

  using stan::math::var;
  using stan::math::promote_scalar;

  EXPECT_THROW_MSG(stan::math::integrate_ode_bdf(harm_osc, promote_scalar<T_y0>(y0), t0,
                                                 ts, promote_scalar<T_theta>(theta), x,
                                                 x_int),
                   std::runtime_error,
                   error_msg);
}

// TODO(carpenter): g++6 failure
TEST(StanAgradRevOde_integrate_ode, harmonic_oscillator_finite_diff) {
  sho_finite_diff_test(0);
  sho_finite_diff_test(1.0);
  sho_finite_diff_test(-1.0);

  sho_data_finite_diff_test(0);
  sho_data_finite_diff_test(1.0);
  sho_data_finite_diff_test(-1.0);
}

TEST(StanAgradRevOde_integrate_ode, harmonic_oscillator_error) {
  using stan::math::var;
  harm_osc_ode_wrong_size_1_fun harm_osc;

  std::vector<double> theta;
  theta.push_back(0.15);

  std::vector<double> y0;
  y0.push_back(1.0);
  y0.push_back(0.0);

  double t0 = 0;
  std::vector<double> ts;
  for (int i = 0; i < 100; i++)
    ts.push_back(t0 + 0.1 * (i + 1));

  std::vector<double> x(3,1);
  std::vector<int> x_int(2,0);

  std::string error_msg
    = "ode_system: size of state vector y (2) and derivative vector dy_dt (3)"
    " in the ODE functor do not match in size.";

  sho_error_test<double,var>(harm_osc, y0, t0, ts,
                             theta, x, x_int, error_msg);
  sho_error_test<var,double>(harm_osc, y0, t0, ts,
                             theta, x, x_int, error_msg);
  sho_error_test<var,var>(harm_osc, y0, t0, ts,
                          theta, x, x_int, error_msg);
}

// TODO(carpenter): g++6 failure
TEST(StanAgradRevOde_integrate_ode, lorenz_finite_diff) {
  lorenz_ode_fun lorenz;

  std::vector<double> y0;
  std::vector<double> theta;
  double t0;
  std::vector<double> ts;

  t0 = 0;

  theta.push_back(10.0);
  theta.push_back(28.0);
  theta.push_back(8.0/3.0);
  y0.push_back(10.0);
  y0.push_back(1.0);
  y0.push_back(1.0);

  std::vector<double> x;
  std::vector<int> x_int;

  for (int i = 0; i < 100; i++)
    ts.push_back(0.1*(i+1));

  test_ode_cvode(lorenz, t0, ts, y0, theta, x, x_int, 1e-8, 1e-1);
}


//    Copyright 2017 Rainer Gemulla
//
//    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.
/** \file
 *
 * Illustrates matrix factorization with DSGD. We first creates factors and then a data matrix
 * from these factors. This process ensures that we know the best factorization of the input.
 * These matrices are distributed across a cluster. We then try to reconstruct the factors
 * using DSGD.
 *
 * Run with: mpirun --hosts localhost,localhost dsgd
 * (make sure to use a production build, otherwise it will be slow)
 */
#include <iostream>

#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/math/distributions/normal.hpp>
#include <boost/random/uniform_real.hpp>

#include <util/evaluation.h>

#include <mpi2/mpi2.h>
#include <mf/mf.h>

log4cxx::LoggerPtr logger(log4cxx::Logger::getLogger("main"));

using namespace std;
using namespace mf;
using namespace mpi2;
using namespace rg;
using namespace boost::numeric::ublas;

// type of SGD
typedef UpdateTruncate<UpdateNzslL2> Update;
//typedef UpdateTruncate<UpdateNzsl> Update;
typedef RegularizeNone Regularize;
//typedef RegularizeL2 Regularize;
//typedef RegularizeNzl2 Regularize;
typedef SumLoss<NzslLoss, L2Loss> Loss;
//typedef SumLoss<NzslLoss, Nzl2Loss> Loss;
typedef NzslLoss TestLoss;

int main(int argc, char* argv[]) {
	// initialize mf library and mpi2
	boost::mpi::communicator& world = mfInit(argc, argv);

	// parameters for the factorization
	mf_size_type size1 = 10000;
	mf_size_type size2 = 10000;
	mf_size_type nnz = 1000000;
	double sigma = sqrt(10); // standard deviation
	double lambda = 1/sigma/sigma;
	mf_size_type r = 10;

	// parameters for distribution
	int tasksPerRank = 4;
	mf_size_type blocks1 = world.size() * tasksPerRank;
	mf_size_type blocks2 = world.size() * tasksPerRank;

	// parameters for SGD
	double eps0 = 0.01;
	mf_size_type epochs = 10;
	SgdOrder order = SGD_ORDER_WOR;
	StratumOrder stratumOrder = STRATUM_ORDER_RSEQ;
	Update update = Update(UpdateNzslL2(lambda), -10*sigma, 10*sigma); // truncate for numerical stability
//	Update update = Update(UpdateNzsl(), -10*sigma, 10*sigma); // truncate for numerical stability
	Regularize regularize;
//	Regularize regularize = Regularize(lambda);
	Loss loss((NzslLoss()), L2Loss(lambda));
//	Loss loss((NzslLoss()), Nzl2Loss(lambda));
	TestLoss testLoss;
	mf_size_type testNnz = nnz/100;
	BalanceType balanceType = BALANCE_NONE;
	BalanceMethod balanceMethod = BALANCE_OPTIMAL;
	bool mapReduce = false;

	// start mf library
	mfStart();

	if (world.rank() == 0)
	{
#ifndef NDEBUG
		LOG4CXX_WARN(logger, "Warning: Debug mode activated (runtimes may be slow).");
#endif

		// TODO: distribute matrix generation
		// generate original factors by sampling from a normal(0,sigma) distribution
		Random32 random; // note: this takes a default seed (not randomized!)
		DenseMatrix wIn(size1, r);
		DenseMatrixCM hIn(r, size2);
		generateRandom(wIn, random, boost::normal_distribution<>(0, sigma));
		generateRandom(hIn, random, boost::normal_distribution<>(0, sigma));

		// generate a sparse matrix by selecting random entries from the generated factors
		// and add small Gaussian noise
		SparseMatrix v;
		generateRandom(v, nnz, wIn, hIn, random);
		addRandom(v, random, boost::normal_distribution<>(0, 0.1));
		LOG4CXX_INFO(logger, "Data matrix: "
			<< v.size1() << " x " << v.size2() << ", " << v.nnz() << " nonzeros");
		LOG4CXX_INFO(logger, "Loss with original factors: " << loss((FactorizationData<>(v, wIn, hIn))));

		// create a test matrix (without noise)
		SparseMatrix vTest;
		generateRandom(vTest, testNnz, wIn, hIn, random);
		LOG4CXX_INFO(logger, "Test matrix: "
			<< v.size1() << " x " << v.size2() << ", " << vTest.nnz() << " nonzeros");

		// take a small sample and remove empty rows/columns
		ProjectedSparseMatrix Vsample;
		projectRandomSubmatrix(random, v, Vsample, v.size1()/5, v.size2()/5);
		projectFrequent(Vsample, 0);
		LOG4CXX_INFO(logger, "Sample matrix: "
			<< Vsample.data.size1() << " x " << Vsample.data.size2()
			<< ", " << Vsample.data.nnz() << " nonzeros");

		// generate initial factors by sampling from a uniform[-0.5,0.5] distribution
		DenseMatrix w(size1, r);
		DenseMatrixCM h(r, size2);
		generateRandom(w, random, boost::uniform_real<>(-0.5, 0.5));
		generateRandom(h, random, boost::uniform_real<>(-0.5, 0.5));

		// distribute the input matrices and test matrix (cheaper than fetching W/H)
		DistributedSparseMatrix dv = distributeMatrix("V", blocks1, blocks2, true, v);
		LOG4CXX_INFO(logger, "Distributed data matrix: "
					<< dv.blocks1() << " x " << dv.blocks2() << " blocks");
		DistributedSparseMatrix dvTest = distributeMatrix("VTest", blocks1, blocks2, true, vTest);
		DistributedDenseMatrix dw = distributeMatrix("W", blocks1, 1, true, w);
		DistributedDenseMatrixCM dh = distributeMatrix("H", 1, blocks2, false, h);
		LOG4CXX_INFO(logger, "Distributed factor matrices");

		// initialize the DSGD
		Timer t;
		DsgdRunner dsgdRunner(random);
		DsgdJob<Update,Regularize> dsgdJob(dv, dw, dh, update, regularize, order,
				stratumOrder, mapReduce, tasksPerRank);
		//BoldDriver decay(eps0);
		DistributedDecayAuto<Update,Regularize,Loss> decay(dsgdJob, loss, Vsample, "decay", eps0,
				world.size()*tasksPerRank);
		Trace trace;

		//std::string s1("Loss"),s2("nzsl"),s3("Regularize"),s4("L2"),s5("nodes"),s6("threads");

		trace.addField("Loss","nzsl");
		trace.addField("Regularize","L2");

		trace.addField("nodes",3);
		trace.addField("threads",4);

		// print the test loss
		DsgdFactorizationData<> testData(dvTest,dw,dh,tasksPerRank);
		LOG4CXX_INFO(logger, "Initial test loss: " << testLoss(testData));

		// run DSGD to try to reconstruct the original factors
		t.start();
		dsgdRunner.run(dsgdJob, loss, epochs, decay, trace, balanceType, balanceMethod, &testData, &testLoss);
		//dsgdRunner.run(dsgdJob, loss, epochs, decay, trace);

		t.stop();
		LOG4CXX_INFO(logger, "Total time: " << t);

		// print the test loss
		LOG4CXX_INFO(logger, "Final test loss: " << testLoss(testData));

		// write trace to an R file
		LOG4CXX_INFO(logger, "Writing trace to " << "/tmp/dsgd-trace.R");
		trace.toRfile("/tmp/dsgd-trace.R", "dsgd");

//		cout << squaredSums2(dsgdJob.dw, tasksPerRank) << endl;
//		cout << squaredSums1(dsgdJob.dh, tasksPerRank) << endl;
	}

	mfStop();
	mfFinalize();

	return 0;
}


#include <Eigen/Core>
#include <Eigen/LU>
#include <elasty/alembic-manager.hpp>
#include <elasty/fem.hpp>
#include <timer.hpp>
#include <vector>

namespace
{
    constexpr size_t k_num_dims = 2;

    constexpr double k_youngs_modulus = 800.0;
    constexpr double k_poisson_ratio  = 0.40;

    constexpr double k_first_lame  = elasty::fem::calcFirstLame(k_youngs_modulus, k_poisson_ratio);
    constexpr double k_second_lame = elasty::fem::calcSecondLame(k_youngs_modulus, k_poisson_ratio);

    constexpr unsigned k_num_substeps = 10;
    constexpr double   k_delta_time   = 1.0 / 60.0;

    constexpr double k_damping_factor = 0.1;

    enum class Model
    {
        CoRotational,
        StVenantKirchhoff
    };

    constexpr Model k_model = Model::CoRotational;
} // namespace

template <int N> struct Mesh
{
    using ElemList = Eigen::Matrix<std::int32_t, N + 1, Eigen::Dynamic>;

    ElemList elems;

    Eigen::VectorXd x_rest;
    Eigen::VectorXd x;
    Eigen::VectorXd v;
    Eigen::VectorXd f;

    double mass = 1.0;

    /// \details Should be precomputed
    Eigen::VectorXd lumped_mass;

    /// \details Should be precomputed
    std::vector<double> volume_array;

    /// \details Should be precomputed
    std::vector<Eigen::Matrix<double, N, N>> rest_shape_mat_inv_array;

    /// \details Should be precomputed
    std::vector<Eigen::Matrix<double, N * N, N*(N + 1)>> vec_PFPx_array;
};

using TriangleMesh = Mesh<2>;

template <typename Derived> typename Derived::Scalar calcEnergyDensity(const Eigen::MatrixBase<Derived>& deform_grad)
{
    switch (k_model)
    {
        case Model::StVenantKirchhoff:
            return elasty::fem::calcStVenantKirchhoffEnergyDensity(deform_grad, k_first_lame, k_second_lame);
        case Model::CoRotational:
            return elasty::fem::calcCoRotationalEnergyDensity(deform_grad, k_first_lame, k_second_lame);
    }
}

template <typename Derived>
Eigen::Matrix<typename Derived::Scalar, 2, 2> calcPiolaStress(const Eigen::MatrixBase<Derived>& deform_grad)
{
    switch (k_model)
    {
        case Model::StVenantKirchhoff:
            return elasty::fem::calcStVenantKirchhoffPiolaStress(deform_grad, k_first_lame, k_second_lame);
        case Model::CoRotational:
            return elasty::fem::calcCoRotationalPiolaStress(deform_grad, k_first_lame, k_second_lame);
    }
}

class Explicit2dEngine
{
public:
    Explicit2dEngine() {}

    void proceedFrame()
    {
        const std::size_t              num_verts         = m_mesh.x_rest.size() / k_num_dims;
        const std::vector<std::size_t> constrained_verts = {0, 1, 2, 3, 4};

        // Reset forces
        m_mesh.f = Eigen::VectorXd::Zero(2 * num_verts);

        // Calculate forces
        for (std::size_t i = 0; i < m_mesh.elems.cols(); ++i)
        {
            const auto& indices = m_mesh.elems.col(i);

            // Retrieve precomputed values
            const auto& D_m_inv  = m_mesh.rest_shape_mat_inv_array[i];
            const auto& area     = m_mesh.volume_array[i];
            const auto& vec_PFPx = m_mesh.vec_PFPx_array[i];

            // Calculate the deformation gradient $\mathbf{F}$
            const auto F = elasty::fem::calc2dTriangleDeformGrad(m_mesh.x.segment<2>(2 * indices[0]),
                                                                 m_mesh.x.segment<2>(2 * indices[1]),
                                                                 m_mesh.x.segment<2>(2 * indices[2]),
                                                                 D_m_inv);

            // Calculate $\frac{\partial \Phi}{\partial \mathbf{x}}$ and related values
            const auto P      = calcPiolaStress(F);
            const auto vec_P  = Eigen::Map<const Eigen::Vector4d>(P.data(), P.size());
            const auto PPsiPx = vec_PFPx.transpose() * vec_P;

            // Calculate the internal force
            const auto force = -area * PPsiPx;

            assert(indices.size() == 3);
            assert(F.rows() == 2 && F.cols() == 2);
            assert(vec_PFPx.rows() == 4 && vec_PFPx.cols() == 6);
            assert(PPsiPx.rows() == 6 && PPsiPx.cols() == 1);

            // Accumulate forces
            m_mesh.f.segment(2 * indices[0], 2) += force.segment(2 * 0, 2);
            m_mesh.f.segment(2 * indices[1], 2) += force.segment(2 * 1, 2);
            m_mesh.f.segment(2 * indices[2], 2) += force.segment(2 * 2, 2);
        }

        // Apply gravity force
        for (size_t i = 0; i < num_verts; ++i)
        {
            m_mesh.f[i * 2 + 1] += m_mesh.lumped_mass(i * 2 + 1) * (-9.80665);
        }

        // Calculate the "modified" inverse lumped mass matrix
        // Note: This "mass modification" is described in "Large Steps in Cloth Simulation" (SIGGRAPH '98)
        Eigen::VectorXd W_diags = m_mesh.lumped_mass.cwiseInverse();
        for (size_t i : constrained_verts)
        {
            W_diags.segment(i * 2, 2) = Eigen::Vector2d::Zero();
        }
        const auto W = W_diags.asDiagonal();

        // Note: Explicit Euler integration (for velocities)
        m_mesh.v = m_mesh.v + m_delta_physics_time * W * m_mesh.f;

        // Note: Explicit Euler integration (for positions)
        m_mesh.x = m_mesh.x + m_delta_physics_time * m_mesh.v;

        // Apply naive damping
        m_mesh.v *= std::exp(-k_damping_factor * m_delta_physics_time);
    }

    void initializeScene()
    {
        // A simple cantilever

        constexpr std::size_t num_cols  = 20;
        constexpr std::size_t num_rows  = 4;
        constexpr std::size_t num_verts = (num_cols + 1) * (num_rows + 1);
        constexpr std::size_t num_elems = (num_cols * num_rows) * 2;
        constexpr double      size      = 1.0;

        m_mesh.elems.resize(3, num_elems);
        m_mesh.x_rest.resize(num_verts * k_num_dims);

        // Generate a triangle mesh
        for (std::size_t col = 0; col < num_cols; ++col)
        {
            for (std::size_t row = 0; row < num_rows; ++row)
            {
                const auto base = col * (num_rows + 1) + row;

                m_mesh.elems.col(2 * num_rows * col + 2 * row + 0) << 0 + base, 1 + base, (num_rows + 1) + 1 + base;
                m_mesh.elems.col(2 * num_rows * col + 2 * row + 1) << 0 + base, (num_rows + 1) + 1 + base,
                    (num_rows + 1) + base;

                m_mesh.x_rest.segment(k_num_dims * ((num_rows + 1) * col + row), k_num_dims) =
                    Eigen::Vector2d{col * 1.0, -1.0 * row};
            }
            m_mesh.x_rest.segment(k_num_dims * ((num_rows + 1) * col + num_rows), k_num_dims) =
                Eigen::Vector2d{col * 1.0, -1.0 * num_rows};
        }
        for (std::size_t row = 0; row < num_rows; ++row)
        {
            m_mesh.x_rest.segment(k_num_dims * ((num_rows + 1) * num_cols + row), k_num_dims) =
                Eigen::Vector2d{num_cols * 1.0, -1.0 * row};
        }
        m_mesh.x_rest.segment(k_num_dims * ((num_rows + 1) * num_cols + num_rows), k_num_dims) =
            Eigen::Vector2d{num_cols * 1.0, -1.0 * num_rows};

        // Set transform
        m_mesh.x_rest *= 1.0 / static_cast<double>(num_rows);
        for (std::size_t vert = 0; vert < num_verts; ++vert)
        {
            m_mesh.x_rest[2 * vert + 1] += 0.5;
        }
        m_mesh.x_rest *= size;

        // Initialize other values
        m_mesh.x = m_mesh.x_rest;
        m_mesh.v = Eigen::VectorXd::Zero(k_num_dims * num_verts);
        m_mesh.f = Eigen::VectorXd::Zero(k_num_dims * num_verts);

        // Perform precomputation
        m_mesh.volume_array.resize(num_elems);
        m_mesh.rest_shape_mat_inv_array.resize(num_elems);
        m_mesh.vec_PFPx_array.resize(num_elems);
        for (std::size_t elem_index = 0; elem_index < num_elems; ++elem_index)
        {
            const auto& indices = m_mesh.elems.col(elem_index);

            m_mesh.volume_array[elem_index] = elasty::fem::calc2dTriangleArea(m_mesh.x_rest.segment<2>(2 * indices[0]),
                                                                              m_mesh.x_rest.segment<2>(2 * indices[1]),
                                                                              m_mesh.x_rest.segment<2>(2 * indices[2]));
            m_mesh.rest_shape_mat_inv_array[elem_index] =
                elasty::fem::calc2dShapeMatrix(m_mesh.x_rest.segment<2>(2 * indices[0]),
                                               m_mesh.x_rest.segment<2>(2 * indices[1]),
                                               m_mesh.x_rest.segment<2>(2 * indices[2]))
                    .inverse();
            m_mesh.vec_PFPx_array[elem_index] =
                elasty::fem::calcVecTrianglePartDeformGradPartPos(m_mesh.rest_shape_mat_inv_array[elem_index]);
        }
        m_mesh.lumped_mass = elasty::fem::calcTriangleMeshLumpedMass(m_mesh.x_rest, m_mesh.elems, m_mesh.mass);
    }

    /// \brief Getter of the delta physics time.
    ///
    /// \details The value equals to the delta frame time devided by the number of substeps.
    double getDeltaPhysicsTime() const { return m_delta_physics_time; }

    void setDeltaPhysicsTime(const double delta_physics_time) { m_delta_physics_time = delta_physics_time; }

    const TriangleMesh* getMesh() const { return &m_mesh; }

private:
    double m_delta_physics_time = 1.0 / 60.0;

    TriangleMesh m_mesh;
};

int main(int argc, char** argv)
{
    Explicit2dEngine engine;

    engine.initializeScene();
    engine.setDeltaPhysicsTime(k_delta_time / static_cast<double>(k_num_substeps));

    const auto        mesh      = engine.getMesh();
    const std::size_t num_verts = mesh->x_rest.size() / 2;
    const std::size_t num_elems = mesh->elems.cols();

    auto alembic_manager = elasty::createTriangleMesh2dAlembicManager(
        "./out.abc", k_delta_time, num_verts, num_elems, mesh->x.data(), mesh->elems.data());

    for (unsigned int frame = 0; frame < 240; ++frame)
    {
        timer::Timer t(std::to_string(frame));

        alembic_manager->submitCurrentStatus();

        for (int i = 0; i < k_num_substeps; ++i)
        {
            engine.proceedFrame();
        }
    }

    return 0;
}


/*! \file demo_2d_limits.cpp
    \brief Demonstration of some 2D values including NaN and + and - infinity.
    \details Quickbook markup to include in documentation.
    \author Paul A. Bristow
*/

// Copyright Paul A Bristow 2008, 2009, 2013, 2020

// Use, modification and distribution are 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)

// An example to demonstrate plotting 2D 'at limts' values including NaN and + and - infinity.

// This file is written to be included from a Quickbook .qbk document.
// It can be compiled by the C++ compiler, and run. Any output can
// also be added here as comment or included or pasted in elsewhere.

// Caution: this file contains Quickbook markup as well as code
// and comments: don't change any of the special comment markups!

//[demo_2d_limits_1

/*`An example to demonstrate plotting 2D 'at limits' values
including NaN and + and - infinity.

As ever, we need the usual includes to use Boost.Plot.

*/
//] [demo_2d_limits_1]

#include <boost/svg_plot/svg_2d_plot.hpp>
  using namespace boost::svg;
  using boost::svg::svg_2d_plot;
#include <iostream>
  using std::cout;
  using std::endl;

#include <map>
  using std::map;

#include <limits>
  using std::numeric_limits;

#include <utility>
  using std::make_pair;

int main()
{
//[demo_2d_limits_2
  /*`Some fictional data is pushed into an STL container, here map:*/

  map<double, double> my_data;
  /*`
  Inserting some fictional values also sorts the data.
  The map index value in [ ] is the x value, so mydata[x] = y.

  First some normal valued points, not 'at limits'.
  */
  my_data[1.1] = 3.2;
  my_data[4.3] = 3.1;
  my_data[0.25] = 1.4;
/*`
  Now some values including + and - infinity:
*/
  my_data[3] = numeric_limits<double>::quiet_NaN(); // marker at x = 3, y = 0
  my_data[0.] = numeric_limits<double>::quiet_NaN(); // Marker at 0,0
  my_data[1.] = numeric_limits<double>::infinity(); // Marker at 1, top
  my_data[-1] = -numeric_limits<double>::infinity(); // Marker at -1, bottom
  my_data[+numeric_limits<double>::infinity()] = +1.; // Marker at right, 1
  my_data[-numeric_limits<double>::infinity()] = -1.; // Marker at left, -1
  my_data[+(numeric_limits<double>::max)()] = +2.;  // Marker at right, 2
  my_data[-(numeric_limits<double>::max)()] = +2.; // Marker at left, 2
  my_data[-(numeric_limits<double>::max)() /2] = +3.; // Value near to max, marker left, 3
  my_data[numeric_limits<double>::infinity()] = numeric_limits<double>::infinity(); // Top right.
  my_data[-numeric_limits<double>::infinity()] = -numeric_limits<double>::infinity(); // Bottom left.

/*`
  [caution Using map (rather than multimap that allows duplicates) some assignments values overwrite,
  and so not all display as they do individually.
  In particular, an X value of quiet_NaN() causes a overwrite of the lowerest value (because NaNs never compare equal).
  So avoid NaN as an X value.]
*/

  try
  { // try'n'catch blocks are needed to ensure error messages from any exceptions are shown.
    svg_2d_plot my_2d_plot; // Construct a plot with all the default constructor values.

    my_2d_plot.title("Default 2D 'at limits' NaN and infinities Demo") // Add a string title of the plot.
      .x_range(-5, 5) // Add a range for the X-axis.
      .y_range(-5, 5) // Add a range for the Y-axis
      .x_label("time (s)"); // Add a label for the X-axis.

/*`
Add the one data series, `my_data` and a description, and how the data points are to be marked,
here a circle with a diameter of 5 pixels.
*/
    svg_2d_plot_series& my_series = my_2d_plot.plot(my_data, "2D limits").shape(circlet).size(5);

/*`
We can also keep note of the plot series and use this to interrogate how many normal and how many 'at limit' values.
*/
    cout << my_series.values_count() << " normal data values in series." << endl;
    cout << my_series.limits_count() << " 'at limits' data values in series."<< endl;

/*`To put a value label against each data point, switch on the option:
*/
    //  my_2d_plot.x_values_on(true).y_values_on(true).x_values_font_size(12).y_values_font_size(12); // Add the X-axis and Y-axis values.
    // This displays x horizontally and Y downward.
    my_2d_plot.xy_values_on(true).x_values_font_size(12).y_values_font_size(12); // Add the X-axis and Y-axis values.
    // This displays X above and Y below.

/*`To change the default colors (lightgray and whitesmoke) for the 'at limit' point marker
to something more conspicuous for this demonstration:
*/
    my_2d_plot.plus_inf_limit_color(blue);
    my_2d_plot.plus_inf_limit_color(pink);

/*`To use all these settings, finally write the plot to file.
*/
    my_2d_plot.write("demo_2d_limits.svg");

/*` Note the +infinity point is marked on the far right of the plot, the -infinity on the far left, but the NaN (Not A Number is at zero).

To echo the new marker colors chosen:
*/
    cout << "limit points stroke color " << my_2d_plot.plus_inf_limit_color() << endl;
    cout << "limit points fill color " << my_2d_plot.plus_inf_limit_color() << endl;

//] [/demo_2d_limits_2]
  }
  catch(const std::exception& e)
  {
    std::cout <<
      "\n""Message from thrown exception was:\n   " << e.what() << std::endl;
  }
  return 0;
} // int main()

/*

//[demo_2d_limits_output

Output:

Autorun "j:\Cpp\SVG\Debug\demo_2d_limits.exe"
3 normal data values in series.
9 'at limits' data values in series.
limit points stroke color RGB(0,0,255)
limit points fill color RGB(255,192,203)

*/
//] //[/demo_2d_limits_output]


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

using namespace std;

int main()
{
    boost::random::mt19937 rng;     // produces randomness out of thin air
                                    // see pseudo-random number generators
    boost::random::uniform_int_distribution<> six(1,6);
                                    // distribution that maps to 1..6
                                    // see random number distributions
    cout <<six(rng) << endl;;       // simulate rolling a die

    double probabilities[] = {
    0.5, 0.1, 0.1, 0.1, 0.1, 0.1
    };
    boost::random::discrete_distribution<> dist(probabilities);
    for (int i = 0; i < 10; i++) cout << dist(rng) << endl;

    boost::random::uniform_smallint<>si(1, 100);
    for (int i = 0; i < 10; i++) cout << si(rng) << endl;

    boost::random::uniform_01<>uniform;
    for (int i = 0; i < 10; i++) cout << uniform(rng) << endl;

    boost::random::uniform_real_distribution<>rd(0.1, 0.6);
    for (int i = 0; i < 10; i++) cout << rd(rng) << endl;
    return 0;
}


/*
@copyright Louis Dionne 2014
Distributed under the Boost Software License, Version 1.0.
(See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
 */

#include <boost/hana/comparable/equal_mcd.hpp>
#include <boost/hana/core/datatype.hpp>
#include <boost/hana/detail/assert.hpp>
#include <boost/hana/functional.hpp>
#include <boost/hana/functor/fmap_mcd.hpp>
#include <boost/hana/integral.hpp>
#include <boost/hana/list/instance.hpp>
#include <boost/hana/range.hpp>
#include <boost/hana/type.hpp>
using namespace boost::hana;

struct Matrix;

template <typename Storage, typename = operators<Comparable>>
struct matrix_type {
    using hana_datatype = Matrix;

    Storage rows_;
    constexpr auto ncolumns() const
    { return length(head(rows_)); }

    constexpr auto nrows() const
    { return length(rows_); }

    constexpr auto size() const
    { return nrows() * ncolumns(); }

    template <typename I, typename J>
    constexpr auto at(I i, J j) const
    { return boost::hana::at(j, boost::hana::at(i, rows_)); }
};

auto transpose = [](auto m) {
    auto new_storage = unpack(m.rows_, zip);
    return matrix_type<decltype(new_storage)>{new_storage};
};

auto rows = [](auto m) {
    return m.rows_;
};

auto columns = [](auto m) {
    return rows(transpose(m));
};

auto elementwise = [](auto f) {
    return [=](auto ...matrices) {
        auto new_storage = zip_with(partial(zip_with, f), matrices.rows_...);
        return matrix_type<decltype(new_storage)>{new_storage};
    };
};

template <typename S1, typename S2>
constexpr auto operator+(matrix_type<S1> m1, matrix_type<S2> m2)
{ return elementwise(_+_)(m1, m2); }

template <typename S1, typename S2>
constexpr auto operator-(matrix_type<S1> m1, matrix_type<S2> m2)
{ return elementwise(_-_)(m1, m2); }

auto scalar_prod = [](auto v1, auto v2) {
    return sum(zip_with(_*_, v1, v2));
};
auto repeat_n = [](auto n, auto x) {
    return unpack(range(int_<0>, n), on(list, always(x)));
};

template <typename S1, typename S2>
constexpr auto operator*(matrix_type<S1> m1, matrix_type<S2> m2) {
    auto storage = fmap(
        [=](auto row) {
            return zip_with(
                    scalar_prod,
                    repeat_n(m2.ncolumns(), row),
                    columns(m2));
        }
    , rows(m1));
    return matrix_type<decltype(storage)>{storage};
}

auto row = [](auto ...entries) {
    return list(entries...);
};

auto matrix = [](auto ...rows) {
    auto storage = list(rows...);
    auto all_same_length = all(tail(storage), [=](auto row) {
        return length(row) == length(head(storage));
    });
    static_assert(all_same_length, "");

    return matrix_type<decltype(storage)>{storage};
};

auto vector = on(matrix, row);

constexpr int exponent(int x, unsigned int n) {
    if (n == 0) return 1;
    else        return x * exponent(x, n - 1);
}

auto remove_at = [](auto n, auto xs) {
    auto with_indices = zip(xs, range(int_<0>, length(xs)));
    auto removed = filter(with_indices, compose(n != _, last));
    return fmap(head, removed);
};

template <typename Matrix>
struct _det; // remove circular dependency between matrix_minor and det

auto matrix_minor = [](auto m, auto i, auto j) {
    auto submatrix_storage = fmap(partial(remove_at, j), remove_at(i, rows(m)));
    matrix_type<decltype(submatrix_storage)> submatrix{submatrix_storage};
    return _det<decltype(submatrix)>{}(submatrix);
};

auto cofactor = [](auto m, auto i, auto j) {
    auto i_plus_j = i + j;
    return int_<exponent(-1, i_plus_j())> * matrix_minor(m, i, j);
};

template <typename Matrix>
struct _det {
    constexpr auto operator()(Matrix m) const {
        return eval_if(m.size() == int_<1>,
            always(m.at(int_<0>, int_<0>)),
            [=](auto _) {
                auto cofactors_1st_row = unpack(_(range)(int_<0>, m.ncolumns()),
                    on(list, partial(cofactor, m, int_<0>))
                );
                return scalar_prod(head(rows(m)), cofactors_1st_row);
            }
        );
    }
};

auto det = [](auto m) {
    return _det<decltype(m)>{}(m);
};

namespace boost { namespace hana {
    template <>
    struct Functor::instance<Matrix> : Functor::fmap_mcd {
        template <typename F, typename M>
        static constexpr auto fmap_impl(F f, M mat) {
            auto new_storage = fmap(partial(fmap, f), mat.rows_);
            return matrix_type<decltype(new_storage)>{new_storage};
        }
    };

    template <>
    struct Comparable::instance<Matrix, Matrix> : Comparable::equal_mcd {
        template <typename M1, typename M2>
        static constexpr auto equal_impl(M1 m1, M2 m2) {
            return m1.nrows() == m2.nrows() &&
                   m1.ncolumns() == m2.ncolumns() &&
                   all_of(zip_with(_==_, m1.rows_, m2.rows_));
        }
    };
}}

void test_sizes() {
    auto m = matrix(
        row(1, '2', 3),
        row('4', char_<'5'>, 6)
    );
    BOOST_HANA_CONSTEXPR_ASSERT(m.size() == 6);
    BOOST_HANA_CONSTEXPR_ASSERT(m.ncolumns() == 3);
    BOOST_HANA_CONSTEXPR_ASSERT(m.nrows() == 2);
}

void test_at() {
    auto m = matrix(
        row(1, '2', 3),
        row('4', char_<'5'>, 6),
        row(int_<7>, '8', 9.3)
    );
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<0>, int_<0>) == 1);
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<0>, int_<1>) == '2');
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<0>, int_<2>) == 3);

    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<1>, int_<0>) == '4');
    BOOST_HANA_CONSTANT_ASSERT(m.at(int_<1>, int_<1>) == char_<'5'>);
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<1>, int_<2>) == 6);

    BOOST_HANA_CONSTANT_ASSERT(m.at(int_<2>, int_<0>) == int_<7>);
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<2>, int_<1>) == '8');
    BOOST_HANA_CONSTEXPR_ASSERT(m.at(int_<2>, int_<2>) == 9.3);
}

void test_comparable() {
    BOOST_HANA_CONSTEXPR_ASSERT(matrix(row(1, 2)) == matrix(row(1, 2)));
    BOOST_HANA_CONSTEXPR_ASSERT(matrix(row(1, 2)) != matrix(row(1, 5)));

    BOOST_HANA_CONSTEXPR_ASSERT(matrix(row(1, 2), row(3, 4)) == matrix(row(1, 2), row(3, 4)));
    BOOST_HANA_CONSTEXPR_ASSERT(matrix(row(1, 2), row(3, 4)) != matrix(row(1, 2), row(0, 4)));
    BOOST_HANA_CONSTEXPR_ASSERT(matrix(row(1, 2), row(3, 4)) != matrix(row(0, 2), row(3, 4)));

    BOOST_HANA_CONSTANT_ASSERT(matrix(row(1), row(2)) != matrix(row(3, 4), row(5, 6)));
    BOOST_HANA_CONSTANT_ASSERT(matrix(row(1), row(2)) != matrix(row(3, 4)));
}

void test_functor() {
    auto m = matrix(
        row(1, int_<2>, 3),
        row(int_<4>, 5, 6),
        row(7, 8, int_<9>)
    );
    BOOST_HANA_CONSTEXPR_ASSERT(fmap(_ + int_<1>, m) ==
        matrix(
            row(2, int_<3>, 4),
            row(int_<5>, 6, 7),
            row(8, 9, int_<10>)
        )
    );
}

void test_operators() {
    auto m = matrix(row(1, 2), row(3, 4));
    BOOST_HANA_CONSTEXPR_ASSERT(m + m == matrix(row(2, 4), row(6, 8)));
    BOOST_HANA_CONSTEXPR_ASSERT(m - m == matrix(row(0, 0), row(0, 0)));
}

void test_matrix_multiplication() {
    auto A = matrix(
        row(1, 2, 3),
        row(4, 5, 6)
    );

    auto B = matrix(
        row(1, 2),
        row(3, 4),
        row(5, 6)
    );

    BOOST_HANA_CONSTEXPR_ASSERT(A * B == matrix(
        row(1*1 + 2*3 + 5*3, 1*2 + 2*4 + 3*6),
        row(4*1 + 3*5 + 5*6, 4*2 + 5*4 + 6*6)
    ));
}

void test_vector() {
    auto v = vector(1, '2', int_<3>, 4.2f);
    BOOST_HANA_CONSTEXPR_ASSERT(v.size() == 4);
    BOOST_HANA_CONSTEXPR_ASSERT(v.nrows() == 4);
    BOOST_HANA_CONSTEXPR_ASSERT(v.ncolumns() == 1);
}

void test_transpose() {
    auto m = matrix(
        row(1, 2.2, '3'),
        row(4, '5', 6)
    );
    BOOST_HANA_CONSTEXPR_ASSERT(transpose(m) ==
        matrix(
            row(1, 4),
            row(2.2, '5'),
            row('3', 6)
        )
    );
}

void test_repeat_n() {
    struct T;
    BOOST_HANA_CONSTANT_ASSERT(repeat_n(int_<0>, type<T>) == list());
    BOOST_HANA_CONSTANT_ASSERT(repeat_n(int_<1>, type<T>) == list(type<T>));
    BOOST_HANA_CONSTANT_ASSERT(repeat_n(int_<2>, type<T>) == list(type<T>, type<T>));
    BOOST_HANA_CONSTANT_ASSERT(repeat_n(int_<3>, type<T>) == list(type<T>, type<T>, type<T>));
    BOOST_HANA_CONSTANT_ASSERT(repeat_n(int_<4>, type<T>) == list(type<T>, type<T>, type<T>, type<T>));
}

void test_determinant() {
    BOOST_HANA_CONSTEXPR_ASSERT(det(matrix(row(1))) == 1);
    BOOST_HANA_CONSTEXPR_ASSERT(det(matrix(row(2))) == 2);

    BOOST_HANA_CONSTEXPR_ASSERT(det(matrix(row(1, 2), row(3, 4))) == -2);

    BOOST_HANA_CONSTEXPR_ASSERT(
        det(matrix(
            row(1, 5, 6),
            row(3, 2, 4),
            row(7, 8, 9)
        ))
        == 51
    );

    BOOST_HANA_CONSTEXPR_ASSERT(
        det(matrix(
            row(1, 5, 6, -3),
            row(3, 2, 4, -5),
            row(7, 8, 9, -1),
            row(8, 2, 1, 10)
        )) == 214
    );

    BOOST_HANA_CONSTEXPR_ASSERT(
        det(matrix(
            row(1,  5,  6, -3, 92),
            row(3,  2,  4, -5, 13),
            row(7,  8,  9, -1, 0),
            row(8,  2,  1, 10, 41),
            row(3, 12, 92, -7, -4)
        )) == -3115014
    );
}

void test_remove_at() {
    BOOST_HANA_CONSTANT_ASSERT(remove_at(int_<0>, list(1)) == list());
    BOOST_HANA_CONSTEXPR_ASSERT(remove_at(int_<0>, list(1, '2')) == list('2'));
    BOOST_HANA_CONSTEXPR_ASSERT(remove_at(int_<0>, list(1, '2', 3.3)) == list('2', 3.3));

    BOOST_HANA_CONSTEXPR_ASSERT(remove_at(int_<1>, list(1, '2')) == list(1));
    BOOST_HANA_CONSTEXPR_ASSERT(remove_at(int_<1>, list(1, '2', 3.3)) == list(1, 3.3));

    BOOST_HANA_CONSTEXPR_ASSERT(remove_at(int_<2>, list(1, '2', 3.3)) == list(1, '2'));
}

void test_exponent() {
    BOOST_HANA_CONSTEXPR_ASSERT(exponent(3, 0) == 1);
    BOOST_HANA_CONSTEXPR_ASSERT(exponent(3, 1) == 3);
    BOOST_HANA_CONSTEXPR_ASSERT(exponent(3, 2) == 3 * 3);
    BOOST_HANA_CONSTEXPR_ASSERT(exponent(3, 3) == 3 * 3 * 3);
    BOOST_HANA_CONSTEXPR_ASSERT(exponent(3, 4) == 3 * 3 * 3 * 3);
}

int main() {
    test_repeat_n();
    test_remove_at();
    test_exponent();

    test_sizes();
    test_at();
    test_comparable();
    test_functor();
    test_operators();
    test_vector();
    test_transpose();
    test_matrix_multiplication();
    test_determinant();
}


//
// Created by wei on 1/14/19.
//

#include <Eigen/Eigen>
#include <gtest/gtest.h>
#include <iostream>
#include <random>
#include <Cuda/Container/Array2DCuda.h>

TEST(Eigen, RawData) {
    Eigen::Matrix<float, 3, 5, Eigen::RowMajor> matrix;
    matrix << 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14;

    std::cout << matrix << std::endl;
    for (int i = 0; i < 15; ++i) {
        std::cout << *(matrix.data() + i) << " ";
    }
}

TEST(Eigen, UploadAndDownload) {
    std::random_device rd;
    std::default_random_engine engine(rd());
    std::uniform_int_distribution<int> uniform(100, 10000);
    int rows = uniform(engine);
    int cols = 30;

    std::cout << "rows: " << rows << std::endl;
    Eigen::Matrix<float, -1, -1, Eigen::RowMajor> matrix(rows, cols);
    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            matrix(i, j) = i + j;
            EXPECT_EQ(*(matrix.data() + i * cols + j), i + j);
        }
    }

    open3d::cuda::Array2DCuda<float> matrix_cuda;
    matrix_cuda.Create(rows, cols);
    matrix_cuda.Upload(matrix);

    EXPECT_EQ(matrix_cuda.Download(), matrix);
}

int main(int argc, char **argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}


#ifndef MONTECARLO_HPP_
#define MONTECARLO_HPP_

#include <mpi.h>
#include <boost/random.hpp>
#include <iostream>
#include <numeric>
#include <cmath>
#include <cstring>
#include <fstream>

#include "helper.hpp"

/* Monte Carlo method for UQ */
 class MonteCarlo
 {
 private:
	/* the Raynolds number will be of the form mean + stddev*sigma, where
	sigma is ~N(0,1) or ~U(0,1) */
	double par1, par2;

	/* mersenne twister random number generator */
	boost::mt19937 rng;
	/* normal(Gaussian) distribution */
	boost::normal_distribution<> normal_distr;
	/* uniform distribution */
	boost::uniform_real<> uniform_distr;

	boost::variate_generator<boost::mt19937&,
	boost::normal_distribution<> >* var_normal;

	boost::variate_generator<boost::mt19937&,
	boost::uniform_real<> >* var_uniform;

public:

	/* constructor */
	MonteCarlo(double par1, double par2, int distr_flag);

	/** Uncertainty (i.e. Random variables) related methods **/

	/* generate nsamples samples of normal distributed random variables */
	std::vector<double> generate_nd_samples(double mean_nd, double sttdev_nd, int* nsamples);
	/* generate nsamples samples of uniformly distributed random variables */
	std::vector<double> generate_ud_samples(double mean_ud, double sttdev_ud, int* nsamples);

	/* get a normal and uniform distributed RV */
	double get_normal() const;
	double get_uniform() const;

	/***********************************************************/

	/** Parallelization related methods **/

	/* data decomposition among processes*/
	void data_decomposition(int samples_per_proc, int* nsampels, int* nprocs, int *myrank, int *il, int *ir);

	/* call the NS solver for each generated sample */
	void monte_carlo_simulation(int *myrank, int* nsamples, int* samples_per_proc, int *il, int *ir
		, std::vector<double> &rv, int rv_flag, int imax, int jmax, double *mean, double* variance, int* flag_prog);

	/***********************************************************/

	/* destructor */
	~MonteCarlo();
};

#endif /* MONTECARLO_HPP_ */


// Boost.Units - A C++ library for zero-overhead dimensional analysis and
// unit/quantity manipulation and conversion
//
// Copyright (C) 2003-2008 Matthias Christian Schabel
// Copyright (C) 2008 Steven Watanabe
//
// 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)

// $Id: lambda.cpp 27 2008-06-16 14:50:58Z maehne $

////////////////////////////////////////////////////////////////////////
///
/// \file lambda.cpp
///
/// \brief Example demonstrating the usage of Boost.Units' quantity,
///        unit, and absolute types in functors created with the
///        Boost.Lambda library and stored in Boost.Function objects.
///
/// \author Torsten Maehne
/// \date   2008-06-04
///
/// A mechanical, electrical, geometrical, and thermal example
/// demonstrate how to use Boost.Units' quantity, unit, and absolute
/// types in lambda expressions. The resulting functors can be stored
/// in boost::function objects. It is also shown how to work around a
/// limitation of Boost.Lambda's bind() to help it to find the correct
/// overloaded function by specifying its signature with a
/// static_cast.
///
////////////////////////////////////////////////////////////////////////

#include <iostream>
#include <boost/function.hpp>
#include <boost/units/io.hpp>
#include <boost/units/cmath.hpp>
#include <boost/units/pow.hpp>
#include <boost/units/systems/si.hpp>
#include <boost/units/absolute.hpp>

// Include boost/units/lambda.hpp instead of boost/lambda/lambda.hpp
// for a convenient usage of Boost.Units' quantity, unit, and absolute
// types in lambda expressions. The header augments Boost.Lambda's
// return type detuction system to recognize the new types so that not
// for each arithmetic operation the return type needs to be
// explicitely specified.
#include <boost/units/lambda.hpp>

#include <boost/lambda/bind.hpp>

static const double pi = 3.14159265358979323846;

//[lambda_snippet_1

int main(int argc, char **argv) {

   using namespace std;
   namespace bl = boost::lambda;
   namespace bu = boost::units;
   namespace si = boost::units::si;

   ////////////////////////////////////////////////////////////////////////
   // Mechanical example: linear accelerated movement
   ////////////////////////////////////////////////////////////////////////

   // Initial condition variables for acceleration, speed, and displacement
   bu::quantity<si::acceleration> a = 2.0 * si::meters_per_second_squared;
   bu::quantity<si::velocity> v = 1.0 * si::meters_per_second;
   bu::quantity<si::length> s0 = 0.5 * si::meter;

   // Displacement over time
   boost::function<bu::quantity<si::length> (bu::quantity<si::time>) >
       s = 0.5 * bl::var(a) * bl::_1 * bl::_1
           + bl::var(v) * bl::_1
           + bl::var(s0);

   cout << "Linear accelerated movement:" << endl
        << "a = " << a << ", v = " << v << ", s0 = " << s0 << endl
        << "s(1.0 * si::second) = " << s(1.0 * si::second) << endl
        << endl;

   // Change initial conditions
   a = 1.0 * si::meters_per_second_squared;
   v = 2.0 * si::meters_per_second;
   s0 = -1.5 * si::meter;

   cout << "a = " << a << ", v = " << v << ", s0 = " << s0 << endl
        << "s(1.0 * si::second) = " << s(1.0 * si::second) << endl
        << endl;

   ////////////////////////////////////////////////////////////////////////
   // Electrical example: oscillating current
   ////////////////////////////////////////////////////////////////////////

   // Constants for the current amplitude, frequency, and offset current
   const bu::quantity<si::current> iamp = 1.5 * si::ampere;
   const bu::quantity<si::frequency> f = 1.0e3 * si::hertz;
   const bu::quantity<si::current> i0 = 0.5 * si::ampere;

   // The invocation of the sin function needs to be postponed using
   // bind to specify the oscillation function. A lengthy static_cast
   // to the function pointer referencing boost::units::sin() is needed
   // to avoid an "unresolved overloaded function type" error.
   boost::function<bu::quantity<si::current> (bu::quantity<si::time>) >
       i = iamp
           * bl::bind(static_cast<bu::dimensionless_quantity<si::system, double>::type (*)(const bu::quantity<si::plane_angle>&)>(bu::sin),
                      2.0 * pi * si::radian * f * bl::_1)
           + i0;

   cout << "Oscillating current:" << endl
        << "iamp = " << iamp << ", f = " << f << ", i0 = " << i0 << endl
        << "i(1.25e-3 * si::second) = " << i(1.25e-3 * si::second) << endl
        << endl;

   ////////////////////////////////////////////////////////////////////////
   // Geometric example: area calculation for a square
   ////////////////////////////////////////////////////////////////////////

   // Length constant
   const bu::quantity<si::length> l = 1.5 * si::meter;

   // Again an ugly static_cast is needed to bind pow<2> to the first
   // function argument.
   boost::function<bu::quantity<si::area> (bu::quantity<si::length>) >
       A = bl::bind(static_cast<bu::quantity<si::area> (*)(const bu::quantity<si::length>&)>(bu::pow<2>),
                    bl::_1);

   cout << "Area of a square:" << endl
        << "A(" << l <<") = " << A(l) << endl << endl;

   ////////////////////////////////////////////////////////////////////////
   // Thermal example: temperature difference of two absolute temperatures
   ////////////////////////////////////////////////////////////////////////

   // Absolute temperature constants
   const bu::quantity<bu::absolute<si::temperature> >
       Tref = 273.15 * bu::absolute<si::temperature>();
   const bu::quantity<bu::absolute<si::temperature> >
       Tamb = 300.00 * bu::absolute<si::temperature>();

   boost::function<bu::quantity<si::temperature> (bu::quantity<bu::absolute<si::temperature> >,
                                                  bu::quantity<bu::absolute<si::temperature> >)>
       dT = bl::_2 - bl::_1;

   cout << "Temperature difference of two absolute temperatures:" << endl
        << "dT(" << Tref << ", " << Tamb << ") = " << dT(Tref, Tamb) << endl
        << endl;

   return 0;
}
//]


#include <NTL/mat_RR.h>

#include <NTL/new.h>

NTL_START_IMPL

void add(mat_RR& X, const mat_RR& A, const mat_RR& B)
{
   long n = A.NumRows();
   long m = A.NumCols();

   if (B.NumRows() != n || B.NumCols() != m)
      Error("matrix add: dimension mismatch");

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         add(X(i,j), A(i,j), B(i,j));
}

void sub(mat_RR& X, const mat_RR& A, const mat_RR& B)
{
   long n = A.NumRows();
   long m = A.NumCols();

   if (B.NumRows() != n || B.NumCols() != m)
      Error("matrix sub: dimension mismatch");

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         sub(X(i,j), A(i,j), B(i,j));
}

void mul_aux(mat_RR& X, const mat_RR& A, const mat_RR& B)
{
   long n = A.NumRows();
   long l = A.NumCols();
   long m = B.NumCols();

   if (l != B.NumRows())
      Error("matrix mul: dimension mismatch");

   X.SetDims(n, m);

   long i, j, k;
   RR acc, tmp;

   for (i = 1; i <= n; i++) {
      for (j = 1; j <= m; j++) {
         clear(acc);
         for(k = 1; k <= l; k++) {
            mul(tmp, A(i,k), B(k,j));
            add(acc, acc, tmp);
         }
         X(i,j) = acc;
      }
   }
}

void mul(mat_RR& X, const mat_RR& A, const mat_RR& B)
{
   if (&X == &A || &X == &B) {
      mat_RR tmp;
      mul_aux(tmp, A, B);
      X = tmp;
   }
   else
      mul_aux(X, A, B);
}

static
void mul_aux(vec_RR& x, const mat_RR& A, const vec_RR& b)
{
   long n = A.NumRows();
   long l = A.NumCols();

   if (l != b.length())
      Error("matrix mul: dimension mismatch");

   x.SetLength(n);

   long i, k;
   RR acc, tmp;

   for (i = 1; i <= n; i++) {
      clear(acc);
      for (k = 1; k <= l; k++) {
         mul(tmp, A(i,k), b(k));
         add(acc, acc, tmp);
      }
      x(i) = acc;
   }
}

void mul(vec_RR& x, const mat_RR& A, const vec_RR& b)
{
   if (&b == &x || A.position1(x) != -1) {
      vec_RR tmp;
      mul_aux(tmp, A, b);
      x = tmp;
   }
   else
      mul_aux(x, A, b);
}

static
void mul_aux(vec_RR& x, const vec_RR& a, const mat_RR& B)
{
   long n = B.NumRows();
   long l = B.NumCols();

   if (n != a.length())
      Error("matrix mul: dimension mismatch");

   x.SetLength(l);

   long i, k;
   RR acc, tmp;

   for (i = 1; i <= l; i++) {
      clear(acc);
      for (k = 1; k <= n; k++) {
         mul(tmp, a(k), B(k,i));
         add(acc, acc, tmp);
      }
      x(i) = acc;
   }
}

void mul(vec_RR& x, const vec_RR& a, const mat_RR& B)
{
   if (&a == &x) {
      vec_RR tmp;
      mul_aux(tmp, a, B);
      x = tmp;
   }
   else
      mul_aux(x, a, B);
}

void ident(mat_RR& X, long n)
{
   X.SetDims(n, n);
   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i == j)
            set(X(i, j));
         else
            clear(X(i, j));
}

void determinant(RR& d, const mat_RR& M_in)
{
   long k, n;
   long i, j;
   long pos;
   RR t1, t2;
   RR *x, *y;

   n = M_in.NumRows();

   if (M_in.NumCols() != n)
      Error("determinant: nonsquare matrix");

   if (n == 0) {
      set(d);
      return;
   }

   mat_RR M;

   M = M_in;

   RR det;
   set(det);

   RR maxval;

   for (k = 0; k < n; k++) {
      pos = -1;
      clear(maxval);
      for (i = k; i < n; i++) {
         abs(t1, M[i][k]);
         if (t1 > maxval) {
            pos = i;
            maxval = t1;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            negate(det, det);
         }

         mul(det, det, M[k][k]);

         // make M[k, k] == -1

         inv(t1, M[k][k]);
         negate(t1, t1);
         for (j = k+1; j < n; j++) {
            mul(M[k][j], M[k][j], t1);
         }

         for (i = k+1; i < n; i++) {
            // M[i] = M[i] + M[k]*M[i,k]

            t1 = M[i][k];

            x = M[i].elts() + (k+1);
            y = M[k].elts() + (k+1);

            for (j = k+1; j < n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }
      }
      else {
         clear(d);
         return;
      }
   }

   d = det;
}

RR determinant(const mat_RR& a)
   { RR x; determinant(x, a); NTL_OPT_RETURN(RR, x); }

long IsIdent(const mat_RR& A, long n)
{
   if (A.NumRows() != n || A.NumCols() != n)
      return 0;

   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i != j) {
            if (!IsZero(A(i, j))) return 0;
         }
         else {
            if (!IsOne(A(i, j))) return 0;
         }

   return 1;
}

void transpose(mat_RR& X, const mat_RR& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   long i, j;

   if (&X == & A) {
      if (n == m)
         for (i = 1; i <= n; i++)
            for (j = i+1; j <= n; j++)
               swap(X(i, j), X(j, i));
      else {
         mat_RR tmp;
         tmp.SetDims(m, n);
         for (i = 1; i <= n; i++)
            for (j = 1; j <= m; j++)
               tmp(j, i) = A(i, j);
         X.kill();
         X = tmp;
      }
   }
   else {
      X.SetDims(m, n);
      for (i = 1; i <= n; i++)
         for (j = 1; j <= m; j++)
            X(j, i) = A(i, j);
   }
}

void solve(RR& d, vec_RR& X,
           const mat_RR& A, const vec_RR& b)

{
   long n = A.NumRows();
   if (A.NumCols() != n)
      Error("solve: nonsquare matrix");

   if (b.length() != n)
      Error("solve: dimension mismatch");

   if (n == 0) {
      set(d);
      X.SetLength(0);
      return;
   }

   long i, j, k, pos;
   RR t1, t2;
   RR *x, *y;

   mat_RR M;
   M.SetDims(n, n+1);

   for (i = 0; i < n; i++) {
      for (j = 0; j < n; j++)
         M[i][j] = A[j][i];
      M[i][n] = b[i];
   }

   RR det;
   set(det);

   RR maxval;

   for (k = 0; k < n; k++) {
      pos = -1;
      clear(maxval);
      for (i = k; i < n; i++) {
         abs(t1, M[i][k]);
         if (t1 > maxval) {
            pos = i;
            maxval = t1;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            negate(det, det);
         }

         mul(det, det, M[k][k]);

         // make M[k, k] == -1

         inv(t1, M[k][k]);
         negate(t1, t1);
         for (j = k+1; j <= n; j++) {
            mul(M[k][j], M[k][j], t1);
         }

         for (i = k+1; i < n; i++) {
            // M[i] = M[i] + M[k]*M[i,k]

            t1 = M[i][k];

            x = M[i].elts() + (k+1);
            y = M[k].elts() + (k+1);

            for (j = k+1; j <= n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }
      }
      else {
         clear(d);
         return;
      }
   }

   X.SetLength(n);
   for (i = n-1; i >= 0; i--) {
      clear(t1);
      for (j = i+1; j < n; j++) {
         mul(t2, X[j], M[i][j]);
         add(t1, t1, t2);
      }
      sub(t1, t1, M[i][n]);
      X[i] = t1;
   }

   d = det;
}

void inv(RR& d, mat_RR& X, const mat_RR& A)
{
   long n = A.NumRows();
   if (A.NumCols() != n)
      Error("inv: nonsquare matrix");

   if (n == 0) {
      set(d);
      X.SetDims(0, 0);
      return;
   }

   long i, j, k, pos;
   RR t1, t2;
   RR *x, *y;

   mat_RR M;
   M.SetDims(n, 2*n);

   for (i = 0; i < n; i++) {
      for (j = 0; j < n; j++) {
         M[i][j] = A[i][j];
         clear(M[i][n+j]);
      }
      set(M[i][n+i]);
   }

   RR det;
   set(det);

   RR maxval;

   for (k = 0; k < n; k++) {
      pos = -1;
      clear(maxval);
      for (i = k; i < n; i++) {
         abs(t1, M[i][k]);
         if (t1 > maxval) {
            pos = i;
            maxval = t1;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            negate(det, det);
         }

         mul(det, det, M[k][k]);

         // make M[k, k] == -1

         inv(t1, M[k][k]);
         negate(t1, t1);
         for (j = k+1; j < 2*n; j++) {
            mul(M[k][j], M[k][j], t1);
         }

         for (i = k+1; i < n; i++) {
            // M[i] = M[i] + M[k]*M[i,k]

            t1 = M[i][k];

            x = M[i].elts() + (k+1);
            y = M[k].elts() + (k+1);

            for (j = k+1; j < 2*n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }
      }
      else {
         clear(d);
         return;
      }
   }

   X.SetDims(n, n);
   for (k = 0; k < n; k++) {
      for (i = n-1; i >= 0; i--) {
         clear(t1);
         for (j = i+1; j < n; j++) {
            mul(t2, X[j][k], M[i][j]);
            add(t1, t1, t2);
         }
         sub(t1, t1, M[i][n+k]);
         X[i][k] = t1;
      }
   }

   d = det;
}

void mul(mat_RR& X, const mat_RR& A, const RR& b_in)
{
   RR b = b_in;
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 0; i < n; i++)
      for (j = 0; j < m; j++)
         mul(X[i][j], A[i][j], b);
}

void mul(mat_RR& X, const mat_RR& A, double b_in)
{
   static RR b;
   b = b_in;
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 0; i < n; i++)
      for (j = 0; j < m; j++)
         mul(X[i][j], A[i][j], b);
}

void diag(mat_RR& X, long n, const RR& d_in)
{
   RR d = d_in;
   X.SetDims(n, n);
   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i == j)
            X(i, j) = d;
         else
            clear(X(i, j));
}

long IsDiag(const mat_RR& A, long n, const RR& d)
{
   if (A.NumRows() != n || A.NumCols() != n)
      return 0;

   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i != j) {
            if (!IsZero(A(i, j))) return 0;
         }
         else {
            if (A(i, j) != d) return 0;
         }

   return 1;
}

void negate(mat_RR& X, const mat_RR& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         negate(X(i,j), A(i,j));
}

long IsZero(const mat_RR& a)
{
   long n = a.NumRows();
   long i;

   for (i = 0; i < n; i++)
      if (!IsZero(a[i]))
         return 0;

   return 1;
}

void clear(mat_RR& x)
{
   long n = x.NumRows();
   long i;
   for (i = 0; i < n; i++)
      clear(x[i]);
}

mat_RR operator+(const mat_RR& a, const mat_RR& b)
{
   mat_RR res;
   add(res, a, b);
   NTL_OPT_RETURN(mat_RR, res);
}

mat_RR operator*(const mat_RR& a, const mat_RR& b)
{
   mat_RR res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(mat_RR, res);
}

mat_RR operator-(const mat_RR& a, const mat_RR& b)
{
   mat_RR res;
   sub(res, a, b);
   NTL_OPT_RETURN(mat_RR, res);
}

mat_RR operator-(const mat_RR& a)
{
   mat_RR res;
   negate(res, a);
   NTL_OPT_RETURN(mat_RR, res);
}

vec_RR operator*(const mat_RR& a, const vec_RR& b)
{
   vec_RR res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_RR, res);
}

vec_RR operator*(const vec_RR& a, const mat_RR& b)
{
   vec_RR res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_RR, res);
}

void inv(mat_RR& X, const mat_RR& A)
{
   RR d;
   inv(d, X, A);
   if (d == 0) Error("inv: non-invertible matrix");
}

void power(mat_RR& X, const mat_RR& A, const ZZ& e)
{
   if (A.NumRows() != A.NumCols()) Error("power: non-square matrix");

   if (e == 0) {
      ident(X, A.NumRows());
      return;
   }

   mat_RR T1, T2;
   long i, k;

   k = NumBits(e);
   T1 = A;

   for (i = k-2; i >= 0; i--) {
      sqr(T2, T1);
      if (bit(e, i))
         mul(T1, T2, A);
      else
         T1 = T2;
   }

   if (e < 0)
      inv(X, T1);
   else
      X = T1;
}

NTL_END_IMPL


//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//

#include <boost/test/unit_test.hpp>
#include "ParserFlatbuffersSerializeFixture.hpp"
#include "../Deserializer.hpp"

#include <string>
#include <iostream>

BOOST_AUTO_TEST_SUITE(Deserializer)

struct RsqrtFixture : public ParserFlatbuffersSerializeFixture
{
    explicit RsqrtFixture(const std::string & inputShape,
                          const std::string & outputShape,
                          const std::string & dataType)
    {
        m_JsonString = R"(
        {
                inputIds: [0],
                outputIds: [2],
                layers: [
                {
                    layer_type: "InputLayer",
                    layer: {
                          base: {
                                layerBindingId: 0,
                                base: {
                                    index: 0,
                                    layerName: "InputLayer",
                                    layerType: "Input",
                                    inputSlots: [{
                                        index: 0,
                                        connection: {sourceLayerIndex:0, outputSlotIndex:0 },
                                    }],
                                    outputSlots: [ {
                                        index: 0,
                                        tensorInfo: {
                                            dimensions: )" + inputShape + R"(,
                                            dataType: )" + dataType + R"(
                                        },
                                    }],
                                 },}},
                },
                {
                layer_type: "RsqrtLayer",
                layer : {
                        base: {
                             index:1,
                             layerName: "RsqrtLayer",
                             layerType: "Rsqrt",
                             inputSlots: [
                                            {
                                             index: 0,
                                             connection: {sourceLayerIndex:0, outputSlotIndex:0 },
                                            }
                             ],
                             outputSlots: [ {
                                 index: 0,
                                 tensorInfo: {
                                     dimensions: )" + outputShape + R"(,
                                     dataType: )" + dataType + R"(
                                 },
                             }],
                            }},
                },
                {
                layer_type: "OutputLayer",
                layer: {
                        base:{
                              layerBindingId: 0,
                              base: {
                                    index: 2,
                                    layerName: "OutputLayer",
                                    layerType: "Output",
                                    inputSlots: [{
                                        index: 0,
                                        connection: {sourceLayerIndex:1, outputSlotIndex:0 },
                                    }],
                                    outputSlots: [ {
                                        index: 0,
                                        tensorInfo: {
                                            dimensions: )" + outputShape + R"(,
                                            dataType: )" + dataType + R"(
                                        },
                                }],
                            }}},
                }]
         }
        )";
        Setup();
    }
};

struct Rsqrt2dFixture : RsqrtFixture
{
    Rsqrt2dFixture() : RsqrtFixture("[ 2, 2 ]",
                                    "[ 2, 2 ]",
                                    "Float32") {}
};

BOOST_FIXTURE_TEST_CASE(Rsqrt2d, Rsqrt2dFixture)
{
  RunTest<2, armnn::DataType::Float32>(
      0,
      {{"InputLayer", { 1.0f,  4.0f,
                        16.0f, 25.0f }}},
      {{"OutputLayer",{ 1.0f,  0.5f,
                        0.25f, 0.2f }}});
}

BOOST_AUTO_TEST_SUITE_END()


#include <random>
#include <cmath>
#include <cassert>

#include <omp.h>

#include <NTL/ZZ.h>
#include <NTL/tools.h>

#include "gsw.hpp"

using namespace std;
using namespace NTL;

GSW::GSW() {
    GSW(80, 1);
}

GSW::GSW(const int kappa, const int L) {
    // Search for suitable parameters:
    // n >= log(q/sigma)(k+110)/7.2
    // q/sigma6 > 8(N + 1)^L
    BigInt lower_bound;
    n = (kappa+110)/7.2;
    quotient = 4;
    l = floor(log(quotient)/log(2)) + 1;
    N = (n + 1) * l;
    while (true) {
        power(lower_bound, N+1,L);
        lower_bound *= 8 * sigma6;
        if (quotient <= lower_bound) {
            NextPrime(quotient, lower_bound);
        } else {
            break;
        }
        n = log(quotient/ceil(sigma))*(kappa+110)/(7.2*log(2));
        l = floor(log(quotient)/log(2)) + 1;
        N = (n + 1) * l;
    }

    n_1 = n+1;
    m = ceil(n * log(quotient)/log(2));

    gaussSampler = new GaussSampler(sigma);

    omp_set_num_threads(4);
}

GSW::~GSW() {
    delete gaussSampler;
}

BIVector GSW::secret_key_gen() const {
    BIVector secret_key(n+1);
    // sample uniformly
    for (size_t i = 1; i < secret_key.size(); i++) {
        secret_key[i] = RandomBnd(quotient);
    }
    secret_key[0] = 1;
    return secret_key;
}

BIMatrix GSW::public_key_gen(const BIVector& sk) const {
    // recovering t from sk, defined as t = (-s_2,...,-s_n) in Z_q
    BIVector t(n);
    for (unsigned int i = 0; i < n; i++) {
        t[i] = quotient - sk[i+1];
    }

    // Uniformaly generated matrix (part of pk)
    BIMatrix B(m * n);
    for (unsigned int i = 0; i < m*n; i++)
        B[i] = RandomBnd(quotient);

    // First column of public key  b = B*t + e
    BIVector b(m);
    //BIVector e(m); // Error vector, optimized out
    BigInt temp;
    for (unsigned int i = 0; i < m; i++) {
        for (unsigned int j = 0; j < n; j++) {
            MulMod(temp, B[i*n+j], t[j], quotient);
            AddMod(b[i], b[i], temp, quotient);
            //b[i] += B[i*n + j] * t[j];
            //b[i] = b[i] % quotient;
        }
        int bit = gaussSampler->sample() % sigma6;
        //e[i] = bit;
        b[i] += bit;
    }

    // Observe that pk * sk = e
    BIMatrix pk(m * n_1);
    for (unsigned int i = 0; i < m; i++) {
        pk[i*n_1] = b[i];
    }
    for (unsigned int i = 0; i < m*n; i++) {
        pk[1+i+(i/n)] = B[i];
    }

    // this is satisfied, the check proves it
#define DEBUG
#ifdef DEBUG
    BIVector e_1(m);
    for (unsigned int i = 0; i < m; i++) {
        e_1[i] = 0;
        for (unsigned int j = 0; j < n_1; j++) {
            e_1[i] += pk[i*n_1 + j] * sk[j];
            e_1[i] = e_1[i] % quotient;
        }
        //assert(e_1[i] == e[i]);
    }
#endif

    return pk;
}

BitMatrix GSW::encrypt(const BIMatrix& public_key, const BigInt& message) const {
    bernoulli_distribution bernoulli(0.5);

    BitMatrix R(N * m);
    for (size_t i = 0; i < R.size(); i++) {
        R[i] = bernoulli(generator);
    }
    BIMatrix RA(N * n_1);
    BigInt temp;
# pragma omp parallel for shared (R, public_key, RA) schedule(guided)
    for (unsigned int i = 0; i < N; i++) {
        if (omp_get_thread_num() == 0)
            cerr << "Calc RA matrix " << i << " out of " << N << "\r";
        for (unsigned int j = 0; j < n_1; j++) {
            RA[i*n_1 + j] = 0;
            for (unsigned int k = 0; k < m; k++) {
                MulMod(temp, R[i*m + k], public_key[k*n_1 + j], quotient);
                AddMod(RA[i*n_1 + j], RA[i*n_1 + j], temp, quotient);
                //RA[i*n_1 + j] += R[i*m + k] * public_key[k*n_1 + j];
                //RA[i*n_1 + j] = RA[i*n_1 + j] % quotient;
            }
        }
    }
    cerr << endl;
    const BitMatrix RAbits = bit_decomp(RA);
    BIMatrix C(N * N);
# pragma omp parallel for shared (C, message) schedule(guided)
    for (unsigned int i = 0; i < N; i++) {
        if (omp_get_thread_num() == 0)
            cerr << "Calc ciphertext matrix " << i << " out of " << N << "\r";
        for (unsigned int j = 0; j < N; j++) {
            C[i*N + j] = RAbits[i*N + j];
            // message * identity
            if (i == j) {
                C[i*N + j] += message;
            }
            //C[i*N + j] = C[i*N + j] % quotient;
        }
    }

    cerr << endl << "Now to flatten" << endl;

    return flatten(C);
}

BigInt GSW::decrypt(const BIVector& sk, const BitMatrix& C) const {
    BigInt m, it, fract;
    const auto v = powers_of_2(sk);
    BIVector powered_m_bits(l-1);
    for (unsigned int i = 0; i < l-1; i++) {
        powered_m_bits[i] = 0;
        for (unsigned int j = 0; j < N; j++) {
            powered_m_bits[i] += C[i*(l-1) + j] * v[j];
            powered_m_bits[i] = powered_m_bits[i] % quotient;
        }
    }
    m = 0;
    for (int i = l-2; i >= 0; i--) {
        it = (powered_m_bits[i] - pow(2, i)*m);
        fract = it % quotient/2;
        bool bit = fract >= quotient/4;

        m += bit << (l-2 - i);
    }

    return m;
}

bool GSW::decrypt_bit(const BIVector& sk, const BitMatrix& C) const {
    unsigned int i;
    const auto v = powers_of_2(sk);
    BigInt q_4, q_2; q_4 = quotient/4; q_2 = quotient/2;

    for(i = 0; i < l; i++) {
        if(v[i] > q_4 && v[i] <= q_2) break;
    }

    BigInt xi, temp;
    xi = 0;
    for (unsigned int j = 0; j < N; j++) {
        MulMod(temp, C[i*N + j], v[j], quotient);
        AddMod(xi, xi, temp, quotient);
        //xi += (C[i*N + j] * v[j]);
        //xi = xi % quotient;
    }

    return xi >= v[i]/2;
}

BitMatrix GSW::nand(const BitMatrix& a, const BitMatrix& b) const {
    BIMatrix res(a.size());

    for (unsigned int i = 0; i < N; i++) {
        res[i*N + i] = 1;
    }

    BigInt temp;
# pragma omp parallel for shared (a, b, res) schedule(guided)
    for (unsigned int i = 0; i < N; i++) {
        if (omp_get_thread_num() == 0)
            cerr << "Performing a NAND, hold on " << i << " out of " << N << "\r";
        for (unsigned int j = 0; j < N; j++) {
            for (unsigned int k = 0; k < N; k++) {
                res[i*N + j] = res[i*N + j] + (a[i*N + k] & b[k*N + j]);
            }
            res[i*N + j] = res[i*N + j] % quotient;
        }
    }
    cerr << endl;

    return flatten(res);
}

//////////////////////////////////////////////
// Utility Functions
//////////////////////////////////////////////

BIVector GSW::powers_of_2(const BIVector& a) const {
    BIVector result(N);
# pragma omp parallel for shared (result, a) schedule(guided)
    for (unsigned int i = 0; i < n+1; i++) {
        for (unsigned int j = 0; j < l; j++) {
            MulMod(result[i*l + j], power2_ZZ(j), a[i], quotient);
            //result[i*l + j] = (BigInt) pow(2, j)*a[i] % quotient;
        }
    }

    return result;
}

BitVector GSW::bit_decomp(const BIVector& a) const {
    unsigned int num_rows = a.size() / n_1;
    BitVector result(n_1*l*num_rows);
    BigInt mask;
    mask = 1;
# pragma omp parallel for shared (result, a) schedule(guided)
    for (unsigned int j = 0; j < l; j++) {
        for (unsigned int k = 0; k < num_rows; k++) {
            for (unsigned int i = 0; i < n_1; i++) {
                result[k*l*n_1 + i*l + j] = (a[k*n_1 + i] & mask) > 0;
            }
        }
        mask <<= 1;
    }

    return result;
}

BIVector GSW::inverse_bit_decomp(const BitVector& a) const {
    BIVector result(n_1 * N);
    BigInt multiplier;
    multiplier = 1;

# pragma omp parallel for shared (result, a) schedule(guided)
    for (unsigned int j = 0; j < l; j++) {
        for (unsigned int row = 0; row < N; row++) {
            for (unsigned int i = 0; i < n_1; i++) {
                result[row * n_1 + i] += a[row*n_1*l + i*l + j] * multiplier;
            }
        }
        multiplier <<= 1;
    }

    return result;
}

BIVector GSW::inverse_bit_decomp(const BIVector& a) const {
    BIVector result(n_1 * N);
    BigInt multiplier;
    multiplier = 1;

# pragma omp parallel for shared (result, a) schedule(guided)
    for (unsigned int j = 0; j < l; j++) {
        for (unsigned int row = 0; row < N; row++) {
            for (unsigned int i = 0; i < n_1; i++) {
                result[row * n_1 + i] += a[row*n_1*l + i*l + j] * multiplier;
            }
        }
        multiplier <<= 1;
    }

    return result;
}

BitVector GSW::flatten(const BitVector& a) const {
    return bit_decomp(inverse_bit_decomp(a));
}

BitVector GSW::flatten(const BIVector& a) const {
    return bit_decomp(inverse_bit_decomp(a));
}