/* ---------------------------------------------------------------------
 *
 * Copyright (C) 1999 - 2021 by the deal.II authors
 *
 * This file is part of the deal.II library.
 *
 * The deal.II library is free software; you can use it, redistribute
 * it, and/or modify it under the terms of the GNU Lesser General
 * Public License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * The full text of the license can be found in the file LICENSE.md at
 * the top level directory of deal.II.
 *
 * ---------------------------------------------------------------------
 *
 *

 *
 *
 * This program was contributed by Peter Munch. This work and the required
 * generalizations of the internal data structures of deal.II form part of the
 * project "Virtual Materials Design" funded by the Helmholtz Association of
 * German Research Centres.
 *
 *
 *
 *
 *
Introduction

 *
 *
Motivation

 *
 * The motivation for using simplex meshes (as done in step-3simplex) is
 * straightforward: many freely available mesh-generation tools are very good in
 * creating good-quality meshes in such a format, while they struggle with
 * hex-only meshes. Hex-only meshes, on the other hand, are characterized with
 * better numerical properties (e.g., less degrees of freedom for the same
 * degree of accuracy and possibly better performance since the tensor-product
 * structure can be exploited) and are, as a consequence, the natural choice for
 * rather simple geometries and for meshes described by a coarse mesh with a few
 * cells (like a hyperball) and obtained in their final form through iterative
 * local refinement.
 *
 * Mixed meshes try to combine the best of both worlds by partitioning the
 * geometry in parts that can be easily meshed by hypercube cells
 * (quadrilaterals in 2D, hexahedrons in 3D) and in parts that can not be meshed
 * easily, requiring simplices (triangles in 2D, tetrahedrons in 3D). Since one
 * assumes that the region requiring simplices is rather small compared to the
 * rest of the domain where more efficient and accurate methods can be applied,
 * one can expect that the overall efficiency is hardly impacted by such an
 * approach.
 *
 * One should note that in 3D, one also needs a transition region between
 * hypercube and simplex regions. Here, one can use wedges/prisms and/or
 * pyramids.
 *
 *
 *
Working with mixed meshes

 *
 *
 * In the following, we concentrate, for the sake of simplicity, on 2D meshes:
 * they can only contain triangles and quadrilaterals. However, as detailed in
 * the outlook, an extension of the presented approach to the 3D case is
 * straightforward.
 *
 *
 * The complexity of working with mixed meshes in 2D results from the fact
 * that it contains of two
 * types of geometrical objects: quadrilaterals and triangles. How to deal with
 * quadrilaterals, we have discussed in step-3: we selected an appropriate
 * finite element, quadrature rule and mapping object, e.g., FE_Q, QGauss, and
 * MappingFE (initialized with FE_Q). For simplex meshes, we selected in
 * step-3simplex FE_SimplexP, QGaussSimplex, and MappingFE (intialized with
 * FE_SimplexP).
 *
 * For mixed meshes, we need multiple finite elements, quadrature rules, and
 * mapping objects (one set for triangles and one set for quadrilaterals) in the
 * same program. To ease the work with the multitude of objects (in particular
 * in 3D, we need at least four of each), you can collect the objects and group
 * them together in hp::FECollection, hp::QCollection, and
 * hp::MappingCollection.
 *
 * Just like in the context of finite elements, quadrature rules, and mapping
 * objects, we need multiple FEValues objects: the collection of FEValues is
 * called hp::FEValues. It returns the FEValues object needed for the current
 * cell via the method hp::FEValues::get_present_fe_values().
 *
 * For hp::FEValues, to be able to select the right finite element/quadrature
 * rule/ mapping object set, it queries the active_fe_index of the given cell
 * during hp::FEValues::reinit(). The indices have to be set - as shown below -
 * before calling DoFHandler::distribute_dofs() by the user.
 *
 *
 * The namespace name of hp::FECollection, hp::QCollection,
 * hp::MappingCollection, and hp::FEValues indicates that these classes have not
 * been written for mixed meshes in the first place, but for problems where each
 * (hypercube) cell could have a different type of finite element assigned - in
 * the simplest case, all cells have the same element type but different
 * polynomial degrees p (the reason for the letter "p" in "hp"). An extension of
 * this infrastructure to work not only on different element types but also on
 * different geometrical objects was a natural choice. For further details on
 * hp-methods, see step-27.
 *
 *
 *
Mesh generation

 *
 * Just like in step-3simplex, we read an externally generated mesh. For this
 * tutorial, we have created the mesh (square with width and height of one;
 * quadrilaterals on the left half and triangles on the right half) with Gmsh
 * with the following journal file "box_2D_mixed.geo":
 *
 * @code
 * Rectangle(1) = {0.0, 0, 0, 0.5, 1, 0};
 * Rectangle(2) = {0.5, 0, 0, 0.5, 1, 0};
 * Recombine Surface{1};
 * Physical Surface("All") = {1, 2};
 * Mesh 2;
 * Coherence Mesh;
 * Save "box_2D_mixed.msh";
 * @endcode
 *
 * The journal file can be processed by Gmsh generating the actual mesh with the
 * ending ".msh":
 *
 * @code
 * gmsh box_2D_mixed.geo
 * @endcode
 *
 * We have included in the tutorial folder both the journal file and the mesh
 * file in the event that one does not have access to Gmsh.
 *
 */

// @sect3{Include files}

// Include files, as used in step-3:
#include <deal.II/base/function.h>

#include <deal.II/dofs/dof_handler.h>
#include <deal.II/dofs/dof_tools.h>

#include <deal.II/fe/fe_q.h>
#include <deal.II/fe/fe_values.h>

#include <deal.II/grid/tria.h>

#include <deal.II/lac/dynamic_sparsity_pattern.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/precondition.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/sparse_matrix.h>
#include <deal.II/lac/vector.h>

#include <deal.II/numerics/data_out.h>
#include <deal.II/numerics/matrix_tools.h>
#include <deal.II/numerics/vector_tools.h>

#include <fstream>
#include <iostream>

// Include files, as added in step-3simplex:
#include <deal.II/base/quadrature_lib.h>

#include <deal.II/fe/fe_simplex_p.h>
#include <deal.II/fe/mapping_fe.h>

#include <deal.II/grid/grid_in.h>

// Include files that we need in this tutorial to be able to deal with
// collections of finite elements, quadrature rules, mapping objects, and
// FEValues.
#include <deal.II/hp/fe_collection.h>
#include <deal.II/hp/fe_values.h>
#include <deal.II/hp/mapping_collection.h>
#include <deal.II/hp/q_collection.h>

using namespace dealii;

// @sect3{The Step3 class}
//
// This is the main class of the tutorial. Since it is very similar to the
// version from step-3 and step-3simplex, we will only point out and explain
// the relevant differences that allow to perform simulations on mixed meshes.

class Step3
{
public:
  Step3();

  void run();

private:
  void make_grid();
  void setup_system();
  void assemble_system();
  void solve();
  void output_results() const;

  Triangulation<2> triangulation;

  // As already explained, we are not working with mapping objects, finite
  // elements, and quadrature rules directly but with collections of them.
  const hp::MappingCollection<2> mapping;
  const hp::FECollection<2>      fe;
  const hp::QCollection<2>       quadrature_formula;

  DoFHandler<2> dof_handler;

  SparsityPattern      sparsity_pattern;
  SparseMatrix<double> system_matrix;

  Vector<double> solution;
  Vector<double> system_rhs;
};

// @sect4{Step3::Step3}
//
// In the constructor of the Step3 class, we fill the collections. Here, we
// position the objects related to triangles in the first place (index 0) and
// the ones related to quadrilaterals in the second place (index 1).
Step3::Step3()
  : mapping(MappingFE<2>(FE_SimplexP<2>(1)), MappingFE<2>(FE_Q<2>(1)))
  , fe(FE_SimplexP<2>(2), FE_Q<2>(2))
  , quadrature_formula(QGaussSimplex<2>(3), QGauss<2>(3))
  , dof_handler(triangulation)
{}

// @sect4{Step3::make_grid}
//
// Read the external mesh file "box_2D_mixed.msh" as in step-3simplex.
void Step3::make_grid()
{
  GridIn<2>(triangulation).read("box_2D_mixed.msh");

  std::cout << "Number of active cells: " << triangulation.n_active_cells()
            << std::endl;
}

// @sect4{Step3::setup_system}
//
// In contrast to step-3 and step-3simplex, we need here a preprocessing step
// that assigns an active_fe_index to each cell consistently according to the
// indices in the collections and the cell type.
void Step3::setup_system()
{
  for (const auto &cell : dof_handler.active_cell_iterators())
    {
      if (cell->reference_cell() == ReferenceCells::Triangle)
        cell->set_active_fe_index(0);
      else if (cell->reference_cell() == ReferenceCells::Quadrilateral)
        cell->set_active_fe_index(1);
      else
        Assert(false, ExcNotImplemented());
    }

  dof_handler.distribute_dofs(fe);
  std::cout << "Number of degrees of freedom: " << dof_handler.n_dofs()
            << std::endl;
  DynamicSparsityPattern dsp(dof_handler.n_dofs());
  DoFTools::make_sparsity_pattern(dof_handler, dsp);
  sparsity_pattern.copy_from(dsp);

  system_matrix.reinit(sparsity_pattern);

  solution.reinit(dof_handler.n_dofs());
  system_rhs.reinit(dof_handler.n_dofs());
}

// @sect4{Step3::assemble_system}
//
// The following function looks similar to the version in step-3 and
// step-3simplex with the following two differences:
//  - We do not work with FEValues directly but with the collection class
//    hp::FEValues. It gives us - after it has been initialized with the current
//    cell - a reference to the right FEValues object (constructed
//    with the correct mapping object, finite element, and quadrature rule),
//    which can be used as usual to compute the cell integrals.
//  - The cell-local stiffness matrix and the right-hand-side vector have
//    different sizes depending on the cell type (6 DoFs vs. 9 DoFs) so that
//    they might need to be resized for each cell.
//
// Apart from these two changes, the code has not changed. Not even, the
// cell integrals have been changed depending on whether one operates on
// hypercube, simplex, or mixed meshes.
void Step3::assemble_system()
{
  hp::FEValues<2> hp_fe_values(mapping,
                               fe,
                               quadrature_formula,
                               update_values | update_gradients |
                                 update_JxW_values);

  FullMatrix<double>                   cell_matrix;
  Vector<double>                       cell_rhs;
  std::vector<types::global_dof_index> local_dof_indices;

  for (const auto &cell : dof_handler.active_cell_iterators())
    {
      hp_fe_values.reinit(cell);

      const auto &fe_values = hp_fe_values.get_present_fe_values();

      const unsigned int dofs_per_cell = cell->get_fe().n_dofs_per_cell();
      cell_matrix.reinit(dofs_per_cell, dofs_per_cell);
      cell_rhs.reinit(dofs_per_cell);
      local_dof_indices.resize(dofs_per_cell);

      cell_matrix = 0;
      cell_rhs    = 0;

      for (const unsigned int q_index : fe_values.quadrature_point_indices())
        {
          for (const unsigned int i : fe_values.dof_indices())
            for (const unsigned int j : fe_values.dof_indices())
              cell_matrix(i, j) +=
                (fe_values.shape_grad(i, q_index) * // grad phi_i(x_q)
                 fe_values.shape_grad(j, q_index) * // grad phi_j(x_q)
                 fe_values.JxW(q_index));           // dx

          for (const unsigned int i : fe_values.dof_indices())
            cell_rhs(i) += (fe_values.shape_value(i, q_index) * // phi_i(x_q)
                            1. *                                // f(x_q)
                            fe_values.JxW(q_index));            // dx
        }
      cell->get_dof_indices(local_dof_indices);

      for (const unsigned int i : fe_values.dof_indices())
        for (const unsigned int j : fe_values.dof_indices())
          system_matrix.add(local_dof_indices[i],
                            local_dof_indices[j],
                            cell_matrix(i, j));

      for (const unsigned int i : fe_values.dof_indices())
        system_rhs(local_dof_indices[i]) += cell_rhs(i);
    }

  std::map<types::global_dof_index, double> boundary_values;
  VectorTools::interpolate_boundary_values(
    mapping, dof_handler, 0, Functions::ZeroFunction<2>(), boundary_values);
  MatrixTools::apply_boundary_values(boundary_values,
                                     system_matrix,
                                     solution,
                                     system_rhs);
}

// @sect4{Step3::solve}
//
// Nothing has changed here.
void Step3::solve()
{
  SolverControl            solver_control(1000, 1e-12);
  SolverCG<Vector<double>> solver(solver_control);
  solver.solve(system_matrix, solution, system_rhs, PreconditionIdentity());
}

// @sect4{Step3::output_results}
//
// Nothing has changed here.
void Step3::output_results() const
{
  DataOut<2> data_out;

  DataOutBase::VtkFlags flags;
  flags.write_higher_order_cells = true;
  data_out.set_flags(flags);

  data_out.attach_dof_handler(dof_handler);
  data_out.add_data_vector(solution, "solution");
  data_out.build_patches(mapping, 2);
  std::ofstream output("solution.vtk");
  data_out.write_vtk(output);
}

// @sect4{Step3::run}
//
// Nothing has changed here.
void Step3::run()
{
  make_grid();
  setup_system();
  assemble_system();
  solve();
  output_results();
}

// @sect3{The main function}
//
// Nothing has changed here.
int main()
{
  deallog.depth_console(2);

  Step3 laplace_problem;
  laplace_problem.run();

  return 0;
}

/**
 *
Results

 *
 * The following figures show the mesh and the result obtained by executing this
 * program:
 *
 *

 *

 *
 *         @image html step_3_mixed_0.png
 *
 *
 *         @image html step_3_mixed_1.png
 *
 *

 *

 *
 * Not surprisingly, the result looks as expected.
 *
 *
 *
Possibilities for extensions

 *
 * In this tutorial, we presented how to use the deal.II simplex infrastructure
 * to solve a simple Poisson problem on a mixed mesh in 2D. In this scope, we
 * could only present a small section of the capabilities. In the following, we
 * point out further capabilities briefly.
 *
 *
 *
Pure hypercube and simplex meshes

 *
 * In this tutorial, we worked on a mesh consisting both of quadrilaterals and
 * triangles. However, the program and the underlying concepts also work if the
 * mesh only contains either quadrilaterals or triangles. Interested users can
 * try this out: we have provided appropriate journal files and meshes for such
 * cases.
 *
 *
 *
3D meshes

 *
 * In 3D, meshes might also consist of wedges/prisms and pyramids. Therefore,
 * the above introduced collections might consist of four components.
 *
 * For wedge/prism and pyramid cell types, following finite-element and
 * quadrature-rule classes are available:
 *  - wedge: FE_WedgeP, FE_WedgeDGP, QGaussWedge, MappingFE
 *  - pyramid: FE_PyramidP, FE_PyramidDGP, QGaussPyramid, MappingFE
 *
 *
Parallelization, face integrals, discontinuous Galerkin methods, and
 * matrix-free operator evaluation

 *
 * Regarding these aspects, the same comments are valid that are described in
 * step-3simplex for pure simplex meshes.
 *
 */


/**
 * io.cpp
 *
 * 2021 Gabriel Moreira
 *
 * https://github.com/gabmoreira/maks
 *
 * This software and the related documents  are provided as  is,  with no express
 * or implied  warranties,  other  than those  that are  expressly stated  in the
 * License.
 *
 * Copyright © 2021 Gabriel Moreira. All rights reserved.
 */

#include <iostream>
#include <string>
#include <vector>
#include <fstream>

#include <Eigen/Geometry>

#include "io.hpp"

using std::vector;
using std::string;
using std::stod;
using std::stoi;

using Eigen::Quaterniond;
using Eigen::Vector3d;
using Eigen::VectorXi;
using Eigen::Matrix3d;
using Eigen::MatrixXd;
using Eigen::Ref;

#define G2O_PRECISION 15

/**
 * Splits a string based on a delimiter.
 *
 * @param str Input string to be split.
 * @param delim Input string delimiter.
 * @return vector of tokens.
 */
vector<string> strSplit(const string& str, const string& delim) {
    vector<string> tokens;
    size_t prev = 0, pos = 0;
    do {
        pos = str.find(delim, prev);
        if (pos == string::npos)
            pos = str.length();
        string token = str.substr(prev, pos-prev);
        if (!token.empty())
            tokens.push_back(token);
        prev = pos + delim.length();
    } while (pos < str.length() && prev < str.length());

    return tokens;
};

/**
 * Reads g2o file.
 *
 * @param path (input) char*  - g2o file name.
 * @param node_i (output) Eigen::VectorXi - node indices.
 * @param node_r (output) Eigen::MatrixXd - SO(3) node rotations.
 * @param node_t (output) Eigen::MatrixXd - node translations.
 * @param edge_i (output) Eigen::VectorXi - edge indices (from).
 * @param edge_j (output) Eigen::VectorXi - edge indices (to).
 * @param edge_r (output) Eigen::MatrixXd - SO(3) edge rotations.
 * @param edge_t (output) Eigen::MatrixXd - edge translations.
 * @param num_nodes (output) int - number of nodes loaded.
 * @param num_edges (output) int - number of edges loaded.
 */
void readG2O(const char*   path,
             Ref<VectorXi> node_i,
             Ref<MatrixXd> node_r,
             Ref<MatrixXd> node_t,
             Ref<VectorXi> edge_i,
             Ref<VectorXi> edge_j,
             Ref<MatrixXd> edge_r,
             Ref<MatrixXd> edge_t,
             int&          num_nodes,
             int&          num_edges) {

    std::ifstream infile(path);
    if(!infile)
        return;

    string edge_tag = "EDGE_SE3:QUAT";
    string node_tag = "VERTEX_SE3:QUAT";

    num_nodes = 0;
    num_edges = 0;

    string line;
    while (std::getline(infile, line)) {
        vector<string> tokens = strSplit(line, " ");
        string tag = tokens[0];

        if (0 == tag.compare(edge_tag)) {
            int edge_i_id = stoi(tokens[1]);
            int edge_j_id = stoi(tokens[2]);

            // Edge translation
            Vector3d t( stod(tokens[3]), stod(tokens[4]), stod(tokens[5]) );

            // Edge quaternion
            Quaterniond q( stod(tokens[9]), stod(tokens[6]), stod(tokens[7]), stod(tokens[8]) );

            edge_i(num_edges) += edge_i_id;
            edge_j(num_edges) += edge_j_id;
            edge_r.block<3,3>(0,num_edges*3) += q.normalized().toRotationMatrix();
            edge_t.block<3,1>(0,num_edges) += t;
            num_edges++;

        } else if (0 == tag.compare(node_tag)) {
            int node_id = stoi(tokens[1]);

            // Node translation
            Vector3d t( stod(tokens[2]), stod(tokens[3]), stod(tokens[4]) );

            // Node quaternion
            Quaterniond q( stod(tokens[8]), stod(tokens[5]), stod(tokens[6]), stod(tokens[7]) );

            node_i(num_nodes) += node_id;
            node_r.block<3,3>(0,num_nodes*3) += q.normalized().toRotationMatrix();
            node_t.block<3,1>(0,num_nodes) += t;
            num_nodes++;
        };
    };
};

/**
 * Reads edge rotations from g2o file.
 *
 * @param path (input) char*  - g2o file name.
 * @param edge_i (output) Eigen::VectorXi - edge id (from).
 * @param edge_j (output) Eigen::VectorXi - edge id (to).
 * @param edge_r (output) Eigen::MatrixXd - Row block matrix containing SO(3) rotations.
 * @param num_edges (output) int - number of edges.
 */
void readG2O(const char*          path,
             Eigen::Ref<VectorXi> edge_i,
             Eigen::Ref<VectorXi> edge_j,
             Eigen::Ref<MatrixXd> edge_r,
             int&                 num_edges) {

    std::ifstream infile(path);
    if(!infile)
        return;

    string edge_tag = "EDGE_SE3:QUAT";

    num_edges = 0;

    string line;
    while (std::getline(infile, line)) {
        vector<string> tokens = strSplit(line, " ");
        string tag = tokens[0];

        if (0 == tag.compare(edge_tag)) {
            int edge_i_id = stoi(tokens[1]);
            int edge_j_id = stoi(tokens[2]);

            /* Edge quaternion */
            Quaterniond q( stod(tokens[9]), stod(tokens[6]), stod(tokens[7]), stod(tokens[8]) );

            edge_i(num_edges) += edge_i_id;
            edge_j(num_edges) += edge_j_id;
            edge_r.block<3,3>(0,num_edges*3) += q.normalized().toRotationMatrix();
            num_edges++;
        };
    };
};

void writeG2O(const char* filename,
              const vector<int>&         node_i,
              const vector<Quaterniond>& node_q,
              const vector<Vector3d>&    node_t,
              const vector<int>&         edge_i,
              const vector<int>&         edge_j,
              const vector<Quaterniond>& edge_q,
              const vector<Vector3d>&    edge_t) {

    int num_nodes = (int) node_i.size();
    int num_edges = (int) edge_i.size();

    string filepath = string(filename) + ".g2o";

    std::ostringstream buffer;
    buffer.clear();

    /* Write node data */
    for (int i = 0; i < num_nodes; ++i) {
        /* Write tag */
        buffer << "VERTEX_SE3:QUAT";

        /* Write node id */
        buffer << " " << std::to_string(node_i[i]);

        /* Write node translation */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](0); // x
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](1); // y
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_t[i](2); // z

        /* Write node quaternion */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].x();           // qx
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].y();           // qy
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].z();           // qz
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << node_q[i].w() << "\n";   // qw
    };

    /* Write edge data */
    for (int i = 0; i < num_edges; ++i) {
        /* Write tag */
        buffer << "EDGE_SE3:QUAT";

        /* Write edge ids */
        buffer << " " << std::to_string(edge_i[i]);
        buffer << " " << std::to_string(edge_j[i]);

        /* Write node translation */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](0); // x
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](1); // y
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_t[i](2); // z

        /* Write node quaternion */
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].x();           // qx
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].y();           // qy
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].z();           // qz
        buffer << " " << std::fixed << std::setprecision(G2O_PRECISION) << edge_q[i].w() << "\n";   // qw
    };

    std::ofstream outfile;
    outfile.open(filepath);
    if(!outfile) {
        return;
    };

    outfile << buffer.str();
    outfile.close();
    printf("Pose graph saved to %s\n", filepath.c_str());
};


#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>
#include <boost/multiprecision/cpp_int.hpp>
using namespace std;
using namespace boost::multiprecision;
int main() {
    int n, k; cin >> n >> k;
    vector<long long int> a(n, 0);
    for (int i = 0; i < n; i++) cin >> a[i];
    sort(a.begin(), a.end());
    cpp_int sum = 0, j = 0;
    for (int i = 0; i < k; i++) {
        sum += a[i] + j;
        j++;
    }
    cout << sum << endl;
}


/*
 * rosenbrock4.cpp
 *
 * Copyright 2010-2012 Mario Mulansky
 * Copyright 2011-2012 Karsten Ahnert
 * Copyright 2012 Andreas Angelopoulos
 *
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE_1_0.txt or
 * copy at http://www.boost.org/LICENSE_1_0.txt)
 */

#include <iostream>
#include <fstream>
#include <utility>

#include <boost/numeric/odeint.hpp>

#include <boost/phoenix/core.hpp>

#include <boost/phoenix/core.hpp>
#include <boost/phoenix/operator.hpp>

using namespace std;
using namespace boost::numeric::odeint;
namespace phoenix = boost::phoenix;

//[ stiff_system_definition
typedef boost::numeric::ublas::vector< double > vector_type;
typedef boost::numeric::ublas::matrix< double > matrix_type;

struct stiff_system
{
    void operator()( const vector_type &x , vector_type &dxdt , double /* t */ )
    {
        dxdt[ 0 ] = -101.0 * x[ 0 ] - 100.0 * x[ 1 ];
        dxdt[ 1 ] = x[ 0 ];
    }
};

struct stiff_system_jacobi
{
    void operator()( const vector_type & /* x */ , matrix_type &J , const double & /* t */ , vector_type &dfdt )
    {
        J( 0 , 0 ) = -101.0;
        J( 0 , 1 ) = -100.0;
        J( 1 , 0 ) = 1.0;
        J( 1 , 1 ) = 0.0;
        dfdt[0] = 0.0;
        dfdt[1] = 0.0;
    }
};
//]

/*
//[ stiff_system_alternative_definition
typedef boost::numeric::ublas::vector< double > vector_type;
typedef boost::numeric::ublas::matrix< double > matrix_type;

struct stiff_system
{
    template< class State >
    void operator()( const State &x , State &dxdt , double t )
    {
        ...
    }
};

struct stiff_system_jacobi
{
    template< class State , class Matrix >
    void operator()( const State &x , Matrix &J , const double &t , State &dfdt )
    {
        ...
    }
};
//]
 */

int main( int argc , char **argv )
{
//    typedef rosenbrock4< double > stepper_type;
//    typedef rosenbrock4_controller< stepper_type > controlled_stepper_type;
//    typedef rosenbrock4_dense_output< controlled_stepper_type > dense_output_type;
    //[ integrate_stiff_system
    vector_type x( 2 , 1.0 );

    size_t num_of_steps = integrate_const( make_dense_output< rosenbrock4< double > >( 1.0e-6 , 1.0e-6 ) ,
            make_pair( stiff_system() , stiff_system_jacobi() ) ,
            x , 0.0 , 50.0 , 0.01 ,
            cout << phoenix::arg_names::arg2 << " " << phoenix::arg_names::arg1[0] << "\n" );
    //]
    clog << num_of_steps << endl;

//    typedef runge_kutta_dopri5< vector_type > dopri5_type;
//    typedef controlled_runge_kutta< dopri5_type > controlled_dopri5_type;
//    typedef dense_output_runge_kutta< controlled_dopri5_type > dense_output_dopri5_type;
    //[ integrate_stiff_system_alternative

    vector_type x2( 2 , 1.0 );

    size_t num_of_steps2 = integrate_const( make_dense_output< runge_kutta_dopri5< vector_type > >( 1.0e-6 , 1.0e-6 ) ,
            stiff_system() , x2 , 0.0 , 50.0 , 0.01 ,
            cout << phoenix::arg_names::arg2 << " " << phoenix::arg_names::arg1[0] << "\n" );
    //]
    clog << num_of_steps2 << endl;

    return 0;
}


/**
 * -*- c++ -*-
 *
 * \file const_iterator_type.hpp
 *
 * \brief Const iterator to a given container type.
 *
 * Copyright (c) 2009, Marco Guazzone
 *
 * 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)
 *
 * \author Marco Guazzone, marco.guazzone@gmail.com
 */

#ifndef BOOST_NUMERIC_UBLAS_TRAITS_CONST_ITERATOR_TYPE_HPP
#define BOOST_NUMERIC_UBLAS_TRAITS_CONST_ITERATOR_TYPE_HPP

#include <boost/numeric/ublas/fwd.hpp>
#include <boost/numeric/ublas/tags.hpp>
#include <boost/numeric/ublas/traits.hpp>

namespace boost { namespace numeric { namespace ublas {

    namespace detail {

        /**
         * \brief Auxiliary class for retrieving the const iterator to the given
         *  matrix expression according its orientation and to the given dimension tag.
         * \tparam MatrixT A model of MatrixExpression.
         * \tparam TagT A dimension tag type (e.g., tag::major).
         * \tparam OrientationT An orientation category type (e.g., row_major_tag).
         */
        template <typename MatrixT, typename TagT, typename OrientationT>
        struct const_iterator_type_impl;

        /// \brief Specialization of \c const_iterator_type_impl for row-major oriented
        ///  matrices and over the major dimension.
        template <typename MatrixT>
        struct const_iterator_type_impl<MatrixT,tag::major,row_major_tag>
        {
            typedef typename matrix_view_traits<MatrixT>::const_iterator1 type;
        };

        /// \brief Specialization of \c const_iterator_type_impl for column-major
        ///  oriented matrices and over the major dimension.
        template <typename MatrixT>
        struct const_iterator_type_impl<MatrixT,tag::major,column_major_tag>
        {
            typedef typename matrix_view_traits<MatrixT>::const_iterator2 type;
        };

        /// \brief Specialization of \c const_iterator_type_impl for row-major oriented
        ///  matrices and over the minor dimension.
        template <typename MatrixT>
        struct const_iterator_type_impl<MatrixT,tag::minor,row_major_tag>
        {
            typedef typename matrix_view_traits<MatrixT>::const_iterator2 type;
        };

        /// \brief Specialization of \c const_iterator_type_impl for column-major
        ///  oriented matrices and over the minor dimension.
        template <typename MatrixT>
        struct const_iterator_type_impl<MatrixT,tag::minor,column_major_tag>
        {
            typedef typename matrix_view_traits<MatrixT>::const_iterator1 type;
        };

    } // Namespace detail

    /**
     * \brief A const iterator for the given container type over the given
     *  dimension.
     * \tparam ContainerT A container expression type.
     * \tparam TagT A dimension tag type (e.g., tag::major).
     */
    template <typename ContainerT, typename TagT=void>
    struct const_iterator_type;

    /**
     * \brief Specialization of \c const_iterator_type for vector expressions.
     * \tparam VectorT A model of VectorExpression type.
     */
    template <typename VectorT>
    struct const_iterator_type<VectorT, void>
    {
        typedef typename vector_view_traits<VectorT>::const_iterator type;
    };

    /**
     * \brief Specialization of \c const_iterator_type for matrix expressions and
     *  over the major dimension.
     * \tparam MatrixT A model of MatrixExpression type.
     */
    template <typename MatrixT>
    struct const_iterator_type<MatrixT,tag::major>
    {
        typedef typename detail::const_iterator_type_impl<MatrixT,tag::minor,typename matrix_view_traits<MatrixT>::orientation_category>::type type;
    };

    /**
     * \brief Specialization of \c const_iterator_type for matrix expressions and
     *  over the minor dimension.
     * \tparam MatrixT A model of MatrixExpression type.
     */
    template <typename MatrixT>
    struct const_iterator_type<MatrixT,tag::minor>
    {
        typedef typename detail::const_iterator_type_impl<MatrixT,tag::minor,typename matrix_view_traits<MatrixT>::orientation_category>::type type;
    };

}}} // Namespace boost::numeric::ublas

#endif // BOOST_NUMERIC_UBLAS_TRAITS_CONST_ITERATOR_TYPE_HPP


//
//  Copyright (c) 2018-2020, Cem Bassoy, cem.bassoy@gmail.com
//  Copyright (c) 2019-2020, Amit Singh, amitsingh19975@gmail.com
//
//  Distributed under the Boost Software License, Version 1.0. (See
//  accompanying file LICENSE_1_0.txt or copy at
//  http://www.boost.org/LICENSE_1_0.txt)
//
//  The authors gratefully acknowledge the support of
//  Google and Fraunhofer IOSB, Ettlingen, Germany
//

#include <boost/numeric/ublas/tensor/expression_evaluation.hpp>
#include <boost/numeric/ublas/tensor/expression.hpp>
#include <boost/numeric/ublas/tensor/extents.hpp>
#include <boost/numeric/ublas/tensor/tensor.hpp>
#include "utility.hpp"
#include <boost/test/unit_test.hpp>

#include <functional>

BOOST_AUTO_TEST_SUITE(test_tensor_expression)
using test_types = zip<int,float,std::complex<float>>::with_t<boost::numeric::ublas::layout::first_order, boost::numeric::ublas::layout::last_order>;

struct fixture
{
  using extents_t = boost::numeric::ublas::extents<>;

  const std::vector<extents_t> extents =
    {
//      extents_t{},            // 0

      extents_t{1,1},         // 1
      extents_t{1,2},         // 2
      extents_t{2,1},         // 3

      extents_t{2,3},         // 4
      extents_t{2,3,1},       // 5
      extents_t{1,2,3},       // 6
      extents_t{1,1,2,3},     // 7
      extents_t{1,2,3,1,1},   // 8

      extents_t{4,2,3},       // 9
      extents_t{4,2,1,3},     // 10
      extents_t{4,2,1,3,1},   // 11
      extents_t{1,4,2,1,3,1}  // 12
  };
};

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_expression_retrieve_extents, value,  test_types, fixture)
{
    namespace ublas = boost::numeric::ublas;
    using value_t  = typename value::first_type;
    using layout_t = typename value::second_type;
    using tensor_t = ublas::tensor_dynamic<value_t, layout_t>;

    auto uplus1 = [](auto const& a){ return a + value_t(1);  };
    auto uplus2 = [](auto const& a){ return value_t(2) + a;  };
    auto bplus  = std::plus <value_t>{};
    auto bminus = std::minus<value_t>{};

    for(auto const& e : extents) {

        auto t = tensor_t(e);
        auto v = value_t{};
        for(auto& tt: t){ tt = v; v+=value_t{1}; }

        BOOST_CHECK( ublas::detail::retrieve_extents( t ) == e );

        // uexpr1 = t+1
        // uexpr2 = 2+t
        auto uexpr1 = ublas::detail::make_unary_tensor_expression<tensor_t>( t, uplus1 );
        auto uexpr2 = ublas::detail::make_unary_tensor_expression<tensor_t>( t, uplus2 );

        BOOST_CHECK( ublas::detail::retrieve_extents( uexpr1 ) == e );
        BOOST_CHECK( ublas::detail::retrieve_extents( uexpr2 ) == e );

        // bexpr_uexpr = (t+1) + (2+t)
        auto bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( uexpr1, uexpr2, bplus );

        BOOST_CHECK( ublas::detail::retrieve_extents( bexpr_uexpr ) == e );

        // bexpr_bexpr_uexpr = ((t+1) + (2+t)) - t
        auto bexpr_bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( bexpr_uexpr, t, bminus );

        BOOST_CHECK( ublas::detail::retrieve_extents( bexpr_bexpr_uexpr ) == e );

    }

    for(auto i = 0u; i < extents.size()-1; ++i)
    {

        auto v = value_t{};

        auto t1 = tensor_t(extents[i]);
        for(auto& tt: t1){ tt = v; v+=value_t{1}; }

        auto t2 = tensor_t(extents[i+1]);
        for(auto& tt: t2){ tt = v; v+=value_t{2}; }

        BOOST_CHECK( ublas::detail::retrieve_extents( t1 ) != ublas::detail::retrieve_extents( t2 ) );

        // uexpr1 = t1+1
        // uexpr2 = 2+t2
        auto uexpr1 = ublas::detail::make_unary_tensor_expression<tensor_t>( t1, uplus1 );
        auto uexpr2 = ublas::detail::make_unary_tensor_expression<tensor_t>( t2, uplus2 );

        BOOST_CHECK( ublas::detail::retrieve_extents( t1 )     == ublas::detail::retrieve_extents( uexpr1 ) );
        BOOST_CHECK( ublas::detail::retrieve_extents( t2 )     == ublas::detail::retrieve_extents( uexpr2 ) );
        BOOST_CHECK( ublas::detail::retrieve_extents( uexpr1 ) != ublas::detail::retrieve_extents( uexpr2 ) );

        // bexpr_uexpr = (t1+1) + (2+t2)
        auto bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( uexpr1, uexpr2, bplus );

        BOOST_CHECK( ublas::detail::retrieve_extents( bexpr_uexpr ) == ublas::detail::retrieve_extents(t1) );

        // bexpr_bexpr_uexpr = ((t1+1) + (2+t2)) - t2
        auto bexpr_bexpr_uexpr1 = ublas::detail::make_binary_tensor_expression<tensor_t>( bexpr_uexpr, t2, bminus );

        BOOST_CHECK( ublas::detail::retrieve_extents( bexpr_bexpr_uexpr1 ) == ublas::detail::retrieve_extents(t2) );

        // bexpr_bexpr_uexpr = t2 - ((t1+1) + (2+t2))
        auto bexpr_bexpr_uexpr2 = ublas::detail::make_binary_tensor_expression<tensor_t>( t2, bexpr_uexpr, bminus );

        BOOST_CHECK( ublas::detail::retrieve_extents( bexpr_bexpr_uexpr2 ) == ublas::detail::retrieve_extents(t2) );
    }
}

BOOST_FIXTURE_TEST_CASE_TEMPLATE( test_tensor_expression_all_extents_equal, value,  test_types, fixture)
{
    namespace ublas = boost::numeric::ublas;
    using value_t  = typename value::first_type;
    using layout_t = typename value::second_type;
    using tensor_t = ublas::tensor_dynamic<value_t, layout_t>;

    auto uplus1 = [](auto const& a){ return a + value_t(1);  };
    auto uplus2 = [](auto const& a){ return value_t(2) + a;  };
    auto bplus  = std::plus <value_t>{};
    auto bminus = std::minus<value_t>{};

    for(auto const& e : extents) {

        auto t = tensor_t(e);
        auto v = value_t{};
        for(auto& tt: t){ tt = v; v+=value_t{1}; }

        BOOST_CHECK( ublas::detail::all_extents_equal( t , e ) );

        // uexpr1 = t+1
        // uexpr2 = 2+t
        auto uexpr1 = ublas::detail::make_unary_tensor_expression<tensor_t>( t, uplus1 );
        auto uexpr2 = ublas::detail::make_unary_tensor_expression<tensor_t>( t, uplus2 );

        BOOST_CHECK( ublas::detail::all_extents_equal( uexpr1, e ) );
        BOOST_CHECK( ublas::detail::all_extents_equal( uexpr2, e ) );

        // bexpr_uexpr = (t+1) + (2+t)
        auto bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( uexpr1, uexpr2, bplus );

        BOOST_CHECK( ublas::detail::all_extents_equal( bexpr_uexpr, e ) );

        // bexpr_bexpr_uexpr = ((t+1) + (2+t)) - t
        auto bexpr_bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( bexpr_uexpr, t, bminus );

        BOOST_CHECK( ublas::detail::all_extents_equal( bexpr_bexpr_uexpr , e ) );

    }

    for(auto i = 0u; i < extents.size()-1; ++i)
    {

        auto v = value_t{};

        auto t1 = tensor_t(extents[i]);
        for(auto& tt: t1){ tt = v; v+=value_t{1}; }

        auto t2 = tensor_t(extents[i+1]);
        for(auto& tt: t2){ tt = v; v+=value_t{2}; }

        BOOST_CHECK( ublas::detail::all_extents_equal( t1, ublas::detail::retrieve_extents(t1) ) );
        BOOST_CHECK( ublas::detail::all_extents_equal( t2, ublas::detail::retrieve_extents(t2) ) );

        // uexpr1 = t1+1
        // uexpr2 = 2+t2
        auto uexpr1 = ublas::detail::make_unary_tensor_expression<tensor_t>( t1, uplus1 );
        auto uexpr2 = ublas::detail::make_unary_tensor_expression<tensor_t>( t2, uplus2 );

        BOOST_CHECK( ublas::detail::all_extents_equal( uexpr1, ublas::detail::retrieve_extents(uexpr1) ) );
        BOOST_CHECK( ublas::detail::all_extents_equal( uexpr2, ublas::detail::retrieve_extents(uexpr2) ) );

        // bexpr_uexpr = (t1+1) + (2+t2)
        auto bexpr_uexpr = ublas::detail::make_binary_tensor_expression<tensor_t>( uexpr1, uexpr2, bplus );

        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_uexpr, ublas::detail::retrieve_extents( bexpr_uexpr  ) ) );

        // bexpr_bexpr_uexpr = ((t1+1) + (2+t2)) - t2
        auto bexpr_bexpr_uexpr1 = ublas::detail::make_binary_tensor_expression<tensor_t>( bexpr_uexpr, t2, bminus );

        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_bexpr_uexpr1, ublas::detail::retrieve_extents( bexpr_bexpr_uexpr1  ) ) );

        // bexpr_bexpr_uexpr = t2 - ((t1+1) + (2+t2))
        auto bexpr_bexpr_uexpr2 = ublas::detail::make_binary_tensor_expression<tensor_t>( t2, bexpr_uexpr, bminus );

        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_bexpr_uexpr2, ublas::detail::retrieve_extents( bexpr_bexpr_uexpr2  ) ) );

        // bexpr_uexpr2 = (t1+1) + t2
        auto bexpr_uexpr2 = ublas::detail::make_binary_tensor_expression<tensor_t>( uexpr1, t2, bplus );
        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_uexpr2, ublas::detail::retrieve_extents( bexpr_uexpr2  ) ) );

        // bexpr_uexpr2 = ((t1+1) + t2) + t1
        auto bexpr_bexpr_uexpr3 = ublas::detail::make_binary_tensor_expression<tensor_t>( bexpr_uexpr2, t1, bplus );
        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_bexpr_uexpr3, ublas::detail::retrieve_extents( bexpr_bexpr_uexpr3  ) ) );

        // bexpr_uexpr2 = t1 + (((t1+1) + t2) + t1)
        auto bexpr_bexpr_uexpr4 = ublas::detail::make_binary_tensor_expression<tensor_t>( t1, bexpr_bexpr_uexpr3, bplus );
        BOOST_CHECK( ! ublas::detail::all_extents_equal( bexpr_bexpr_uexpr4, ublas::detail::retrieve_extents( bexpr_bexpr_uexpr4  ) ) );

    }
}

BOOST_AUTO_TEST_SUITE_END()


#include <boost/format.hpp>
#include <cmath>
#include "biggles/detail/fun.hpp"
#include "tools/debug.hpp"
#include <time.h>
namespace biggles {

    using boost::format;
    using boost::str;

    const std::string SI_PREFIX(" kMGTPEZY");

    int num_digits(const float number) {
        return int(std::floor(std::log10(std::abs(number))+1));
    }

    /// \brief converts a size_t to a string using SI prefixes
    std::string count2str(size_t cnt) {
        if (cnt < 1000) {
            return str(format("%6d") % cnt);
        }
        size_t ind = 0;
        float num(cnt);
        while (num >= 1000.f) {
            num /= 1000.f;
            ++ind;
        }
        return str(format("%5.1f%c") % num % SI_PREFIX[ind]);
    }

    /// \brief converts a float to a string using SI prefixes
    std::string float2str(float num) {
        if (std::abs(num) < 1000.f) {
            return str(format("%7.1f") % num);
        }
        size_t ind = 0;
        while (std::abs(num) >= 1000.f) {
            num /= 1000.f;
            ++ind;
        }
        std::string formstr(str(format("%%6.%df%%c") % (4 - num_digits(num))));
        return str(format(formstr) % num % SI_PREFIX[ind]);
    }

    std::string time2str(float seconds) {
        size_t sec = static_cast<size_t>(seconds);
        size_t minutes = sec/size_t(60);
        size_t hours = minutes/size_t(60);
        return str(format("%3d:%02d:%02d") % hours % (minutes % 60) % (sec % 60));
    }

    /// \brief returns the local time as a string formated by \e strftime
    std::string local_time(const std::string& fmt) {
        time_t rawtime;
        struct tm * timeinfo;
        char buffer [160];
        time (&rawtime);
        timeinfo = localtime (&rawtime);
        size_t n = strftime (buffer,160, fmt.c_str(),timeinfo);
        return std::string(buffer).substr(0, n);
    }
}


#ifndef BART_SRC_DOMAIN_FINITE_ELEMENT_I_H
#define BART_SRC_DOMAIN_FINITE_ELEMENT_I_H

#include <deal.II/fe/fe_values.h>

#include "domain/domain_types.hpp"
#include "utility/has_description.h"

//! Classes that provide a finite-element basis
namespace bart::domain::finite_element {

/*! \brief Interface for a finite element object based on the dealii library.
 *
 * This object is responsible for returning finite element data such as: basis
 * function values at quadrature points, basis function values on faces, numbers
 * of quadrature points and dofs per cell.
 *
 */
template <int dim>
class FiniteElementI : public utility::HasDescription {
 public:
  using CellPtr = domain::CellPtr<dim>;
  using DealiiVector = dealii::Vector<double>;
  using Tensor = dealii::Tensor<1, dim>;
  virtual ~FiniteElementI() = default;

  // Basic FE properties
  /*! \brief Gets polynomial degree */
  virtual auto polynomial_degree() const -> int = 0;
  /*! \brief Gets degrees of freedom per cell */
  virtual auto dofs_per_cell() const -> int = 0;
  /*! \brief Gets number of quadrature points per cell */
  virtual auto n_cell_quad_pts() const -> int = 0;
  /*! \brief Gets number of quadrature points per face */
  virtual auto n_face_quad_pts() const -> int = 0;

  // Various methods to access the underlying finite element object
  /*! \brief Sets the cell for finite element values.
   *
   * Does nothing if the cell is already set to the passed value.
   *
   * \param to_set cell to set
   * \return bool indicating if the cell was changed.
   */
  virtual auto SetCell(const CellPtr &to_set) -> bool = 0;

  /*! \brief Sets the face and cell.
   *
   * \param to_set cell to set
   * \param face face number to set
   * \return indicating if the cell was changed.
   */
  virtual auto SetFace(const CellPtr &to_set, const domain::FaceIndex face) -> bool = 0;

  /*! \brief Get the value of shape functions.
   *
   * Returns the value of the \f$i\f$-th shape function, \f$\varphi_i\f$, that
   * corresponds to the \f$i\f$-th cell degree of freedom, evaluated at cell
   * quadrature point, \f$q\f$.
   *
   * \param cell_degree_of_freedom value of \f$i\f$.
   * \param cell_quadrature_point cell quadrature point \f$q\f$ to evaluate the
   * shape function.
   * \return double corresponding to the value.
   */
  virtual auto ShapeValue(const int cell_degree_of_freedom, const int cell_quadrature_point) const -> double = 0;

  /*! \brief Get the value of face shape functions.
   *
   * Returns the value of the \f$i\f$-th shape function, \f$\varphi_i\f$, that
   * corresponds to the \f$i\f$-th cell degree of freedom, evaluated at face
   * quadrature point, \f$q\f$.
   *
   * \param cell_degree_of_freedom value of \f$i\f$.
   * \param face_quadrature_point face quadrature point \f$q\f$ to evaluate the
   * shape function.
   * \return double corresponding to the value.
   */
  virtual auto FaceShapeValue(const int cell_degree_of_freedom, const int face_quadrature_point) const -> double = 0;

  /*! \brief Get the value of gradient of the shape function.
   *
   * Returns the value of the gradient of the \f$i\f$-th shape function,
   * \f$\nabla\varphi_i\f$, that corresponds to the \f$i\f$-th cell degree of
   * freedom, evaluated at cell quadrature point, \f$q\f$.
   *
   * \param cell_degree_of_freedom value of \f$i\f$.
   * \param cell_quadrature_point cell quadrature point \f$q\f$ to evaluate the
   * shape function.
   * \return dealii::Tensor with the value of the gradient.
   */
  virtual auto ShapeGradient(const int cell_degree_of_freedom, const int cell_quadrature_point) const -> Tensor = 0;

  /*! \brief Get the value of the Jacobian for the current cell.
   *
   * Returns the Jacobian evaluated at the specified quadrature point, needed
   * for all integrations.
   *
   * \param cell_quadrature_point cell quadrature point \f$q\f$ to evaluate the
   * Jacobian function.
   * \return double corresponding to the value.
   */
  virtual auto Jacobian(const int cell_quadrature_point) const -> double = 0;

  /*! \brief Get the value of the Jacobian for the current face.
 *
 * Returns the Jacobian evaluated at the specified face quadrature point, needed
 * for all integrations.
 *
 * \param face_quadrature_point face quadrature point \f$q\f$ to evaluate the
 * Jacobian function.
 * \return double corresponding to the value.
 */
  virtual auto FaceJacobian(const int face_quadrature_point) const -> double = 0;

  /*! \brief Get the value of the face normal for the current face.
   *
   * Returns the normal vector for the specified face.
   *
   */
  virtual auto FaceNormal() const -> Tensor = 0;

  /*! \brief Get the value of a vector at the interior cell quadrature points.
   *
   * \param vector_at_dofs vector value at dofs
   * \return a vector holding the value of the vector at each quadrature point.
   */
  virtual auto ValueAtQuadrature(const DealiiVector& vector_at_dofs) const -> std::vector<double> = 0;

  /*! \brief Get the value of an MPI Vector at the cell face quadrature points.
   *
   * @param mpi_vector mpi vector to get the face values of.
   * @return a vector holding the value of the mpi vector at each face quadrature point.
   */
   virtual auto ValueAtFaceQuadrature(const DealiiVector& values_at_dofs) const -> std::vector<double> = 0;

  // DealII Finite element object access. These methods access the underlying
  // finite element objects.
  /*! \brief Gets pointer to the underlying finite element object */
  virtual auto finite_element() -> dealii::FiniteElement<dim, dim>* = 0;
  /*! \brief Gets pointer to the underlying finite element values object */
  virtual auto values() -> dealii::FEValues<dim>* = 0;
  /*! \brief Gets pointer to underlying finite element faces values object */
  virtual auto face_values() -> dealii::FEFaceValues<dim>* = 0;
  /*! \brief Gets pointer to second face values object.
   * This is often used if you need to get face values for a neighbor cell. */
  virtual auto neighbor_face_values() -> dealii::FEFaceValues<dim>* = 0;
};

} // namespace bart::domain::finite_element

#endif //BART_SRC_DOMAIN_FINITE_ELEMENT_I_H


/*
 *  Authors: Lana Mineh and John Scott
 *  Copyright 2021 Phasecraft Ltd. and John Scott
 *
 *  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 gates.cpp
 * \brief Contains tests for the gate member functions
 */

#include <gtest/gtest.h>
#include <qsl/qubits.hpp>
#include <qsl/concepts.hpp>
#include <armadillo>
#include <complex>
#include <list>
#include "test-utils.hpp"

template<typename T>
struct SimWrapper
{
    using Sim = T;
};

/**
 * \brief Typed test suite for one-qubit gates
 */
template <typename T>
class Gates : public testing::Test {
public:
    using List = std::list<T>;
    static T shared_;
    T value_;
};

template <typename T>
class NPGates : public testing::Test {
public:
    using List = std::list<T>;
    static T shared_;
    T value_;
};

/// List of simulator types to check
using Sim1 = SimWrapper<qsl::Qubits<qsl::Type::Default, float>>;
using Sim2 = SimWrapper<qsl::Qubits<qsl::Type::Default, double>>;
using Sim3 = SimWrapper<qsl::Qubits<qsl::Type::Omp, float>>;
using Sim4 = SimWrapper<qsl::Qubits<qsl::Type::Omp, double>>;
using Sim5 = SimWrapper<qsl::Qubits<qsl::Type::Resize, float>>;
using Sim6 = SimWrapper<qsl::Qubits<qsl::Type::Resize, double>>;
using SimTypes = ::testing::Types<Sim1, Sim2, Sim3, Sim4, Sim5, Sim6>;

using Sim7 = SimWrapper<qsl::Qubits<qsl::Type::NP, float>>;
using Sim8 = SimWrapper<qsl::Qubits<qsl::Type::NP, double>>;
using Sim9 = SimWrapper<qsl::Qubits<qsl::Type::OmpNP, float>>;
using Sim10 = SimWrapper<qsl::Qubits<qsl::Type::OmpNP, double>>;

using NPSimTypes = ::testing::Types<Sim7,Sim8,Sim9,Sim10>;

TYPED_TEST_SUITE(Gates, SimTypes);
TYPED_TEST_SUITE(NPGates, NPSimTypes);

TEST(GateTests, MakeMatrixTestsPauliX)
{
    // Pauli gate, targ = 4
    const unsigned num_qubits{ 7 };
    const unsigned targ{ 4 };
    arma::Mat<std::complex<double>> mat(2, 2, arma::fill::zeros);
    mat(0b0,0b1) = 1;
    mat(0b1,0b0) = 1;

    // Make identity gate
    // Create gate in armadillo
    // Sizes of idenity matrix padding
    std::size_t pre = 1 << targ;
    std::size_t post = 1 << (num_qubits - targ - 1);

    // Tensor to make the gate. Have fun with these lines...
    arma::SpMat<std::complex<double>> gate =
	arma::speye<arma::SpMat<std::complex<double>>>(pre, pre);
    gate = arma::kron(arma::conv_to<arma::SpMat<std::complex<double>>>::from(mat),
		      gate);
    gate = arma::kron(arma::speye<arma::SpMat<std::complex<double>>>(post, post),
		      gate);

    // Compute the same matrix using the makeMatrix function
    const auto gate2{ makeMatrix(mat, num_qubits, {targ}) };
    EXPECT_TRUE(arma::approx_equal(gate2, gate, "absdiff", 1e-8));

}

TEST(GateTests,MakeMatrixTestsCnot)
{

    // CNOT gate, qubit 1 is ctrl, qubit 0 is targ
    arma::Mat<std::complex<double>> gate(4, 4, arma::fill::zeros);
    gate(0b00,0b00) = 1;
    gate(0b01,0b01) = 1;
    gate(0b10,0b11) = 1;
    gate(0b11,0b10) = 1;

    // Apply CNOT to qubits ctrl = 1, targ = 0
    unsigned ctrl = 1;
    unsigned targ = 0;
    const auto m0 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m0 << std::endl;
    // Check all the basis states
    const std::complex<double> one{1,0};
    EXPECT_EQ(m0(0b000,0b000), one); // |000) -> |000)
    EXPECT_EQ(m0(0b001,0b001), one); // |001) -> |001)
    EXPECT_EQ(m0(0b010,0b011), one); // |010) -> |011)
    EXPECT_EQ(m0(0b011,0b010), one); // |011) -> |010)
    EXPECT_EQ(m0(0b100,0b100), one); // |100) -> |100)
    EXPECT_EQ(m0(0b101,0b101), one); // |101) -> |101)
    EXPECT_EQ(m0(0b110,0b111), one); // |110) -> |111)
    EXPECT_EQ(m0(0b111,0b110), one); // |111) -> |110)

    // Apply CNOT to qubits ctrl = 0, targ = 2
    ctrl = 0;
    targ = 2;
    const auto m1 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m1 << std::endl;
    // Check all the basis states
    EXPECT_EQ(m1(0b000,0b000), one);
    EXPECT_EQ(m1(0b001,0b101), one);
    EXPECT_EQ(m1(0b010,0b010), one);
    EXPECT_EQ(m1(0b011,0b111), one);
    EXPECT_EQ(m1(0b100,0b100), one);
    EXPECT_EQ(m1(0b101,0b001), one);
    EXPECT_EQ(m1(0b110,0b110), one);
    EXPECT_EQ(m1(0b111,0b011), one);

    // Check the degnerate cases
    ctrl = 0;
    targ = 1;
    const auto m2 = makeMatrix(gate, 2, {targ, ctrl});
    //std::cout << m2 << std::endl;
    // Check all the basis states
    EXPECT_EQ(m2(0b00,0b00), one);
    EXPECT_EQ(m2(0b01,0b11), one);
    EXPECT_EQ(m2(0b10,0b10), one);
    EXPECT_EQ(m2(0b11,0b01), one);

    // Check the degnerate case the other way round
    ctrl = 1;
    targ = 0;
    const auto m3 = makeMatrix(gate, 2, {targ, ctrl});
    //std::cout << m3 << std::endl;
    // Check all the basis states
    EXPECT_EQ(m3(0b00,0b00), one);
    EXPECT_EQ(m3(0b01,0b01), one);
    EXPECT_EQ(m3(0b10,0b11), one);
    EXPECT_EQ(m3(0b11,0b10), one);
}

TEST(GateTests,MakeMatrixTestsSwap)
{

    // SWAP gate, qubit 1 is ctrl, qubit 0 is targ
    arma::Mat<std::complex<double>> gate(4, 4, arma::fill::zeros);
    gate(0b00,0b00) = 1;
    gate(0b01,0b10) = 1;
    gate(0b10,0b01) = 1;
    gate(0b11,0b11) = 1;

    // Apply SWAP to qubits ctrl = 1, targ = 0
    unsigned ctrl = 1;
    unsigned targ = 0;
    const auto m0 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m0 << std::endl;
    // Check all the basis states
    const std::complex<double> one{1,0};
    EXPECT_EQ(m0(0b000,0b000), one); // |000) -> |000)
    EXPECT_EQ(m0(0b001,0b010), one); // |001) -> |010)
    EXPECT_EQ(m0(0b010,0b001), one); // |010) -> |001)
    EXPECT_EQ(m0(0b011,0b011), one); // |011) -> |011)
    EXPECT_EQ(m0(0b100,0b100), one); // |100) -> |100)
    EXPECT_EQ(m0(0b101,0b110), one); // |101) -> |110)
    EXPECT_EQ(m0(0b110,0b101), one); // |110) -> |101)
    EXPECT_EQ(m0(0b111,0b111), one); // |111) -> |111)

    // Apply SWAP to qubits ctrl = 0, targ = 2
    ctrl = 0;
    targ = 2;
    const auto m1 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m1 << std::endl;
    // Check all the basis states
    EXPECT_EQ(m1(0b000,0b000), one);
    EXPECT_EQ(m1(0b001,0b100), one);
    EXPECT_EQ(m1(0b010,0b010), one);
    EXPECT_EQ(m1(0b011,0b110), one);
    EXPECT_EQ(m1(0b100,0b001), one);
    EXPECT_EQ(m1(0b101,0b101), one);
    EXPECT_EQ(m1(0b110,0b011), one);
    EXPECT_EQ(m1(0b111,0b111), one);

}

TEST(GateTests,MakeMatrixTestsCHadamard)
{
    // Controlled Hadamard gate, qubit 1 is ctrl, qubit 0 is targ
    arma::Mat<std::complex<double>> gate(4, 4, arma::fill::zeros);
    double one_sqrt2{1/std::sqrt(2)};
    gate(0b00,0b00) = 1;
    gate(0b01,0b01) = 1;
    gate(0b10,0b10) = one_sqrt2;
    gate(0b10,0b11) = one_sqrt2;
    gate(0b11,0b10) = one_sqrt2;
    gate(0b11,0b11) = -one_sqrt2;

    // Apply controlled Hadamard to qubits ctrl = 1, targ = 0
    unsigned ctrl = 1;
    unsigned targ = 0;
    const auto m0 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m0 << std::endl;
    // Check all the basis states
    const std::complex<double> one{1,0};
    const std::complex<double> val{one_sqrt2,0};
    EXPECT_EQ(m0(0b000,0b000), one); // |000) -> |000)
    EXPECT_EQ(m0(0b001,0b001), one); // |001) -> |010)
    EXPECT_EQ(m0(0b010,0b010), val); // |010) -> [ |010) + |011) ] / sqrt2
    EXPECT_EQ(m0(0b010,0b011), val);
    EXPECT_EQ(m0(0b011,0b010), val); // |011) -> [ |010) - |011) ] / sqrt2
    EXPECT_EQ(m0(0b011,0b011), -val);
    EXPECT_EQ(m0(0b100,0b100), one); // |100) -> |100)
    EXPECT_EQ(m0(0b101,0b101), one); // |101) -> |110)
    EXPECT_EQ(m0(0b110,0b110), val); // |110) -> [ |110) + |111) ] / sqrt2
    EXPECT_EQ(m0(0b110,0b111), val);
    EXPECT_EQ(m0(0b111,0b110), val); // |111) -> [ |110) - |111) ] / sqrt2
    EXPECT_EQ(m0(0b111,0b111), -val);

    // Apply controlled Hadamard to qubits ctrl = 1, targ = 0
    ctrl = 0;
    targ = 2;
    const auto m1 = makeMatrix(gate, 3, {targ, ctrl});
    //std::cout << m1 << std::endl;

    // Check all the basis states
    // Identity
    EXPECT_EQ(m1(0b000,0b000), one);
    EXPECT_EQ(m1(0b010,0b010), one);
    EXPECT_EQ(m1(0b100,0b100), one);
    EXPECT_EQ(m1(0b110,0b110), one);

    // To the plus state
    EXPECT_EQ(m1(0b001,0b001), val);
    EXPECT_EQ(m1(0b001,0b101), val);
    EXPECT_EQ(m1(0b011,0b011), val);
    EXPECT_EQ(m1(0b011,0b111), val);

    // To the minus state
    EXPECT_EQ(m1(0b101,0b001), val);
    EXPECT_EQ(m1(0b101,0b101), -val);
    EXPECT_EQ(m1(0b111,0b011), val);
    EXPECT_EQ(m1(0b111,0b111), -val);

}

/// Test the makeMatrix function against tensor product matrix
TEST(GateTests,MakeMatrixTestsTensor)
{
    const unsigned num_qubits{ 8 };
    unsigned ctrl = 6;
    unsigned targ = 5;

    // Controlled Hadamard gate, qubit 1 is ctrl, qubit 0 is targ
    arma::Mat<std::complex<double>> mat(4, 4, arma::fill::zeros);
    double one_sqrt2{1/std::sqrt(2)};
    mat(0b00,0b00) = 1;
    mat(0b01,0b01) = 1;
    mat(0b10,0b10) = one_sqrt2;
    mat(0b10,0b11) = one_sqrt2;
    mat(0b11,0b10) = one_sqrt2;
    mat(0b11,0b11) = -one_sqrt2;

    // Make identity gate
    // Create gate in armadillo
    // Sizes of idenity matrix padding
    std::size_t pre = 1 << targ;
    std::size_t post = 1 << (num_qubits - ctrl - 1);

    // Tensor to make the gate. Have fun with these lines...
    arma::SpMat<std::complex<double>> gate =
	arma::speye<arma::SpMat<std::complex<double>>>(pre, pre);
    gate = arma::kron(arma::conv_to<arma::SpMat<std::complex<double>>>::from(mat),
		      gate);
    gate = arma::kron(arma::speye<arma::SpMat<std::complex<double>>>(post, post),
		      gate);

    // Compute the same matrix using the makeMatrix function
    const auto gate2{ makeMatrix(mat, num_qubits, {targ, ctrl}) };
    EXPECT_TRUE(arma::approx_equal(gate2, gate, "absdiff", 1e-8));
}

TYPED_TEST(Gates, OneQubitGate)
{
    const unsigned num_qubits = 8;
    const unsigned targ = 4;
    using Sim = TypeParam::Sim;
    using Fp = TypeParam::Sim::Fp_type;
    // Generate random angle to test parameterised gates
    Fp angle = arma::datum::pi * arma::randn<Fp>();

    // Create list of gates mapped to matrices
    std::vector<std::pair<std::function<void(Sim &, unsigned)>,
			  arma::Mat<std::complex<Fp>>>> gates;

    // PauliX
    arma::Mat<std::complex<Fp>> pauliX(2, 2, arma::fill::zeros);
    pauliX(0, 1) = 1;
    pauliX(1, 0) = 1;
    auto fn_pauliX = [](Sim & sim, unsigned targ) {
			 sim.pauliX(targ);
		     };
    gates.push_back({fn_pauliX, pauliX});

    // PauliY
    arma::Mat<std::complex<Fp>> pauliY(2, 2, arma::fill::zeros);
    pauliY(0, 1) = std::complex<Fp>{0, -1};
    pauliY(1, 0) = std::complex<Fp>{0, 1};
    auto fn_pauliY = [](Sim & sim, unsigned targ) {
    			 sim.pauliY(targ);
    		     };
    gates.push_back({fn_pauliY, pauliY});

    // PauliZ
    arma::Mat<std::complex<Fp>> pauliZ(2, 2, arma::fill::zeros);
    pauliZ(0, 0) = 1;
    pauliZ(1, 1) = -1;
    auto fn_pauliZ = [](Sim & sim, unsigned targ) {
			 sim.pauliZ(targ);
		     };
    gates.push_back({fn_pauliZ, pauliZ});

    // Hadamard
    arma::Mat<std::complex<Fp>> hadamard(2, 2, arma::fill::zeros);
    Fp sqrt2 = 1/std::sqrt(2);
    hadamard(0, 0) = sqrt2;
    hadamard(0, 1) = sqrt2;
    hadamard(1, 0) = sqrt2;
    hadamard(1, 1) = -sqrt2;
    auto fn_hadamard = [](Sim & sim, unsigned targ) {
			   sim.hadamard(targ);
		       };
    gates.push_back({fn_hadamard, hadamard});

    // phase shift
    arma::Mat<std::complex<Fp>> phase(2, 2, arma::fill::zeros);
    phase(0, 0) = 1;
    phase(1, 1) = std::complex<Fp>{std::cos(angle), std::sin(angle)};
    auto fn_phase = [=](Sim & sim, unsigned targ) {
			sim.phase(targ, angle);
		       };
    gates.push_back({fn_phase, phase});

    // rotateX
    arma::Mat<std::complex<Fp>> rotateX(2, 2, arma::fill::zeros);
    rotateX(0, 0) = std::cos(angle/2);
    rotateX(0, 1) = std::complex<Fp>{0,-std::sin(angle/2)};
    rotateX(1, 0) = std::complex<Fp>{0,-std::sin(angle/2)};
    rotateX(1, 1) = std::cos(angle/2);
    auto fn_rotateX = [=](Sim & sim, unsigned targ) {
			sim.rotateX(targ, angle);
		       };
    gates.push_back({fn_rotateX, rotateX});

    // rotateY
    arma::Mat<std::complex<Fp>> rotateY(2, 2, arma::fill::zeros);
    rotateY(0, 0) = std::cos(angle/2);
    rotateY(0, 1) = -std::sin(angle/2);
    rotateY(1, 0) = std::sin(angle/2);
    rotateY(1, 1) = std::cos(angle/2);
    auto fn_rotateY = [=](Sim & sim, unsigned targ) {
			sim.rotateY(targ, angle);
		       };
    gates.push_back({fn_rotateY, rotateY});

    // rotateZ
    arma::Mat<std::complex<Fp>> rotateZ(2, 2, arma::fill::zeros);
    rotateZ(0, 0) = std::complex<Fp>{std::cos(angle/2), std::sin(-angle/2)};
    rotateZ(1, 1) = std::complex<Fp>{std::cos(angle/2), std::sin(angle/2)};
    auto fn_rotateZ = [=](Sim & sim, unsigned targ) {
			sim.rotateZ(targ, angle);
		       };
    gates.push_back({fn_rotateZ, rotateZ});

    for (const auto & [fn, mat] : gates) {

	// Make a random state
	Sim q{num_qubits};
	const std::vector<qsl::complex<Fp>> state
	    = qsl::makeRandomState<Fp>(num_qubits);
	q.setState(state);

	// Set an armadillo vector to the same state
	std::size_t dim = 1 << num_qubits;
	arma::Col<std::complex<Fp>> v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    v(i) = std::complex<Fp>{state[i].real, state[i].imag};
	}

	// Apply gate to qubits
	std::invoke(fn, q, targ);

	// Apply gate in armadillo
	v = makeMatrix(mat, num_qubits, {targ}) * v;

	// Read qubit state into armadillo
	std::vector<qsl::complex<Fp>> res = q.getState();
	arma::Col<std::complex<Fp>> qubit_v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    qubit_v(i) = std::complex<Fp>{res[i].real, res[i].imag};
	}

    	EXPECT_TRUE(arma::approx_equal(v, qubit_v, "both", 1e-6, 1e-8));
    }
}

TYPED_TEST(Gates, TwoQubitGate)
{
    const unsigned num_qubits = 8;
    const unsigned ctrl = 2;
    const unsigned targ = 4;
    using Sim = TypeParam::Sim;
    using Fp = TypeParam::Sim::Fp_type;
    // Generate random angle to test parameterised gates
    Fp angle = arma::datum::pi * arma::randn<Fp>();

    // Create list of gates mapped to matrices
    std::vector<std::pair<std::function<void(Sim &, unsigned, unsigned)>,
			  arma::Mat<std::complex<Fp>>>> gates;

    // controlNot
    arma::Mat<std::complex<Fp>> cnot(4, 4, arma::fill::zeros);
    cnot(0b00,0b00) = 1;
    cnot(0b01,0b01) = 1;
    cnot(0b10,0b11) = 1;
    cnot(0b11,0b10) = 1;

    auto fn_cnot = [](Sim & sim, unsigned ctrl, unsigned targ) {
			 sim.controlNot(ctrl, targ);
		     };
    gates.push_back({fn_cnot, cnot});

    // controlY
    arma::Mat<std::complex<Fp>> cy(4, 4, arma::fill::zeros);
    cy(0b00,0b00) = 1;
    cy(0b01,0b01) = 1;
    cy(0b10,0b11) = std::complex<Fp>{0, -1};
    cy(0b11,0b10) = std::complex<Fp>{0, 1};;
    auto fn_cy = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		     sim.controlY(ctrl, targ);
		   };
    gates.push_back({fn_cy, cy});

    // controlZ
    arma::Mat<std::complex<Fp>> cz(4, 4, arma::fill::zeros);
    cz(0b00,0b00) = 1;
    cz(0b01,0b01) = 1;
    cz(0b10,0b10) = 1;
    cz(0b11,0b11) = -1;
    auto fn_cz = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		     sim.controlZ(ctrl, targ);
		   };
    gates.push_back({fn_cz, cz});

    // controlRotateX
    arma::Mat<std::complex<Fp>> crx(4, 4, arma::fill::zeros);
    crx(0b00,0b00) = 1;
    crx(0b01,0b01) = 1;
    crx(0b10,0b10) = std::complex<Fp>{std::cos(angle/2), 0};
    crx(0b10,0b11) = std::complex<Fp>{0, -std::sin(angle/2)};
    crx(0b11,0b10) = std::complex<Fp>{0, -std::sin(angle/2)};
    crx(0b11,0b11) = std::complex<Fp>{std::cos(angle/2), 0};

    auto fn_crx = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.controlRotateX(ctrl, targ, angle);
		   };
    gates.push_back({fn_crx, crx});

    // controlRotateY
    arma::Mat<std::complex<Fp>> cry(4, 4, arma::fill::zeros);
    cry(0b00,0b00) = 1;
    cry(0b01,0b01) = 1;
    cry(0b10,0b10) = std::cos(angle/2);
    cry(0b10,0b11) = -std::sin(angle/2);
    cry(0b11,0b10) = std::sin(angle/2);
    cry(0b11,0b11) = std::cos(angle/2);

    auto fn_cry = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.controlRotateY(ctrl, targ, angle);
		   };
    gates.push_back({fn_cry, cry});

    // controlRotateZ
    arma::Mat<std::complex<Fp>> crz(4, 4, arma::fill::zeros);
    crz(0b00,0b00) = 1;
    crz(0b01,0b01) = 1;
    crz(0b10,0b10) = std::complex<Fp>{std::cos(angle/2), -std::sin(angle/2)};
    crz(0b11,0b11) = std::complex<Fp>{std::cos(angle/2), std::sin(angle/2)};
    auto fn_crz = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.controlRotateZ(ctrl, targ, angle);
		   };
    gates.push_back({fn_crz, crz});

    // controlPhase
    arma::Mat<std::complex<Fp>> cphase(4, 4, arma::fill::zeros);
    cphase(0b00,0b00) = 1;
    cphase(0b01,0b01) = 1;
    cphase(0b10,0b10) = 1;
    cphase(0b11,0b11) = std::complex<Fp>{std::cos(angle), std::sin(angle)};
    auto fn_cphase = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		       sim.controlPhase(ctrl, targ, angle);
		   };
    gates.push_back({fn_cphase, cphase});

    // controlHadamard
    Fp sqrt2 = 1/std::sqrt(2);
    arma::Mat<std::complex<Fp>> ch(4, 4, arma::fill::zeros);
    ch(0b00,0b00) = 1;
    ch(0b01,0b01) = 1;
    ch(0b10,0b10) = sqrt2;
    ch(0b10,0b11) = sqrt2;
    ch(0b11,0b10) = sqrt2;
    ch(0b11,0b11) = -sqrt2;

    auto fn_ch = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.controlHadamard(ctrl, targ);
		   };
    gates.push_back({fn_ch, ch});

    // swap
    arma::Mat<std::complex<Fp>> swap(4, 4, arma::fill::zeros);
    swap(0b00,0b00) = 1;
    swap(0b01,0b10) = 1;
    swap(0b10,0b01) = 1;
    swap(0b11,0b11) = 1;
    auto fn_swap = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		       sim.swap(ctrl, targ);
		   };
    gates.push_back({fn_swap, swap});

    // fswap
    arma::Mat<std::complex<Fp>> fswap(4, 4, arma::fill::zeros);
    fswap(0b00,0b00) = 1;
    fswap(0b01,0b10) = 1;
    fswap(0b10,0b01) = 1;
    fswap(0b11,0b11) = -1;
    auto fn_fswap = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.fswap(ctrl, targ);
		   };
    gates.push_back({fn_fswap, fswap});

    // npRotateX
    arma::Mat<std::complex<Fp>> nprx(4, 4, arma::fill::zeros);
    nprx(0b00,0b00) = 1;
    nprx(0b01,0b01) = std::complex<Fp>{std::cos(angle/2), 0};
    nprx(0b01,0b10) = std::complex<Fp>{0, -std::sin(angle/2)};
    nprx(0b10,0b01) = std::complex<Fp>{0, -std::sin(angle/2)};
    nprx(0b10,0b10) = std::complex<Fp>{std::cos(angle/2), 0};
    nprx(0b11,0b11) = 1;

    auto fn_nprx = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npRotateX(targ, ctrl, angle);
		   };
    gates.push_back({fn_nprx, nprx});

    // npRotateY
    arma::Mat<std::complex<Fp>> npry(4, 4, arma::fill::zeros);
    npry(0b00,0b00) = 1;
    npry(0b01,0b01) = std::cos(angle/2);
    npry(0b01,0b10) = -std::sin(angle/2);
    npry(0b10,0b01) = std::sin(angle/2);
    npry(0b10,0b10) = std::cos(angle/2);
    npry(0b11,0b11) = 1;

    auto fn_npry = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npRotateY(targ, ctrl, angle);
		   };
    gates.push_back({fn_npry, npry});

    // npHadamard
    arma::Mat<std::complex<Fp>> nph(4, 4, arma::fill::zeros);
    nph(0b00,0b00) = 1;
    nph(0b01,0b01) = sqrt2;
    nph(0b01,0b10) = sqrt2;
    nph(0b10,0b01) = sqrt2;
    nph(0b10,0b10) = -sqrt2;
    nph(0b11,0b11) = 1;

    auto fn_nph = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npHadamard(targ, ctrl);
		   };
    gates.push_back({fn_nph, nph});

    for (const auto & [fn, mat] : gates) {

	// Make a random state
	Sim q{num_qubits};
	const std::vector<qsl::complex<Fp>> state
	    = qsl::makeRandomState<Fp>(num_qubits);
	q.setState(state);

	// Set an armadillo vector to the same state
	std::size_t dim = 1 << num_qubits;
	arma::Col<std::complex<Fp>> v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    v(i) = std::complex<Fp>{state[i].real, state[i].imag};
	}

	// Apply gate to qubits
	std::invoke(fn, q, ctrl, targ);

	// Apply gate in armadillo
	v = makeMatrix(mat, num_qubits, {targ, ctrl}) * v;

	// Read qubit state into armadillo
	std::vector<qsl::complex<Fp>> res = q.getState();
	arma::Col<std::complex<Fp>> qubit_v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    qubit_v(i) = std::complex<Fp>{res[i].real, res[i].imag};
	}

    	EXPECT_TRUE(arma::approx_equal(v, qubit_v, "both", 1e-6, 1e-8));
    }
}

TYPED_TEST(NPGates, TwoQubitGate)
{
    const unsigned num_qubits = 8;
    const unsigned num_ones = 4;
    const unsigned ctrl = 2;
    const unsigned targ = 4;
    using Sim = TypeParam::Sim;
    using Fp = TypeParam::Sim::Fp_type;
    // Generate random angle to test parameterised gates
    Fp angle = arma::datum::pi * arma::randn<Fp>();

    // Create list of gates mapped to matrices
    std::vector<std::pair<std::function<void(Sim &, unsigned, unsigned)>,
			  arma::Mat<std::complex<Fp>>>> gates;

    // controlZ
    arma::Mat<std::complex<Fp>> cz(4, 4, arma::fill::zeros);
    cz(0b00,0b00) = 1;
    cz(0b01,0b01) = 1;
    cz(0b10,0b10) = 1;
    cz(0b11,0b11) = -1;
    auto fn_cz = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		     sim.controlZ(ctrl, targ);
		   };
    gates.push_back({fn_cz, cz});

    // controlRotateZ
    arma::Mat<std::complex<Fp>> crz(4, 4, arma::fill::zeros);
    crz(0b00,0b00) = 1;
    crz(0b01,0b01) = 1;
    crz(0b10,0b10) = std::complex<Fp>{std::cos(angle/2), -std::sin(angle/2)};
    crz(0b11,0b11) = std::complex<Fp>{std::cos(angle/2), std::sin(angle/2)};
    auto fn_crz = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.controlRotateZ(ctrl, targ, angle);
		   };
    gates.push_back({fn_crz, crz});

    // controlPhase
    arma::Mat<std::complex<Fp>> cphase(4, 4, arma::fill::zeros);
    cphase(0b00,0b00) = 1;
    cphase(0b01,0b01) = 1;
    cphase(0b10,0b10) = 1;
    cphase(0b11,0b11) = std::complex<Fp>{std::cos(angle), std::sin(angle)};
    auto fn_cphase = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		       sim.controlPhase(ctrl, targ, angle);
		   };
    gates.push_back({fn_cphase, cphase});

    // swap
    arma::Mat<std::complex<Fp>> swap(4, 4, arma::fill::zeros);
    swap(0b00,0b00) = 1;
    swap(0b01,0b10) = 1;
    swap(0b10,0b01) = 1;
    swap(0b11,0b11) = 1;
    auto fn_swap = [=](Sim & sim, unsigned ctrl, unsigned targ) {
		       sim.swap(ctrl, targ);
		   };
    gates.push_back({fn_swap, swap});

    // fswap
    arma::Mat<std::complex<Fp>> fswap(4, 4, arma::fill::zeros);
    fswap(0b00,0b00) = 1;
    fswap(0b01,0b10) = 1;
    fswap(0b10,0b01) = 1;
    fswap(0b11,0b11) = -1;
    auto fn_fswap = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.fswap(ctrl, targ);
		   };
    gates.push_back({fn_fswap, fswap});

    // npRotateX
    arma::Mat<std::complex<Fp>> nprx(4, 4, arma::fill::zeros);
    nprx(0b00,0b00) = 1;
    nprx(0b01,0b01) = std::complex<Fp>{std::cos(angle/2), 0};
    nprx(0b01,0b10) = std::complex<Fp>{0, -std::sin(angle/2)};
    nprx(0b10,0b01) = std::complex<Fp>{0, -std::sin(angle/2)};
    nprx(0b10,0b10) = std::complex<Fp>{std::cos(angle/2), 0};
    nprx(0b11,0b11) = 1;

    auto fn_nprx = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npRotateX(targ, ctrl, angle);
		   };
    gates.push_back({fn_nprx, nprx});

    // npRotateY
    arma::Mat<std::complex<Fp>> npry(4, 4, arma::fill::zeros);
    npry(0b00,0b00) = 1;
    npry(0b01,0b01) = std::cos(angle/2);
    npry(0b01,0b10) = -std::sin(angle/2);
    npry(0b10,0b01) = std::sin(angle/2);
    npry(0b10,0b10) = std::cos(angle/2);
    npry(0b11,0b11) = 1;

    auto fn_npry = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npRotateY(targ, ctrl, angle);
		   };
    gates.push_back({fn_npry, npry});

    // npHadamard
    arma::Mat<std::complex<Fp>> nph(4, 4, arma::fill::zeros);
    Fp sqrt2 = 1/std::sqrt(2);
    nph(0b00,0b00) = 1;
    nph(0b01,0b01) = sqrt2;
    nph(0b01,0b10) = sqrt2;
    nph(0b10,0b01) = sqrt2;
    nph(0b10,0b10) = -sqrt2;
    nph(0b11,0b11) = 1;

    auto fn_nph = [=](Sim & sim, unsigned ctrl, unsigned targ) {
	sim.npHadamard(targ, ctrl);
		   };
    gates.push_back({fn_nph, nph});

    for (const auto & [fn, mat] : gates) {

	// Make a random state
	Sim q{num_qubits};
	const std::vector<qsl::complex<Fp>> state
	    = qsl::makeRandomNPState<Fp>(num_qubits, num_ones);
	q.setState(state);

	// Set an armadillo vector to the same state
	std::size_t dim = 1 << num_qubits;
	arma::Col<std::complex<Fp>> v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    v(i) = std::complex<Fp>{state[i].real, state[i].imag};
	}

	// Apply gate to qubits
	std::invoke(fn, q, ctrl, targ);

	// Apply gate in armadillo
	v = makeMatrix(mat, num_qubits, {targ, ctrl}) * v;

	// Read qubit state into armadillo
	std::vector<qsl::complex<Fp>> res = q.getState();
	arma::Col<std::complex<Fp>> qubit_v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    qubit_v(i) = std::complex<Fp>{res[i].real, res[i].imag};
	}

    	EXPECT_TRUE(arma::approx_equal(v, qubit_v, "both", 1e-6, 1e-8));
    }
}

TYPED_TEST(NPGates, OneQubitGate)
{
    const unsigned num_qubits = 8;
    const unsigned num_ones = 4;
    const unsigned targ = 4;
    using Sim = TypeParam::Sim;
    using Fp = TypeParam::Sim::Fp_type;
    // Generate random angle to test parameterised gates
    Fp angle = arma::datum::pi * arma::randn<Fp>();

    // Create list of gates mapped to matrices
    std::vector<std::pair<std::function<void(Sim &, unsigned)>,
			  arma::SpMat<std::complex<Fp>>>> gates;

    // PauliZ
    arma::SpMat<std::complex<Fp>> pauliZ(2, 2);
    pauliZ(0, 0) = 1;
    pauliZ(1, 1) = -1;
    auto fn_pauliZ = [](Sim & sim, unsigned targ) {
			 sim.pauliZ(targ);
		     };
    gates.push_back({fn_pauliZ, pauliZ});

    // phase shift
    arma::SpMat<std::complex<Fp>> phase(2, 2);
    phase(0, 0) = 1;
    phase(1, 1) = std::complex<Fp>{std::cos(angle), std::sin(angle)};
    auto fn_phase = [=](Sim & sim, unsigned targ) {
			sim.phase(targ, angle);
		       };
    gates.push_back({fn_phase, phase});

    // rotateZ
    arma::SpMat<std::complex<Fp>> rotateZ(2, 2);
    rotateZ(0, 0) = std::complex<Fp>{std::cos(angle/2), std::sin(-angle/2)};
    rotateZ(1, 1) = std::complex<Fp>{std::cos(angle/2), std::sin(angle/2)};
    auto fn_rotateZ = [=](Sim & sim, unsigned targ) {
			sim.rotateZ(targ, angle);
		       };
    gates.push_back({fn_rotateZ, rotateZ});

    for (const auto & [fn, mat] : gates) {

	// Make a random state
	Sim q{num_qubits};
	const std::vector<qsl::complex<Fp>> state
	    = qsl::makeRandomNPState<Fp>(num_qubits, num_ones);
	q.setState(state);

	// Set an armadillo vector to the same state
	std::size_t dim = 1 << num_qubits;
	arma::Col<std::complex<Fp>> v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    v(i) = std::complex<Fp>{state[i].real, state[i].imag};
	}

	// Apply gate to qubits
	std::invoke(fn, q, targ);

	// Create gate in armadillo
	// Sizes of idenity matrix padding
	std::size_t pre = 1 << targ;
	std::size_t post = 1 << (num_qubits - targ - 1);

	// Tensor to make the gate
	arma::SpMat<std::complex<Fp>> gate =
	    arma::speye<arma::SpMat<std::complex<Fp>>>(pre, pre);
	gate = arma::kron(mat, gate);
	gate = arma::kron(arma::speye<arma::SpMat<std::complex<Fp>>>(post, post),
			  gate);

	// Apply gate in armadillo
	v = gate * v;

	// Read qubit state into armadillo
	std::vector<qsl::complex<Fp>> res = q.getState();
	arma::Col<std::complex<Fp>> qubit_v(dim);
	for (std::size_t i = 0; i < dim; i++) {
	    qubit_v(i) = std::complex<Fp>{res[i].real, res[i].imag};
	}

    	EXPECT_TRUE(arma::approx_equal(v, qubit_v, "both", 1e-6, 1e-8));
    }
}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2012 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2012 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2012 Mateusz Loskot, London, UK.
// Copyright (c) 2013 Adam Wulkiewicz, Lodz, Poland.

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_ALGORITHMS_CONVERT_HPP
#define BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_ALGORITHMS_CONVERT_HPP

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

#include <boost/geometry/algorithms/convert.hpp>

#include <boost/geometry/extensions/algebra/core/tags.hpp>

#include <boost/geometry/extensions/algebra/algorithms/detail.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template <typename RQuaternion, typename RMatrix>
struct convert<RQuaternion, RMatrix, rotation_quaternion_tag, rotation_matrix_tag, 3, false>
{
    static inline void apply(RQuaternion const& q, RMatrix& m)
    {
        typedef typename coordinate_type<RQuaternion>::type T;

        // quaternion should be normalized

        T xx2 = get<1>(q) * get<1>(q) * 2;
        T yy2 = get<2>(q) * get<2>(q) * 2;
        T zz2 = get<3>(q) * get<3>(q) * 2;
        T xy2 = get<1>(q) * get<2>(q) * 2;
        T yz2 = get<2>(q) * get<3>(q) * 2;
        T xz2 = get<1>(q) * get<3>(q) * 2;
        T wx2 = get<0>(q) * get<1>(q) * 2;
        T wy2 = get<0>(q) * get<2>(q) * 2;
        T wz2 = get<0>(q) * get<3>(q) * 2;

        // WARNING!
        // Quaternion (0, 0, 0, 0) is converted to identity matrix!

        set<0, 0>(m, 1-yy2-zz2); set<0, 1>(m, xy2-wz2);   set<0, 2>(m, xz2+wy2);
        set<1, 0>(m, xy2+wz2);   set<1, 1>(m, 1-xx2-zz2); set<1, 2>(m, yz2-wx2);
        set<2, 0>(m, xz2-wy2);   set<2, 1>(m, yz2+wx2);   set<2, 2>(m, 1-xx2-yy2);
    }
};

template <typename RMatrix, typename RQuaternion>
struct convert<RMatrix, RQuaternion, rotation_matrix_tag, rotation_quaternion_tag, 3, false>
{
    static inline void apply(RMatrix const& m, RQuaternion & q)
    {
        typedef typename coordinate_type<RMatrix>::type T;

        // WARNING!
        // Zero matrix is converted to quaternion(0.5, 0, 0, 0)!

        T w = math::sqrt(1 + get<0, 0>(m) + get<1, 1>(m) + get<2, 2>(m)) / 2;
        T iw4 = 0.25 / w;
        set<0>(q, w);
        set<1>(q, (get<2, 1>(m) - get<1, 2>(m)) * iw4);
        set<2>(q, (get<0, 2>(m) - get<2, 0>(m)) * iw4);
        set<3>(q, (get<1, 0>(m) - get<0, 1>(m)) * iw4);
    }
};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_EXTENSIONS_ALGEBRA_ALGORITHMS_ASSIGN_HPP


#include "lin_reg_uni_hierarchy.h"

#include <Eigen/Dense>
#include <stan/math/prim/err.hpp>

#include "src/utils/eigen_utils.h"
#include "src/utils/proto_utils.h"
#include "src/utils/rng.h"

void LinRegUniHierarchy::initialize() {
  if (prior == nullptr) {
    throw std::invalid_argument("Hierarchy prior was not provided");
  }
  state.regression_coeffs = hypers->mean;
  state.var = hypers->scale / (hypers->shape + 1);
  clear_data();
}

void LinRegUniHierarchy::clear_data() {
  mixed_prod = Eigen::VectorXd::Zero(dim);
  data_sum_squares = 0.0;
  covar_sum_squares = Eigen::MatrixXd::Zero(dim, dim);
  card = 0;
  cluster_data_idx = std::set<int>();
}

void LinRegUniHierarchy::update_summary_statistics(
    const Eigen::VectorXd &datum, const Eigen::VectorXd &covariate, bool add) {
  if (add) {
    data_sum_squares += datum(0) * datum(0);
    covar_sum_squares += covariate * covariate.transpose();
    mixed_prod += datum(0) * covariate;
  } else {
    data_sum_squares -= datum(0) * datum(0);
    covar_sum_squares -= covariate * covariate.transpose();
    mixed_prod -= datum(0) * covariate;
  }
}

LinRegUniHierarchy::Hyperparams LinRegUniHierarchy::normal_invgamma_update() {
  Hyperparams post_params;

  post_params.var_scaling = covar_sum_squares + hypers->var_scaling;
  post_params.mean = post_params.var_scaling.llt().solve(
      mixed_prod + hypers->var_scaling * hypers->mean);
  post_params.shape = hypers->shape + 0.5 * card;
  post_params.scale =
      hypers->scale +
      0.5 * (data_sum_squares +
             hypers->mean.transpose() * hypers->var_scaling * hypers->mean -
             post_params.mean.transpose() * post_params.var_scaling *
                 post_params.mean);
  return post_params;
}

void LinRegUniHierarchy::update_hypers(
    const std::vector<bayesmix::MarginalState::ClusterState> &states) {
  auto &rng = bayesmix::Rng::Instance().get();
  if (prior->has_fixed_values()) {
    return;
  }

  else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

double LinRegUniHierarchy::like_lpdf(
    const Eigen::RowVectorXd &datum,
    const Eigen::RowVectorXd &covariate) const {
  return stan::math::normal_lpdf(
      datum(0), state.regression_coeffs.dot(covariate), sqrt(state.var));
}

double LinRegUniHierarchy::marg_lpdf(
    const Eigen::RowVectorXd &datum,
    const Eigen::RowVectorXd &covariate) const {
  double sig_n = sqrt(
      (1 + (covariate * hypers->var_scaling_inv * covariate.transpose())(0)) *
      hypers->scale / hypers->shape);
  return stan::math::student_t_lpdf(datum(0), 2 * hypers->shape,
                                    covariate.dot(hypers->mean), sig_n);
}

void LinRegUniHierarchy::draw() {
  // Generate new state values from their prior centering distribution
  auto &rng = bayesmix::Rng::Instance().get();
  state.var = stan::math::inv_gamma_rng(hypers->shape, hypers->scale, rng);
  state.regression_coeffs = stan::math::multi_normal_prec_rng(
      hypers->mean, hypers->var_scaling / state.var, rng);
}

void LinRegUniHierarchy::sample_given_data() {
  // Update values
  Hyperparams params = normal_invgamma_update();

  // Generate new state values from their prior centering distribution
  auto &rng = bayesmix::Rng::Instance().get();
  state.var = stan::math::inv_gamma_rng(params.shape, params.scale, rng);
  state.regression_coeffs = stan::math::multi_normal_prec_rng(
      params.mean, params.var_scaling / state.var, rng);
}

void LinRegUniHierarchy::sample_given_data(const Eigen::MatrixXd &data,
                                           const Eigen::MatrixXd &covariates) {
  data_sum_squares = data.squaredNorm();
  covar_sum_squares = covariates.transpose() * covariates;
  mixed_prod = covariates.transpose() * data;
  card = data.rows();
  log_card = std::log(card);
  sample_given_data();
}

void LinRegUniHierarchy::set_state_from_proto(
    const google::protobuf::Message &state_) {
  auto &statecast = google::protobuf::internal::down_cast<
      const bayesmix::MarginalState::ClusterState &>(state_);
  state.regression_coeffs =
      bayesmix::to_eigen(statecast.lin_reg_uni_ls_state().regression_coeffs());
  state.var = statecast.lin_reg_uni_ls_state().var();
  set_card(statecast.cardinality());
}

void LinRegUniHierarchy::set_prior(const google::protobuf::Message &prior_) {
  auto &priorcast =
      google::protobuf::internal::down_cast<const bayesmix::LinRegUniPrior &>(
          prior_);
  prior = std::make_shared<bayesmix::LinRegUniPrior>(priorcast);
  hypers = std::make_shared<Hyperparams>();
  if (prior->has_fixed_values()) {
    // Set values
    hypers->mean = bayesmix::to_eigen(prior->fixed_values().mean());
    dim = hypers->mean.size();
    hypers->var_scaling =
        bayesmix::to_eigen(prior->fixed_values().var_scaling());
    hypers->var_scaling_inv = stan::math::inverse_spd(hypers->var_scaling);
    hypers->shape = prior->fixed_values().shape();
    hypers->scale = prior->fixed_values().scale();
    // Check validity
    bayesmix::check_spd(hypers->var_scaling);
    if (hypers->shape <= 0) {
      throw std::invalid_argument("Shape parameter must be > 0");
    }
    if (hypers->scale <= 0) {
      throw std::invalid_argument("Scale parameter must be > 0");
    }
  }

  else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

void LinRegUniHierarchy::write_state_to_proto(
    google::protobuf::Message *out) const {
  bayesmix::LinRegUniLSState state_;
  bayesmix::to_proto(state.regression_coeffs,
                     state_.mutable_regression_coeffs());
  state_.set_var(state.var);

  auto *out_cast = google::protobuf::internal::down_cast<
      bayesmix::MarginalState::ClusterState *>(out);
  out_cast->mutable_lin_reg_uni_ls_state()->CopyFrom(state_);
  out_cast->set_cardinality(card);
}

void LinRegUniHierarchy::write_hypers_to_proto(
    google::protobuf::Message *out) const {
  bayesmix::LinRegUniPrior hypers_;
  bayesmix::to_proto(hypers->mean,
                     hypers_.mutable_fixed_values()->mutable_mean());
  bayesmix::to_proto(hypers->var_scaling,
                     hypers_.mutable_fixed_values()->mutable_var_scaling());
  hypers_.mutable_fixed_values()->set_shape(hypers->shape);
  hypers_.mutable_fixed_values()->set_scale(hypers->scale);

  google::protobuf::internal::down_cast<bayesmix::LinRegUniPrior *>(out)
      ->mutable_fixed_values()
      ->CopyFrom(hypers_.fixed_values());
}


#ifndef REALTIME_MF_OPENNI_HPP_INCLUDED
#define REALTIME_MF_OPENNI_HPP_INCLUDED

#include <string>

#include <pcl/io/openni_grabber.h>
#include <pcl/io/openni_camera/openni_depth_image.h>
#include <pcl/io/openni_camera/openni_image.h>

#include <Eigen/Dense>

#include <mmf/optimizationSO3.hpp>
#include <mmf/optimizationSO3_approx.hpp>
#include <mmf/optimizationSO3_vmf.hpp>

#include <cudaPcl/openniSmoothNormalsGpu.hpp>
#include <cudaPcl/timerLog.hpp>
#include <cudaPcl/gpuMatrix.hpp>

using namespace Eigen;
using namespace std;

class RealtimeMF_openni : public cudaPcl::OpenniSmoothNormalsGpu
{
  public:
    RealtimeMF_openni(std::string mode, uint32_t w, uint32_t h,
        float f_d);
    virtual ~RealtimeMF_openni();

    virtual void normals_cb(float* d_normals, uint8_t* haveData, uint32_t
        w, uint32_t h);

    virtual void visualizeNormals();

    void init(int w, int h);

    static void projectDirections(cv::Mat& I, const MatrixXf& dirs, double f_d,
        const Matrix<uint8_t, Dynamic,Dynamic>& colors);

  protected:

    double residual_, D_KL_;
    uint32_t nCGIter_;

    cudaPcl::TimerLog tLog_;

    uint32_t w_, h_;

    ofstream fout_;
    std::string mode_;

    Matrix3f kRw_;
    MatrixXf mfAxes_;

    mmf::OptSO3 *optSO3_;

    VectorXu z_;
    cv::Mat zIrgb_;
    cv::Mat Icomb_;

    Matrix3d kRv;
    Vector3d dt0;
};

// -------------------------------- impl -----------------------------------

RealtimeMF_openni::RealtimeMF_openni(std::string mode, uint32_t w, uint32_t h,
    float f_d)
  : OpenniSmoothNormalsGpu(f_d, 0.2*0.2, 9, true),
  // TODO check B and eps!
  residual_(0.0), D_KL_(0.0), nCGIter_(10),
  tLog_("./timer.log",3,10,"TimerLog"),
  w_(w), h_(h),
  fout_("./stats.log",ofstream::out),
  mode_(mode), kRw_(Matrix3f::Identity()),
  optSO3_(NULL)
{};

RealtimeMF_openni::~RealtimeMF_openni()
{
  delete optSO3_;
  fout_.close();
}

void RealtimeMF_openni::init(int w, int h)
{
  if(optSO3_ != NULL) return;
  cout<<"inititalizing optSO3"<<endl;
  // TODO: use weights
  // TODO: make sure N is float all over the place
  //optSO3_ = new OptSO3(25.0f*M_PI/180.0f,d_n,w,h,NULL);//,d_weights);
  if(mode_.compare("direct") == 0)
  {
#ifndef WEIGHTED
    optSO3_ = new mmf::OptSO3(5.0f*M_PI/180.0f);//,d_weights);
#else
    optSO3_ = new mmf::OptSO3(5.0f*M_PI/180.0f);
//        normalExtractor_.d_weights());
#endif
    nCGIter_ = 10; // cannot do that many iterations
  }else if (mode_.compare("approx") == 0){
#ifndef WEIGHTED
    optSO3_ = new mmf::OptSO3Approx(5.0f*M_PI/180.0f);//,d_weights);
#else
    optSO3_ = new mmf::OptSO3Approx(5.0f*M_PI/180.0f);
//        normalExtractor_.d_weights());
#endif
    nCGIter_ = 25;
  }else if (mode_.compare("vmf") == 0){
#ifndef WEIGHTED
    optSO3_ = new mmf::OptSO3vMF(5.0f*M_PI/180.0f);//,d_weights);
#else
    optSO3_ = new mmf::OptSO3vMF(5.0f*M_PI/180.0f);
//        normalExtractor_.d_weights());
#endif
    nCGIter_ = 25;
  }
}

void RealtimeMF_openni::projectDirections(cv::Mat& I, const MatrixXf& dirs,
    double f_d, const Matrix<uint8_t,Dynamic,Dynamic>& colors)
{
  double scale = 0.1;
  VectorXf p0(3); p0 << 0.35,0.25,1;
  double u0 = p0(0)/p0(2)*f_d + 320.;
  double v0 = p0(1)/p0(2)*f_d + 240.;
  for(uint32_t k=0; k < dirs.cols(); ++k)
  {
    VectorXf p1 = p0 + dirs.col(k)*scale;
    double u1 = p1(0)/p1(2)*f_d + 320.;
    double v1 = p1(1)/p1(2)*f_d + 240.;
    cv::line(I, cv::Point(u0,v0), cv::Point(u1,v1),
        CV_RGB(colors(k,0),colors(k,1),colors(k,2)), 2, CV_AA);

    double arrowLen = 10.;
    double angle = atan2(v1-v0,u1-u0);

    double ru1 = u1 - arrowLen*cos(angle + M_PI*0.25);
    double rv1 = v1 - arrowLen*sin(angle + M_PI*0.25);
    cv::line(I, cv::Point(u1,v1), cv::Point(ru1,rv1),
        CV_RGB(colors(k,0),colors(k,1),colors(k,2)), 2, CV_AA);
    ru1 = u1 - arrowLen*cos(angle - M_PI*0.25);
    rv1 = v1 - arrowLen*sin(angle - M_PI*0.25);
    cv::line(I, cv::Point(u1,v1), cv::Point(ru1,rv1),
        CV_RGB(colors(k,0),colors(k,1),colors(k,2)), 2, CV_AA);
  }
  cv::circle(I, cv::Point(u0,v0), 2, CV_RGB(0,0,0), 2, CV_AA);
}

void RealtimeMF_openni::normals_cb(float* d_normals, uint8_t* haveData,
    uint32_t w, uint32_t h)
{
  tLog_.tic(-1); // reset all timers
  int32_t nComp = 0;
  float* d_nComp = this->normalExtract->d_normalsComp(nComp);
  Matrix3f kRwBefore = kRw_;
  tLog_.toctic(0,1);

  optSO3_->updateExternalGpuNormals(d_nComp,nComp,3,0);
  double residual = optSO3_->conjugateGradientCUDA(kRw_,nCGIter_);
  double D_KL = optSO3_->D_KL_axisUnif();
  tLog_.toctic(1,2);

  {
    boost::mutex::scoped_lock updateLock(this->updateModelMutex);
    this->normalsImg_ = this->normalExtract->normalsImg();
    if(z_.rows() != w*h) z_.resize(w*h);
    this->normalExtract->uncompressCpu(optSO3_->z().data(), optSO3_->z().rows(),
        z_.data(), z_.rows());

    mfAxes_ = MatrixXf::Zero(3,6);
    for(uint32_t k=0; k<6; ++k){
      int j = k/2; // which of the rotation columns does this belong to
      float sign = (- float(k%2) +0.5f)*2.0f; // sign of the axis
      mfAxes_.col(k) = sign*kRw_.col(j);
    }
    D_KL_= D_KL;
    residual_ = residual;
    this->update_ = true;
    updateLock.unlock();
  }

  tLog_.toc(2); // total time
  tLog_.logCycle();
  cout<<"delta rotation kRw_ = \n"<<kRwBefore*kRw_.transpose()<<endl;
  cout<<"---------------------------------------------------------------------------"<<endl;
  tLog_.printStats();
  cout<<" residual="<<residual_<<"\t D_KL="<<D_KL_<<endl;
  cout<<"---------------------------------------------------------------------------"<<endl;
  fout_<<D_KL_<<" "<<residual_<<endl; fout_.flush();

////  return kRw_;
//  {
//    boost::mutex::scoped_lock updateLock(this->updateModelMutex);
////    pcl::PointCloud<pcl::PointXYZRGB>::ConstPtr nDispPtr =
////      normalExtract->normalsPc();
////    nDisp_ = pcl::PointCloud<pcl::PointXYZRGB>::Ptr( new
////        pcl::PointCloud<pcl::PointXYZRGB>(*nDispPtr));
////    this->normalsImg_ = this->normalExtract->normalsImg();
//    this->update_ = true;
//  }
};

void RealtimeMF_openni::visualizeNormals()
{
  if (normalsImg_.empty() || normalsImg_.rows == 0 || normalsImg_.cols == 0) return;
  cout<<"entering RealtimeMF_openni::visualizeNormals"<<endl;
//  pcl::PointCloud<pcl::PointXYZRGB>::Ptr nDisp(
//      new pcl::PointCloud<pcl::PointXYZRGB>(*nDisp_));
  const w = this->normalsImg_.cols;
  const h = this->normalsImg_.rows;

  cv::Mat nI(h,w,CV_8UC3);
  cv::Mat nIRGB(h,w,CV_8UC3);
  this->normalsImg_.convertTo(nI,CV_8UC3,127.5f,127.5f);
  cv::cvtColor(nI,this->nIRGB_,CV_RGB2BGR);
  cv::imshow("normals",this->nIRGB_);

//  uint32_t k=0;

  Matrix<uint8_t,Dynamic,Dynamic> mfAxCols(6,3);
  mfAxCols << 255,0,0,
              255,0,0,
              0,255,0,
              0,255,0,
              0,0,255,
              0,0,255;

//  cout<<" z shape "<<z_.rows()<<" "<< w<<" " <<h<<endl;
//  cv::Mat Iz(nDisp->height/SUBSAMPLE_STEP,nDisp->width/SUBSAMPLE_STEP,CV_8UC1);
  zIrgb_ = cv::Mat(h,w,CV_8UC3);
//  for(uint32_t i=0; i<w; i+=1)
//    for(uint32_t j=0; j<h; j+=1)
//      if(this->normalsImg_.at<cv::Vec3f>(j,i)[0]==this->normalsImg_.at<cv::Vec3f>(j,i)[0])
//      {
//        uint32_t idz = z_(w*j +i);
//        //            cout<<"k "<<k<<" "<< z_.rows() <<"\t"<<z_(k)<<"\t"<<int32_t(idz)<<endl;
//        cout<<j<<" "<<i<<" id="<<idz<<endl;
//        zIrgb_.at<cv::Vec3b>(j,i)[0] = mfAxCols(idz,0);
//        zIrgb_.at<cv::Vec3b>(j,i)[1] = mfAxCols(idz,1);
//        zIrgb_.at<cv::Vec3b>(j,i)[2] = mfAxCols(idz,2);
////        k++;
//      }else{
//        zIrgb_.at<cv::Vec3b>(j,i)[0] = 255;
//        zIrgb_.at<cv::Vec3b>(j,i)[1] = 255;
//        zIrgb_.at<cv::Vec3b>(j,i)[2] = 255;
//      }

//  cout<<this->rgb_.rows <<" " << this->rgb_.cols<<endl;
  if(this->rgb_.rows>1 && this->rgb_.cols >1)
  {
    cv::addWeighted(this->rgb_ , 0.7, zIrgb, 0.3, 0.0, Icomb_);
    projectDirections(Icomb_,mfAxes_,this->f_d_,mfAxCols);
    cv::imshow("dbg",Icomb);
  }else{
    projectDirections(zIrgb_,mfAxes_,this->f_d_,mfAxCols);
    cv::imshow("dbg",zIrgb);
  }

};

#endif


#pragma once

#include "../stateobject.hh"

#include <Eigen/Geometry>

namespace bold
{
  class OdometryState : public StateObject
  {
  public:
    OdometryState(Eigen::Affine3d transform)
      : d_transform{std::move(transform)}
    {}

    /** Gets the cumulative transform of the agent frame
     *
     * The returned transformation describes the agent frame at t = 0
     * (A0) wrt the current frame (At). i.e. AtA0, transforming a
     * vector v in A0 by AtA0 * v0 gives that vector in the current
     * agent frame.
     *
     * Users of this value can compute their own delta values through
     * $\delta = A_tA_{t-1} = A_tA_0 * A_{t-1}A0^{-1}$
     */
    Eigen::Affine3d const& getTransform() const { return d_transform; };

    void writeJson(rapidjson::Writer<rapidjson::StringBuffer>& writer) const override { writeJsonInternal(writer); }
    void writeJson(rapidjson::Writer<WebSocketBuffer>& writer) const override { writeJsonInternal(writer); }
    void writeJson(rapidjson::PrettyWriter<rapidjson::StringBuffer>& writer) const override { writeJsonInternal(writer); }

  private:
    template<typename TBuffer>
    void writeJsonInternal(rapidjson::Writer<TBuffer> &writer) const;

    Eigen::Affine3d d_transform;
  };

  template<typename TBuffer>
  inline void OdometryState::writeJsonInternal(rapidjson::Writer<TBuffer> &writer) const
  {
    writer.StartObject();
    {
      writer.String("tr");
      writer.StartArray();
      for (unsigned j = 0; j < 4; ++j)
      {
        for (unsigned i = 0; i < 4; ++i)
          writer.Double(d_transform.matrix()(i, j), "%.3f");
        writer.String(" ");
      }
      writer.EndArray();
    }
    writer.EndObject();
  }
}


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

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

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

#ifndef BOOST_SIMD_NO_NANS
#include <boost/simd/function/scalar/is_nan.hpp>
#endif
#include <boost/simd/function/scalar/is_gtz.hpp>
#include <boost/simd/function/scalar/is_ltz.hpp>
#include <boost/simd/function/scalar/shift_right.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/config.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  BOOST_DISPATCH_OVERLOAD ( sign_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::signed_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0) const BOOST_NOEXCEPT
    {
      return shift_right(a0, (sizeof(A0)*8-1)) - shift_right(-a0, (sizeof(A0)*8-1));
    }
  };

  BOOST_DISPATCH_OVERLOAD ( sign_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::unsigned_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0) const BOOST_NOEXCEPT
    {
      return !!a0;
    }
  };

  BOOST_DISPATCH_OVERLOAD ( sign_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0) const BOOST_NOEXCEPT
    {
      A0 r = bool(is_gtz(a0))-bool(is_ltz(a0));
#ifdef BOOST_SIMD_NO_NANS
      return r;
#else
      return is_nan(a0) ? a0 : r;
#endif
    }
  };
} } }

#endif