/* ---------------------------------------------------------------------
 *
 * Copyright (C) 2011 - 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.
 *
 * ---------------------------------------------------------------------
 *
 * Authors: Joerg Frohne, Texas A&M University and
 *                        University of Siegen, 2011, 2012
 *          Wolfgang Bangerth, Texas A&M University, 2012
 */

// @sect3{Include files}

// 像往常一样，在开始的时候，我们把所有我们需要的头文件都包含在这里。除了为Trilinos库提供接口的各种文件外，没有什么意外。

#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/function.h>
#include <deal.II/base/index_set.h>

#include <deal.II/lac/affine_constraints.h>
#include <deal.II/lac/vector.h>
#include <deal.II/lac/full_matrix.h>
#include <deal.II/lac/dynamic_sparsity_pattern.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/lac/trilinos_sparse_matrix.h>
#include <deal.II/lac/trilinos_vector.h>
#include <deal.II/lac/trilinos_precondition.h>

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

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

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

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

#include <fstream>
#include <iostream>

namespace Step41
{
  using namespace dealii;
// @sect3{The ObstacleProblem class template}

// 该类提供了描述障碍问题所需的所有函数和变量。它与我们在 step-4 中要做的事情很接近，所以相对简单。唯一真正的新组件是计算主动集合的update_solution_and_constraints函数和一些描述线性系统原始（无约束）形式所需的变量（ complete_system_matrix 和 complete_system_rhs ），以及主动集合本身和主动集合公式中用于缩放拉格朗日乘数的质量矩阵 $B$ 的对角线。其余的内容与 step-4 相同。

  template <int dim>
  class ObstacleProblem
  {
  public:
    ObstacleProblem();
    void run();

  private:
    void make_grid();
    void setup_system();
    void assemble_system();
    void
         assemble_mass_matrix_diagonal(TrilinosWrappers::SparseMatrix &mass_matrix);
    void update_solution_and_constraints();
    void solve();
    void output_results(const unsigned int iteration) const;

    Triangulation<dim>        triangulation;
    FE_Q<dim>                 fe;
    DoFHandler<dim>           dof_handler;
    AffineConstraints<double> constraints;
    IndexSet                  active_set;

    TrilinosWrappers::SparseMatrix system_matrix;
    TrilinosWrappers::SparseMatrix complete_system_matrix;

    TrilinosWrappers::MPI::Vector solution;
    TrilinosWrappers::MPI::Vector system_rhs;
    TrilinosWrappers::MPI::Vector complete_system_rhs;
    TrilinosWrappers::MPI::Vector diagonal_of_mass_matrix;
    TrilinosWrappers::MPI::Vector contact_force;
  };
// @sect3{Right hand side, boundary values, and the obstacle}

// 在下文中，我们定义了描述右侧函数、Dirichlet边界值以及作为 $\mathbf x$ 函数的障碍物高度的类。在这三种情况下，我们都从函数 @<dim@>, 派生出这些类，尽管在 RightHandSide 和 Obstacle 的情况下，这更多的是出于惯例而非必要，因为我们从未将此类对象传递给库。在任何情况下，鉴于我们选择了 $f=-10$  ,  $u|_{\partial\Omega}=0$  ...，右手和边界值类的定义是显而易见的。

  template <int dim>
  class RightHandSide : public Function<dim>
  {
  public:
    virtual double value(const Point<dim> & /*p*/,
                         const unsigned int component = 0) const override
    {
      (void)component;
      AssertIndexRange(component, 1);

      return -10;
    }
  };

  template <int dim>
  class BoundaryValues : public Function<dim>
  {
  public:
    virtual double value(const Point<dim> & /*p*/,
                         const unsigned int component = 0) const override
    {
      (void)component;
      AssertIndexRange(component, 1);

      return 0;
    }
  };

// 我们用一个级联的障碍物来描述障碍物的功能（想想看：楼梯的阶梯）。

  template <int dim>
  class Obstacle : public Function<dim>
  {
  public:
    virtual double value(const Point<dim> & p,
                         const unsigned int component = 0) const override
    {
      (void)component;
      Assert(component == 0, ExcIndexRange(component, 0, 1));

      if (p(0) < -0.5)
        return -0.2;
      else if (p(0) >= -0.5 && p(0) < 0.0)
        return -0.4;
      else if (p(0) >= 0.0 && p(0) < 0.5)
        return -0.6;
      else
        return -0.8;
    }
  };

//  @sect3{Implementation of the ObstacleProblem class}
// @sect4{ObstacleProblem::ObstacleProblem}

// 对每个看过前几个教程程序的人来说，构造函数是完全显而易见的。

  template <int dim>
  ObstacleProblem<dim>::ObstacleProblem()
    : fe(1)
    , dof_handler(triangulation)
  {}
// @sect4{ObstacleProblem::make_grid}

// 我们在二维的正方形 $[-1,1]\times [-1,1]$ 上解决我们的障碍物问题。因此这个函数只是设置了一个最简单的网格。

  template <int dim>
  void ObstacleProblem<dim>::make_grid()
  {
    GridGenerator::hyper_cube(triangulation, -1, 1);
    triangulation.refine_global(7);

    std::cout << "Number of active cells: " << triangulation.n_active_cells()
              << std::endl
              << "Total number of cells: " << triangulation.n_cells()
              << std::endl;
  }
// @sect4{ObstacleProblem::setup_system}

// 在这个值得注意的第一个函数中，我们设置了自由度处理程序，调整了向量和矩阵的大小，并处理了约束。最初，约束条件当然只是由边界值给出的，所以我们在函数的顶部对它们进行插值。

  template <int dim>
  void ObstacleProblem<dim>::setup_system()
  {
    dof_handler.distribute_dofs(fe);
    active_set.set_size(dof_handler.n_dofs());

    std::cout << "Number of degrees of freedom: " << dof_handler.n_dofs()
              << std::endl
              << std::endl;

    VectorTools::interpolate_boundary_values(dof_handler,
                                             0,
                                             BoundaryValues<dim>(),
                                             constraints);
    constraints.close();

    DynamicSparsityPattern dsp(dof_handler.n_dofs());
    DoFTools::make_sparsity_pattern(dof_handler, dsp, constraints, false);

    system_matrix.reinit(dsp);
    complete_system_matrix.reinit(dsp);

    IndexSet solution_index_set = dof_handler.locally_owned_dofs();
    solution.reinit(solution_index_set, MPI_COMM_WORLD);
    system_rhs.reinit(solution_index_set, MPI_COMM_WORLD);
    complete_system_rhs.reinit(solution_index_set, MPI_COMM_WORLD);
    contact_force.reinit(solution_index_set, MPI_COMM_WORLD);

// 这里唯一要做的事情是计算 $B$ 矩阵中的因子，该矩阵用于缩放残差。正如在介绍中所讨论的，我们将使用一个小技巧来使这个质量矩阵成为对角线，在下文中，首先将所有这些计算成一个矩阵，然后提取对角线元素供以后使用。

    TrilinosWrappers::SparseMatrix mass_matrix;
    mass_matrix.reinit(dsp);
    assemble_mass_matrix_diagonal(mass_matrix);
    diagonal_of_mass_matrix.reinit(solution_index_set);
    for (unsigned int j = 0; j < solution.size(); j++)
      diagonal_of_mass_matrix(j) = mass_matrix.diag_element(j);
  }
// @sect4{ObstacleProblem::assemble_system}

// 这个函数一次就把系统矩阵和右手边集合起来，并把约束条件（由于活动集以及来自边界值）应用到我们的系统中。否则，它在功能上等同于例如  step-4  中的相应函数。

  template <int dim>
  void ObstacleProblem<dim>::assemble_system()
  {
    std::cout << "   Assembling system..." << std::endl;

    system_matrix = 0;
    system_rhs    = 0;

    const QGauss<dim>  quadrature_formula(fe.degree + 1);
    RightHandSide<dim> right_hand_side;

    FEValues<dim> fe_values(fe,
                            quadrature_formula,
                            update_values | update_gradients |
                              update_quadrature_points | update_JxW_values);

    const unsigned int dofs_per_cell = fe.n_dofs_per_cell();
    const unsigned int n_q_points    = quadrature_formula.size();

    FullMatrix<double> cell_matrix(dofs_per_cell, dofs_per_cell);
    Vector<double>     cell_rhs(dofs_per_cell);

    std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);

    for (const auto &cell : dof_handler.active_cell_iterators())
      {
        fe_values.reinit(cell);
        cell_matrix = 0;
        cell_rhs    = 0;

        for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
          for (unsigned int i = 0; i < dofs_per_cell; ++i)
            {
              for (unsigned int j = 0; j < dofs_per_cell; ++j)
                cell_matrix(i, j) +=
                  (fe_values.shape_grad(i, q_point) *
                   fe_values.shape_grad(j, q_point) * fe_values.JxW(q_point));

              cell_rhs(i) +=
                (fe_values.shape_value(i, q_point) *
                 right_hand_side.value(fe_values.quadrature_point(q_point)) *
                 fe_values.JxW(q_point));
            }

        cell->get_dof_indices(local_dof_indices);

        constraints.distribute_local_to_global(cell_matrix,
                                               cell_rhs,
                                               local_dof_indices,
                                               system_matrix,
                                               system_rhs,
                                               true);
      }
  }

//  @sect4{ObstacleProblem::assemble_mass_matrix_diagonal}

// 下一个函数用于计算对角线质量矩阵 $B$ ，用于在主动集方法中缩放变量。正如介绍中所讨论的，我们通过选择正交的梯形规则来获得质量矩阵的对角线。这样一来，我们就不再需要在正交点、指数 $i$ 和指数 $j$ 上进行三重循环，而是可以直接使用双重循环。考虑到我们在以前的许多教程程序中讨论过的内容，该函数的其余部分是显而易见的。

// 注意在调用这个函数的时候，约束对象只包含边界值约束；因此我们在最后的复制-本地-全局步骤中不必注意保留矩阵项的值，这些项以后可能会受到活动集的约束。

// 还需要注意的是，只有在我们拥有 $Q_1$ 元素的情况下，使用梯形规则的技巧才有效。对于更高阶的元素，我们需要使用一个正交公式，在有限元的所有支持点都有正交点。构建这样一个正交公式其实并不难，但不是这里的重点，所以我们只是在函数的顶部断言我们对有限元的隐含假设实际上得到了满足。

  template <int dim>
  void ObstacleProblem<dim>::assemble_mass_matrix_diagonal(
    TrilinosWrappers::SparseMatrix &mass_matrix)
  {
    Assert(fe.degree == 1, ExcNotImplemented());

    const QTrapezoid<dim> quadrature_formula;
    FEValues<dim>         fe_values(fe,
                            quadrature_formula,
                            update_values | update_JxW_values);

    const unsigned int dofs_per_cell = fe.n_dofs_per_cell();
    const unsigned int n_q_points    = quadrature_formula.size();

    FullMatrix<double> cell_matrix(dofs_per_cell, dofs_per_cell);
    std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);

      {
        fe_values.reinit(cell);
        cell_matrix = 0;

        for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
          for (unsigned int i = 0; i < dofs_per_cell; ++i)
            cell_matrix(i, i) +=
              (fe_values.shape_value(i, q_point) *
               fe_values.shape_value(i, q_point) * fe_values.JxW(q_point));

        cell->get_dof_indices(local_dof_indices);

        constraints.distribute_local_to_global(cell_matrix,
                                               local_dof_indices,
                                               mass_matrix);
      }
  }
// @sect4{ObstacleProblem::update_solution_and_constraints}

// 在某种意义上，这是本程序的核心功能。 它更新了介绍中所讨论的受限自由度的活动集，并从中计算出一个AffineConstraints对象，然后可以用来在下一次迭代的解中消除受限自由度。同时，我们将解决方案的受限自由度设置为正确的值，即障碍物的高度。

// 从根本上说，这个函数是相当简单的。我们必须在所有自由度上循环，并检查函数 $\Lambda^k_i + c([BU^k]_i - G_i) = \Lambda^k_i + cB_i(U^k_i - [g_h]_i)$ 的符号，因为在我们的例子中 $G_i = B_i[g_h]_i$  。为此，我们使用介绍中给出的公式，通过该公式我们可以计算出拉格朗日乘数，作为原始线性系统的残差（通过变量 complete_system_matrix and complete_system_rhs 给出。在这个函数的顶部，我们使用一个属于矩阵类的函数来计算这个残差。

  template <int dim>
  void ObstacleProblem<dim>::update_solution_and_constraints()
  {
    std::cout << "   Updating active set..." << std::endl;

    const double penalty_parameter = 100.0;

    TrilinosWrappers::MPI::Vector lambda(
      complete_index_set(dof_handler.n_dofs()));
    complete_system_matrix.residual(lambda, solution, complete_system_rhs);

// 计算 contact_force[i] =

// - lambda[i] * diagonal_of_mass_matrix[i]。

    contact_force = lambda;
    contact_force.scale(diagonal_of_mass_matrix);
    contact_force *= -1;

// 下一步是重置活动集和约束对象，并在所有自由度上开始循环。由于我们不能只是在解向量的所有元素上循环，所以这变得稍微复杂了一些，因为我们没有办法找出一个自由度与哪个位置相关；但是，我们需要这个位置来测试一个自由度的位移是大于还是小于这个位置的障碍物高度。

// 我们通过在所有单元和定义在每个单元上的DoF上循环来解决这个问题。我们在这里使用一个 $Q_1$ 函数来描述位移，对于该函数，自由度总是位于单元格的顶点上；因此，我们可以通过询问顶点来获得每个自由度的索引及其位置。另一方面，这显然对高阶元素不起作用，因此我们添加了一个断言，确保我们只处理所有自由度都位于顶点的元素，以避免万一有人想玩增加解的多项式程度时用非功能性代码绊倒自己。

// 循环单元格而不是自由度的代价是我们可能会多次遇到一些自由度，即每次我们访问与给定顶点相邻的一个单元格时。因此，我们必须跟踪我们已经接触过的顶点和尚未接触的顶点。我们通过使用一个标志数组来做到这一点  dof_touched  。

    constraints.clear();
    active_set.clear();

    const Obstacle<dim> obstacle;
    std::vector<bool>   dof_touched(dof_handler.n_dofs(), false);

    for (const auto &cell : dof_handler.active_cell_iterators())
      for (const auto v : cell->vertex_indices())
        {
          Assert(dof_handler.get_fe().n_dofs_per_cell() == cell->n_vertices(),
                 ExcNotImplemented());

          const unsigned int dof_index = cell->vertex_dof_index(v, 0);

          if (dof_touched[dof_index] == false)
            dof_touched[dof_index] = true;
          else
            continue;

// 现在我们知道我们还没有触及这个DoF，让我们得到那里的位移函数的值以及障碍函数的值，并使用这个来决定当前DoF是否属于活动集。为此，我们使用上面和介绍中给出的函数。

// 如果我们决定该DoF应该是活动集的一部分，我们将其索引添加到活动集中，在AffineConstraints对象中引入一个不均匀的平等约束，并将解的值重置为障碍物的高度。最后，系统的非接触部分的残差作为一个额外的控制（残差等于剩余的、未计算的力，在接触区之外应该为零），所以我们把残差向量的分量（即拉格朗日乘数lambda）清零，这些分量对应于身体接触的区域；在所有单元的循环结束时，残差将因此只包括非接触区的残差。我们在循环结束后输出这个残差的准则和活动集的大小。

          const double obstacle_value = obstacle.value(cell->vertex(v));
          const double solution_value = solution(dof_index);

          if (lambda(dof_index) + penalty_parameter *
                                    diagonal_of_mass_matrix(dof_index) *
                                    (solution_value - obstacle_value) <
              0)
            {
              active_set.add_index(dof_index);
              constraints.add_line(dof_index);
              constraints.set_inhomogeneity(dof_index, obstacle_value);

              solution(dof_index) = obstacle_value;

              lambda(dof_index) = 0;
            }
        }
    std::cout << "      Size of active set: " << active_set.n_elements()
              << std::endl;

    std::cout << "   Residual of the non-contact part of the system: "
              << lambda.l2_norm() << std::endl;

// 在最后一步中，我们将迄今为止从活动集合中得到的对DoF的约束加入到那些由Dirichlet边界值产生的约束中，并关闭约束对象。

    VectorTools::interpolate_boundary_values(dof_handler,
                                             0,
                                             BoundaryValues<dim>(),
                                             constraints);
    constraints.close();
  }
// @sect4{ObstacleProblem::solve}

// 关于求解函数，其实没有什么可说的。在牛顿方法的背景下，我们通常对非常高的精度不感兴趣（为什么要求一个高度精确的线性问题的解，而我们知道它只能给我们一个非线性问题的近似解），所以我们使用ReductionControl类，当达到一个绝对公差（为此我们选择 $10^{-12}$ ）或者当残差减少一定的系数（这里是 $10^{-3}$ ）时停止反复运算。

  template <int dim>
  void ObstacleProblem<dim>::solve()
  {
    std::cout << "   Solving system..." << std::endl;

    ReductionControl                        reduction_control(100, 1e-12, 1e-3);
    SolverCG<TrilinosWrappers::MPI::Vector> solver(reduction_control);
    TrilinosWrappers::PreconditionAMG       precondition;
    precondition.initialize(system_matrix);

    solver.solve(system_matrix, solution, system_rhs, precondition);
    constraints.distribute(solution);

    std::cout << "      Error: " << reduction_control.initial_value() << " -> "
              << reduction_control.last_value() << " in "
              << reduction_control.last_step() << " CG iterations."
              << std::endl;
  }
// @sect4{ObstacleProblem::output_results}

// 我们使用vtk-format进行输出。 该文件包含位移和活动集的数字表示。

  template <int dim>
  void ObstacleProblem<dim>::output_results(const unsigned int iteration) const
  {
    std::cout << "   Writing graphical output..." << std::endl;

    TrilinosWrappers::MPI::Vector active_set_vector(
      dof_handler.locally_owned_dofs(), MPI_COMM_WORLD);
    for (const auto index : active_set)
      active_set_vector[index] = 1.;

    DataOut<dim> data_out;

    data_out.attach_dof_handler(dof_handler);
    data_out.add_data_vector(solution, "displacement");
    data_out.add_data_vector(active_set_vector, "active_set");
    data_out.add_data_vector(contact_force, "lambda");

    data_out.build_patches();

    std::ofstream output_vtk("output_" +
                             Utilities::int_to_string(iteration, 3) + ".vtk");
    data_out.write_vtk(output_vtk);
  }

//  @sect4{ObstacleProblem::run}

// 这是一个对所有事情都有最高级别控制的函数。 它并不长，而且事实上相当直接：在主动集方法的每一次迭代中，我们都要组装线性系统，求解它，更新主动集并将解投射回可行集，然后输出结果。只要主动集在前一次迭代中没有变化，迭代就会终止。

// 唯一比较棘手的部分是，我们必须在第一次迭代组装好线性系统（即矩阵和右手边）后保存它。原因是这是唯一一个我们可以在没有任何接触约束的情况下访问线性系统的步骤。我们需要这个来计算其他迭代中解决方案的残差，但是在其他迭代中，我们形成的线性系统中对应于约束自由度的行和列都被消除了，因此我们不能再访问原始方程的全部残差。

  template <int dim>
  void ObstacleProblem<dim>::run()
  {
    make_grid();
    setup_system();

    IndexSet active_set_old(active_set);
    for (unsigned int iteration = 0; iteration <= solution.size(); ++iteration)
      {
        std::cout << "Newton iteration " << iteration << std::endl;

        assemble_system();

        if (iteration == 0)
          {
            complete_system_matrix.copy_from(system_matrix);
            complete_system_rhs = system_rhs;
          }

        solve();
        update_solution_and_constraints();
        output_results(iteration);

        if (active_set == active_set_old)
          break;

        active_set_old = active_set;

        std::cout << std::endl;
      }
  }
} // namespace Step41
// @sect3{The main function}

// 这就是主函数。它遵循所有其他主函数的模式。调用初始化MPI是因为我们在这个程序中建立线性求解器的Trilinos库需要它。

int main(int argc, char *argv[])
{
  try
    {
      using namespace dealii;
      using namespace Step41;

      Utilities::MPI::MPI_InitFinalize mpi_initialization(
        argc, argv, numbers::invalid_unsigned_int);

// 这个程序只能在串行中运行。否则，将抛出一个异常。

      AssertThrow(Utilities::MPI::n_mpi_processes(MPI_COMM_WORLD) == 1,
                  ExcMessage(
                    "This program can only be run in serial, use ./step-41"));

      ObstacleProblem<2> obstacle_problem;
      obstacle_problem.run();
    }
  catch (std::exception &exc)
    {
      std::cerr << std::endl
                << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Exception on processing: " << std::endl
                << exc.what() << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;

      return 1;
    }
  catch (...)
    {
      std::cerr << std::endl
                << std::endl
                << "----------------------------------------------------"
                << std::endl;
      std::cerr << "Unknown exception!" << std::endl
                << "Aborting!" << std::endl
                << "----------------------------------------------------"
                << std::endl;
      return 1;
    }

  return 0;
}


#include "lum_retinex.h"

#include <Eigen/SparseCore>
#include <Eigen/SparseCholesky>
#include <chrono>

namespace lum {

    // timing helper
	double get_s() {
		using namespace std::chrono;
		auto now = system_clock::now();
		system_clock::duration tse = now.time_since_epoch();
		return duration_cast<nanoseconds>(tse).count() / 1e9;
	}

    // Helper image processing routines
	void log(const float* inimg, float* outimg, int sz) {
#ifdef USE_MKL
		vsLn(sz, inimg, outimg);
#else
		for (int i = 0; i < sz; ++i) {
			float in = inimg[i];
			outimg[i] = std::logf(inimg[i] + 0.00001);
		}
#endif
	}

	void exp(const float* inimg, float* outimg, int sz) {
#ifdef USE_MKL
		vsExp(sz, inimg, outimg);
#else
		for (int i = 0; i < sz; ++i) {
			outimg[i] = std::expf(inimg[i]);
		}
#endif
	}
	void mean(const float* inimg, float* outimg, int outsz) {
		for (int i = 0; i < outsz; ++i) {
			outimg[i] = (inimg[i * 3] + inimg[i * 3 + 1] + inimg[i * 3 + 2]) / 3;
		}
	}

	// helpers for image indices
	class reflshadidx {
	public:
		reflshadidx(int w, int h)
		:m_w(w), m_h(h){}
		int reflidx(int x, int y) const {
			return m_w * y + x;
		}
		int shadidx(int x, int y) const {
			return m_w * m_h + reflidx(x, y);
		}
	private:
		int m_w;
		int m_h;
	};
	class imwrap {
	public:
		imwrap(const float* img, int w, int h)
		:m_w(w), m_h(h), m_img(img){ }

		float operator()(int x, int y) const {
			assert(x >= 0);
			assert(y >= 0);
			assert(x < m_w);
			assert(y < m_h);
			return m_img[m_w * y + x];
		}
	private:
		const float* m_img;
		int m_w;
		int m_h;
	};

	// forward declaration of internal functions
	void reflect_clamp(int w, int h, float* refl_in, float* shading_in);
	void preprocess(int w, int h, const float* img, float* logimg);
	void postprocess(int w, int h, float* refl_in, float* shading_in, float* refl_out, float* shading_out);
	Eigen::VectorXf makeB(float threshold, const float* im, int w, int h);

	using Triplet = Eigen::Triplet<float>;
	int nconstraints(int w, int h) {
		return w*h + 2 * w*(h - 1) + 2 * (w - 1) * h;
	}
	int nentries(int w, int h) {
		return 2 * (w*h + 2 * w*(h - 1) + 2 * (w - 1) * h);
	}
	std::vector<Triplet> makeTriplets(int w, int h) {
		reflshadidx I(w, h);
		printf("Assemble matrix.\n");
		double assemble_start = get_s();
		std::vector <Triplet> entries;
		int cit = 0;
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				if (x < w - 1) {
					// dxR(r, c) = -lR(r, c) + lR(r, c + 1)
					entries.push_back(Triplet(cit, I.reflidx(x, y), -1));
					entries.push_back(Triplet(cit, I.reflidx(x + 1, y), +1));
					cit++;

					// dxS(r, c) = -lS(r, c) + lS(r, c + 1)
					entries.push_back(Triplet(cit, I.shadidx(x, y), -1));
					entries.push_back(Triplet(cit, I.shadidx(x + 1, y), +1));
					cit++;
				}
				if (y < h - 1) {
					entries.push_back(Triplet(cit, I.reflidx(x, y), -1));
					entries.push_back(Triplet(cit, I.reflidx(x, y + 1), +1));
					cit++;

					entries.push_back(Triplet(cit, I.shadidx(x, y), -1));
					entries.push_back(Triplet(cit, I.shadidx(x, y + 1), +1));
					cit++;
					// dyR(r, c) = -lR(r, c) + lR(r + 1)
					// dyS(r, c) = -lS(r, c) + lS(r + 1)
				}

				// reflectance plus shading (log space) == final image
				entries.push_back(Triplet(cit, I.reflidx(x, y), 1));
				entries.push_back(Triplet(cit, I.shadidx(x, y), 1));
				cit++;
			}
		}
		assert(entries.size() == nentries(w, h));
		double assemble_end = get_s();
		printf("Makemtx took %.1fms\n", (assemble_end- assemble_start) * 1000);
		return entries;
	}

	// to log domain
	void preprocess(int w, int h, const float* img, float* logimg) {
		printf("Start Preprocessing.\n");
		double preprocess_start = get_s();
		log(img, logimg, w*h);
		double preprocess_b_end = get_s();
		printf("Preprocess took %.1fms\n", (preprocess_b_end - preprocess_start) * 1000);
	}

	// back to linear
	void postprocess(int w, int h, float* refl_in, float* shading_in, float* refl_out, float* shading_out) {
		printf("Post process start.\n");
		double postprocess_start = get_s();
		reflect_clamp(w, h, refl_in, shading_in);
		exp(refl_in, refl_out, w*h);
		exp(shading_in, shading_out, w*h);
		double postprocess_end = get_s();
		printf("Postprocess took %.1fms\n", (postprocess_end - postprocess_start)*1000 );
	}

	Eigen::VectorXf makeB(float threshold, const float* im, int w, int h) {
		Eigen::VectorXf b(nconstraints(w, h));
		double assemble_b_start = get_s();
		imwrap I(im, w, h);
		int cit = 0;
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				if (x < w - 1) {
					float dx = -I(x, y) + I(x + 1, y);
					float dxR;
					float dxS;
					if (std::abs(dx) > threshold) {
						dxR = dx;
						dxS = 0;
					} else {
						dxR = 0;
						dxS = dx;
					}
					// dxR(r, c) = -lR(r, c) + lR(r, c + 1)
					b(cit++) = dxR;

					// dxS(r, c) = -lS(r, c) + lS(r, c + 1)
					b(cit++) = dxS;
				}
				if (y < h - 1) {
					float dy = -I(x, y) + I(x, y + 1);
					float dyR;
					float dyS;
					if (std::abs(dy) > threshold) {
						dyR = dy;
						dyS = 0;
					} else {
						dyR = 0;
						dyS = dy;
					}
					b(cit++) = dyR;
					b(cit++) = dyS;
					// dyR(r, c) = -lR(r, c) + lR(r + 1)
					// dyS(r, c) = -lS(r, c) + lS(r + 1)
				}
				// reflectance plus shading (log space) == final image
				b(cit++) = I(x, y);
			}
		}
		return b;
	}

	// operates on log-reflectance
	// makes sure that log-reflectane is less than 0
	void reflect_clamp(int w, int h, float* refl_in, float* shading_in) {

		float max_reflectance = -FLT_MIN;
		float min_reflectance = FLT_MAX;
		int nancount = 0;
		int infcount = 0;
		for (int i = 0; i < w * h; ++i) {
			if (refl_in[i] > max_reflectance) {
				max_reflectance = refl_in[i];
			}
			if (refl_in[i] < min_reflectance) {
				min_reflectance = refl_in[i];
			}
		}

		if (max_reflectance > 0) {
			for (int i = 0; i < w * h; ++i) {
				refl_in[i] -= max_reflectance;
			}
			for (int i = 0; i < w * h; ++i) {
				shading_in[i] += max_reflectance;
			}
		}
	}

	void retinex(float threshold, const float* im, int w, int h, float* reflectance, float* shading) {
		assert(reflectance);
		assert(shading);
		assert(im);
		retinex_decomp rdec(w, h);
		rdec.solve(threshold, im, reflectance, shading);
	}
	/* I would prefer solving direction Ax = b.
	  However, this doesn't work with Eigen's QR decomp (why?).
	  I solve A'Ax = A'b, using Cholesky, instead.
	*/
	retinex_decomp::retinex_decomp(int w, int h): m_w(w), m_h(h) {
		std::vector <Triplet> entries = makeTriplets(w, h);
		SpMat A(nconstraints(w, h), w * h * 2);
		A.setFromTriplets(entries.begin(), entries.end());
		m_At = A.transpose();
		{
			printf("factorize %d-by-%d matrix\n", A.cols(), A.cols());
			double decompose_start = get_s();
			m_solver.compute(m_At * A);
			double decompose_end = get_s();
			printf("factorize took %.1fms\n", (decompose_end - decompose_start) * 1000);
		}
	}
	void retinex_decomp::solve(float threshold, const float* im, float* reflectance, float* shading) {
		assert(reflectance);
		assert(shading);
		assert(im);

		std::vector<float> logimg(m_w*m_h);
		preprocess(m_w, m_h, im, logimg.data());
		Eigen::VectorXf b = makeB(threshold, logimg.data(), m_w, m_h);
		double solve_start = get_s();
		Eigen::VectorXf x = m_solver.solve(m_At * b);
		double solve_end = get_s();
		printf("solve took %.1fms\n", (solve_end - solve_start) * 1000);
		postprocess(m_w, m_h, x.data(), x.data() + m_w * m_h, reflectance, shading);
	}

	// does greyscale retinex with post processing.
	void retinex_decomp::solve_rgb(float threshold, const float* im, float* reflectance, float* shading) {
		assert(reflectance);
		assert(shading);
		assert(im);
		const int sz = m_w * m_h;

		std::vector<float> grayimg(sz);
		mean(im, grayimg.data(), grayimg.size());
		std::vector<float> graylog(sz);
		double before = get_s();
		log(grayimg.data(), graylog.data(), sz);
		double after = get_s();

		Eigen::VectorXf b = makeB(threshold, graylog.data(), m_w, m_h);
		double solve_start = get_s();
		Eigen::VectorXf x = m_solver.solve(m_At * b);
		double solve_end = get_s();

		float* log_shading = x.data() + sz;
		reflect_clamp(m_w, m_h, x.data(), log_shading);

		std::vector<float> rgb_logimg(3*sz);
		log(im, rgb_logimg.data(), rgb_logimg.size());

		// do gray to rgb conversion
		// log R = log I - log S
		for (int i = 0; i < sz; ++i) {
			rgb_logimg[i * 3 + 0] = rgb_logimg[i * 3 + 0] - log_shading[i];
			rgb_logimg[i * 3 + 1] = rgb_logimg[i * 3 + 1] - log_shading[i];
			rgb_logimg[i * 3 + 2] = rgb_logimg[i * 3 + 2] - log_shading[i];
		}

		exp(rgb_logimg.data(), reflectance, 3 * sz);
		exp(log_shading, shading, sz);
	}
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2014 Jose Aparicio

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

#include <ql/quantlib.hpp>

#include <boost/timer.hpp>
#include <boost/make_shared.hpp>
#include <boost/function.hpp>

#include <iostream>
#include <iomanip>

using namespace std;
using namespace QuantLib;

#ifdef BOOST_MSVC
#  ifdef QL_ENABLE_THREAD_SAFE_OBSERVER_PATTERN
#    include <ql/auto_link.hpp>
#    define BOOST_LIB_NAME boost_system
#    include <boost/config/auto_link.hpp>
#    undef BOOST_LIB_NAME
#    define BOOST_LIB_NAME boost_thread
#    include <boost/config/auto_link.hpp>
#    undef BOOST_LIB_NAME
#  endif
#endif

#if defined(QL_ENABLE_SESSIONS)
namespace QuantLib {

    Integer sessionId() { return 0; }

}
#endif

/* This sample code shows basic usage of a Latent variable model.
   The data and correlation problem presented is the same as in:
     'Modelling Dependent Defaults: Asset Correlations Are Not Enough!'
     Frey R., A. J. McNeil and M. A. Nyfeler RiskLab publications March 2001
*/
int main(int, char* []) {

    try {

        boost::timer timer;
        std::cout << std::endl;

        Calendar calendar = TARGET();
        Date todaysDate(19, March, 2014);
        // must be a business day
        todaysDate = calendar.adjust(todaysDate);

        Settings::instance().evaluationDate() = todaysDate;

        /* --------------------------------------------------------------
                        SET UP BASKET PORTFOLIO
        -------------------------------------------------------------- */
        // build curves and issuers into a basket of three names
        std::vector<Real> hazardRates(3, -std::log(1.-0.01));
        std::vector<std::string> names;
        for(Size i=0; i<hazardRates.size(); i++)
            names.push_back(std::string("Acme") +
                boost::lexical_cast<std::string>(i));
        std::vector<Handle<DefaultProbabilityTermStructure> > defTS;
        for(Size i=0; i<hazardRates.size(); i++)
            defTS.push_back(Handle<DefaultProbabilityTermStructure>(
                boost::make_shared<FlatHazardRate>(0, TARGET(), hazardRates[i],
                    Actual365Fixed())));
        std::vector<Issuer> issuers;
        for(Size i=0; i<hazardRates.size(); i++) {
            std::vector<QuantLib::Issuer::key_curve_pair> curves(1,
                std::make_pair(NorthAmericaCorpDefaultKey(
                    EURCurrency(), QuantLib::SeniorSec,
                    Period(), 1. // amount threshold
                    ), defTS[i]));
            issuers.push_back(Issuer(curves));
        }

        boost::shared_ptr<Pool> thePool = boost::make_shared<Pool>();
        for(Size i=0; i<hazardRates.size(); i++)
            thePool->add(names[i], issuers[i], NorthAmericaCorpDefaultKey(
                    EURCurrency(), QuantLib::SeniorSec, Period(), 1.));

        std::vector<DefaultProbKey> defaultKeys(hazardRates.size(),
            NorthAmericaCorpDefaultKey(EURCurrency(), SeniorSec, Period(), 1.));
        // Recoveries are irrelevant in this example but must be given as the
        //   lib stands.
        std::vector<boost::shared_ptr<RecoveryRateModel> > rrModels(
            hazardRates.size(), boost::make_shared<ConstantRecoveryModel>(
            ConstantRecoveryModel(0.5, SeniorSec)));
        boost::shared_ptr<Basket> theBskt = boost::make_shared<Basket>(
            todaysDate, names, std::vector<Real>(hazardRates.size(), 100.),
            thePool);
        /* --------------------------------------------------------------
                        SET UP JOINT DEFAULT EVENT LATENT MODELS
        -------------------------------------------------------------- */
        // Latent model factors, corresponds to the first entry in Table1 of the
        //   publication mentioned. It is a single factor model
        std::vector<std::vector<Real> > fctrsWeights(hazardRates.size(),
            std::vector<Real>(1, std::sqrt(0.1)));
        // --- Default Latent models -------------------------------------
        // Gaussian integrable joint default model:
        boost::shared_ptr<GaussianDefProbLM> lmG(new
            GaussianDefProbLM(fctrsWeights,
            LatentModelIntegrationType::GaussianQuadrature,
			GaussianCopulaPolicy::initTraits() // otherwise gcc screams
			));
        // Define StudentT copula
        // this is as far as we can be from the Gaussian, 2 T_3 factors:
        std::vector<Integer> ordersT(2, 3);
        TCopulaPolicy::initTraits iniT;
        iniT.tOrders = ordersT;
        // StudentT integrable joint default model:
        boost::shared_ptr<TDefProbLM> lmT(new TDefProbLM(fctrsWeights,
            // LatentModelIntegrationType::GaussianQuadrature,
            LatentModelIntegrationType::Trapezoid,
            iniT));

        // --- Default Loss models ----------------------------------------
        // Gaussian random joint default model:
        Size numSimulations = 100000;
        // Size numCoresUsed = 4;
        // Sobol, many cores
        boost::shared_ptr<DefaultLossModel> rdlmG(
            boost::make_shared<RandomDefaultLM<GaussianCopulaPolicy> >(lmG,
                std::vector<Real>(), numSimulations, 1.e-6, 2863311530));
        // StudentT random joint default model:
        boost::shared_ptr<DefaultLossModel> rdlmT(
            boost::make_shared<RandomDefaultLM<TCopulaPolicy> >(lmT,
            std::vector<Real>(), numSimulations, 1.e-6, 2863311530));

        /* --------------------------------------------------------------
                        DUMP SOME RESULTS
        -------------------------------------------------------------- */
        /* Default correlations in a T copula should be below those of the
        gaussian for the same factors.
        The calculations on the MC show dispersion on both copulas (thats
        ok) and too large values with very large dispersions on the T case.
        Computations are ok, within the dispersion, for the gaussian; compare
        with the direct integration in both cases.
        However the T does converge to the gaussian value for large value of
        the parameters.
        */
        Date calcDate(TARGET().advance(Settings::instance().evaluationDate(),
            Period(120, Months)));
        std::vector<Probability> probEventsTLatent, probEventsGLatent,
            probEventsTRandLoss, probEventsGRandLoss;
        //
        lmT->resetBasket(theBskt);
        for(Size numEvts=0; numEvts <=theBskt->size(); numEvts++) {
            probEventsTLatent.push_back(lmT->probAtLeastNEvents(numEvts,
                calcDate));
         }
        //
        lmG->resetBasket(theBskt);
        for(Size numEvts=0; numEvts <=theBskt->size(); numEvts++) {
            probEventsGLatent.push_back(lmG->probAtLeastNEvents(numEvts,
                calcDate));
         }
        //
        theBskt->setLossModel(rdlmT);
        for(Size numEvts=0; numEvts <=theBskt->size(); numEvts++) {
            probEventsTRandLoss.push_back(theBskt->probAtLeastNEvents(numEvts,
                calcDate));
         }
        //
        theBskt->setLossModel(rdlmG);
        for(Size numEvts=0; numEvts <=theBskt->size(); numEvts++) {
            probEventsGRandLoss.push_back(theBskt->probAtLeastNEvents(numEvts,
                calcDate));
         }

        Date correlDate = TARGET().advance(
            Settings::instance().evaluationDate(), Period(12, Months));
        std::vector<std::vector<Real> > correlsGlm, correlsTlm, correlsGrand,
            correlsTrand;
        //
        lmG->resetBasket(theBskt);
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            std::vector<Real> tmp;
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                tmp.push_back(lmG->defaultCorrelation(correlDate,
                    iName1, iName2));
            correlsGlm.push_back(tmp);
        }
        //
        lmT->resetBasket(theBskt);
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            std::vector<Real> tmp;
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                tmp.push_back(lmT->defaultCorrelation(correlDate,
                    iName1, iName2));
            correlsTlm.push_back(tmp);
        }
        //
        theBskt->setLossModel(rdlmG);
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            std::vector<Real> tmp;
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                tmp.push_back(theBskt->defaultCorrelation(correlDate,
                    iName1, iName2));
            correlsGrand.push_back(tmp);
        }
        //
        theBskt->setLossModel(rdlmT);
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            std::vector<Real> tmp;
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                tmp.push_back(theBskt->defaultCorrelation(correlDate,
                    iName1, iName2));
            correlsTrand.push_back(tmp);
        }

        std::cout <<
            " Gaussian versus T prob of extreme event (random and integrable)-"
            << std::endl;
        for(Size numEvts=0; numEvts <=theBskt->size(); numEvts++) {
            std::cout << "-Prob of " << numEvts << " events... " <<
                probEventsGLatent[numEvts] << " ** " <<
                probEventsTLatent[numEvts] << " ** " <<
                probEventsGRandLoss[numEvts]<< " ** " <<
                probEventsTRandLoss[numEvts]
            << std::endl;
        }

        cout << endl;
        cout << "-- Default correlations G,T,GRand,TRand--" << endl;
        cout << "-----------------------------------------" << endl;
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                cout <<
                    correlsGlm[iName1][iName2] << " , ";
            ;
                cout << endl;
        }
        cout << endl;
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                cout <<
                    correlsTlm[iName1][iName2] << " , ";
            ;
                cout << endl;
        }
        cout << endl;
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                cout <<
                    correlsGrand[iName1][iName2] << " , ";
            ;
                cout << endl;
        }
        cout << endl;
        for(Size iName1=0; iName1 <theBskt->size(); iName1++) {
            for(Size iName2=0; iName2 <theBskt->size(); iName2++)
                cout <<
                    correlsTrand[iName1][iName2] << " , ";
            ;
                cout << endl;
        }

        Real seconds  = timer.elapsed();
        Integer hours = Integer(seconds/3600);
        seconds -= hours * 3600;
        Integer minutes = Integer(seconds/60);
        seconds -= minutes * 60;
        cout << "Run completed in ";
        if (hours > 0)
            cout << hours << " h ";
        if (hours > 0 || minutes > 0)
            cout << minutes << " m ";
        cout << fixed << setprecision(0)
             << seconds << " s" << endl;

        return 0;
    } catch (exception& e) {
        cerr << e.what() << endl;
        return 1;
    } catch (...) {
        cerr << "unknown error" << endl;
        return 1;
    }
}


//==================================================================================================
/*
  Copyright 2017 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
//! [inverse_hyperbolic]
#include <boost/simd/hyperbolic.hpp>
#include <boost/simd/pack.hpp>
#include <boost/simd/function/enumerate.hpp>
#include <iostream>

namespace bs =  boost::simd;
using pack_ft =  bs::pack <float, 8>;

int main()
{
  pack_ft p = bs::enumerate<pack_ft>(-2.0f, 0.5f);
  std::cout << " p =  " << p << std::endl
            <<  " -> bs::acosh(p) =  " << bs::acosh(p) << std::endl
            <<  " -> bs::asinh(p) =  " << bs::asinh(p) << std::endl
            <<  " -> bs::atanh(p) =  " << bs::atanh(p) << std::endl
            <<  " -> bs::asech(p) =  " << bs::asech(p) << std::endl
            <<  " -> bs::acsch(p) =  " << bs::acsch(p) << std::endl;
  return 0;
}
//! [inverse_hyperbolic]


/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#include <iostream>
#include <fstream>
#include <string>
#include <boost/program_options.hpp>

#include <omp.h>

#include "dpSubclustersSphereMM.hpp"

using namespace Eigen;
using std:string;
using std::ofstream;
using std::ifstream;
namespace po = boost::program_options;

{

  // Declare the supported options.
  po::options_description desc("Allowed options");
  desc.add_options()
    ("help,h", "produce help message")
    ("N,N", po::value<int>(), "number of input datapoints")
    ("D,D", po::value<int>(), "number of dimensions of the data")
    ("T,T", po::value<int>(), "iterations")
    ("alpha,a", po::value<double>(), "alpha parameter of the DP")
    ("base,b", po::value<string>(),
      "which base measure to use (only NIW right now)")
    ("params,p", po::value< vector<double> >()->multitoken(),
      "parameters of the base measure")
    ("input,i", po::value<string>(),
      "path to input dataset .csv file (rows: dimensions; cols: different datapoints)")
    ("output,o", po::value<string>(),
      "path to output labels .csv file (rows: time; cols: different datapoints)")
    ;

  po::variables_map vm;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);

  if (vm.count("help")) {
    cout << desc << "\n";
    return 1;
  }

#ifndef NDEBUG
  uint32_t nThreads = 1;
#else
  uint32_t nThreads = omp_get_max_threads(); // has to match number of threads on computer
#endif
  // number of iterations
  uint32_t T=1000;
  if (vm.count("T")) T = vm["T"].as<int>();
  uint32_t N=100;
  if (vm.count("N")) N = vm["N"].as<int>();
  uint32_t D=2;
  if (vm.count("D")) D = vm["D"].as<int>();
  cout << "T="<<T<<endl;
  // DP alpha parameter
  double alpha = 0.1;
  if (vm.count("alpha")) alpha = vm["alpha"].as<double>();
  cout << "alpha="<<alpha<<endl;
  // which base distribution
  string base = "NIW";
  if(vm.count("base")) base = vm["base"].as<string>();

//  DpMM *dpmm;
//  if(!base.compare("NIW"))
//  {
    MatrixXd Delta(D-1,D-1);
    VectorXd theta(D);
    double nu = D+3.0;
    double kappa = D+3.0;
//    Delta << nu,0.0,0.0,nu;
//    theta << 0.0,0.0;
    if(vm.count("params"))
    {
      vector<double> params = vm["params"].as< vector<double> >();
      cout<<"params length="<<params.size()<<" D="<<D<<endl;
      nu = params[0];
      kappa = params[1];
      for(uint32_t i=0; i<D; ++i)
        theta(i) = params[2+i];
      for(uint32_t i=0; i<D-1; ++i)
        for(uint32_t j=0; j<D-1; ++j)
          Delta(i,j) = params[2+D-1+i+(D-1)*j];
      cout <<"nu="<<nu<<endl;
      cout <<"kappa="<<kappa<<endl;
      cout <<"theta="<<theta<<endl;
      cout <<"Delta="<<Delta<<endl;
    }
    niwSphere_sampled niw(D-1,kappa,nu, theta.data(),Delta.data());
//
//  }else{
//    cout<<"base "<<base<<" not supported"<<endl;
//    return 1;
//  }

  MatrixXd x(D,N);
  string pathIn ="";
  if(vm.count("input")) pathIn = vm["input"].as<string>();
  if (!pathIn.compare(""))
  {
    for(uint32_t i=0; i<N; ++i)
      if(i<N/2)
      {
        x.col(i) << VectorXd::Zero(D);
      }else{
        x.col(i) << 2.0*VectorXd::Ones(D);
      }
  }else{
    cout<<"loading data from "<<pathIn<<endl;
    ifstream fin(pathIn.data(),ifstream::in);
    for (uint32_t j=0; j<D; ++j)
      for (uint32_t i=0; i<N; ++i)
      {
        fin>>x(j,i);
      }
    //cout<<x<<endl;
  }
  string pathOut ="./labels.csv";
  if(vm.count("output"))
    pathOut = vm["output"].as<string>();
  cout<<"output to "<<pathOut<<endl;

  VectorXd z(N);
  z.setZero();
  cout<<"z.shape="<<z.size()<<endl;
  DpSubclustersSphereMM *dpmm = new DpSubclustersSphereMM(N,D,x.data(),
      z.data(), alpha, niw, nThreads,true,true);
  cout<<"-- init"<<endl;
  dpmm->initialize();

  ofstream fout(pathOut.data(),ofstream::out);
  for (uint32_t t=0; t<T; ++t)
  {
    cout<<"------------ t="<<t<<" -------------"<<endl;
    cout<<"-- sampling params"<<endl;
    dpmm->sample_params();
    cout<<"-- sampling superclusters"<<endl;
    dpmm->sample_superclusters();
    cout<<"-- sampling labels"<<endl;
    dpmm->sample_labels();

    if(t%50 == 0)
    {
      cout<<"-- random splits"<<endl;
      dpmm->propose_random_splits();
    }
    cout<<"-- random merges"<<endl;
    dpmm->propose_random_merges();
    cout<<"-- splits"<<endl;
    dpmm->propose_splits();

    cout<<"   logLike=\t"<<dpmm->joint_loglikelihood()<<endl;
    cout<<"   K=\t"<<dpmm->getK()<<endl;
    cout<<"   Nk=\t"<<dpmm->getNK()<<endl;

//    const VectorXi& z = dpmm->getLabels().transpose();
//    //cout<< "z_"<<t<<"= "<<z.transpose()<<endl;
    for (uint32_t i=0; i<z.size()-1; ++i)
      fout<<int(floor(z(i)))<<" ";
    fout<<int(floor(z(z.size()-1)))<<endl;
  }
  fout.close();
}


#ifndef VEC2D_HPP_INCLUDED
#define VEC2D_HPP_INCLUDED

/*! \file vector.hpp
    \ingroup common
    \brief vector data types
    \details this file imports the vector data types from boost.ublas for use under less verbose
            type names.
*/

#include <cassert>
#include <numeric>
#include <boost/numeric/ublas/vector.hpp>

//! \addtogroup common
//! \{

//! vector of doubles with constant maximum length \p N
template<int N>
using c_vector = boost::numeric::ublas::c_vector<double, N>;

/*! normal generic vector type with at most three entries.
    uses c_vector, i.e. preallocated space so it does not
    require any dynamic memory management to create these
    objects.
*/
using gen_vect = c_vector<3>;

/// integer vector of at most three dimension.
using int_vect = boost::numeric::ublas::c_vector<int, 3>;

/// cross product function, very ugly, not generic or easy to use.
/// \todo should be removed or significantly improved.
template<class Result, class Vector>
void crossProduct( Result& target, const Vector& v1, const Vector& v2)
{
    assert( target.size() == 3 );
    assert( v1.size() == 3);
    assert( v2.size() == 3 );
    target[0] =  v1[1]*v2[2] - v1[2]*v2[1];
    target[1] = -v1[0]*v2[2] + v1[2]*v2[0];
    target[2] = v1[0]*v2[1] - v1[1]*v2[0];
}

template<class Vec1, class Vec2>
double dotProduct(Vec1&& v1, Vec2&& v2) {
    using std::begin;
    using std::end;
    return std::inner_product(begin(v1), end(v1), begin(v2), 0.0);
}

//! \}

#endif // VEC2D_HPP_INCLUDED


/**
 * $Id$
 *
 * Software License Agreement (GNU General Public License)
 *
 *  Copyright (C) 2015:
 *
 *    Johann Prankl, prankl@acin.tuwien.ac.at
 *    Aitor Aldoma, aldoma@acin.tuwien.ac.at
 *
 *      Automation and Control Institute
 *      Vienna University of Technology
 *      Gusshausstraße 25-29
 *      1170 Vienn, Austria
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * @author Johann Prankl, Aitor Aldoma
 *
 */

#ifndef KP_INV_POSE_HPP
#define KP_INV_POSE_HPP

#include <v4r/common/rotation.h>
#include <Eigen/Dense>

namespace v4r {

inline void invPose(const Eigen::Matrix4f &pose, Eigen::Matrix4f &inv_pose) {
  inv_pose.setIdentity();
  inv_pose.topLeftCorner<3, 3>() = pose.topLeftCorner<3, 3>().transpose();
  inv_pose.block<3, 1>(0, 3) = -1 * (inv_pose.topLeftCorner<3, 3>() * pose.block<3, 1>(0, 3));
}

inline void invPose(const Eigen::Matrix4d &pose, Eigen::Matrix4d &inv_pose) {
  inv_pose.setIdentity();
  inv_pose.topLeftCorner<3, 3>() = pose.topLeftCorner<3, 3>().transpose();
  inv_pose.block<3, 1>(0, 3) = -1 * (inv_pose.topLeftCorner<3, 3>() * pose.block<3, 1>(0, 3));
}

/**
 * invPose6
 */
template <typename T1, typename T2, typename T3, typename T4>
inline void invPose6(const T1 r[3], const T2 t[3], T3 inv_r[3], T4 inv_t[3]) {
  inv_r[0] = T1(-1) * r[0];
  inv_r[1] = T1(-1) * r[1];
  inv_r[2] = T1(-1) * r[2];

  v4r::AngleAxisRotatePoint(inv_r, t, inv_t);

  inv_t[0] = T4(-1) * inv_t[0];
  inv_t[1] = T4(-1) * inv_t[1];
  inv_t[2] = T4(-1) * inv_t[2];
}

}  // namespace v4r

#endif


#include <iostream>
#include <armadillo>
using namespace std;
using namespace arma;

int main() {
    vec a = {1, 2};
    a.print();
    // mat a = {{1}, {2}};
    rowvec b = {2, 3};
    b.print();
    mat res = a * b ;

    res.print();
    (b * a).print();

    std::vector<double> v = {1, 2, 3};
    vec c = vec({1, 2, 3});
    c.print();
    cout << c(2) << endl;

    return 0;
}


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2015 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#include <iostream>

//#define BOOST_DISABLE_ASSERTS

#include <boost/timer/timer.hpp>

#include "dune/grid/config.h"
#include "dune/grid/uggrid.hh"

#include "fem/assemble.hh"
#include "fem/lagrangespace.hh"
#include "linalg/trivialpreconditioner.hh"
#include "linalg/jacobiPreconditioner.hh"
#include "linalg/uzawa.hh"
#include "linalg/direct.hh"
#include "linalg/cg.hh"
#include "io/vtk.hh"
#include "utilities/kaskopt.hh" // property_tree
#include "utilities/gridGeneration.hh"

using namespace Kaskade;
#include "stokes.hh"

int main(int argc, char *argv[])
{
  using namespace boost::fusion;

  int verbosity = 1;  // print to console if arguments are changed
  bool dump = false; // do not write properties into file
  std::unique_ptr<boost::property_tree::ptree> pt = getKaskadeOptions(argc, argv, verbosity, dump);

  std::cout << "Start stokes tutorial programm." << std::endl;

  boost::timer::cpu_timer totalTimer;

  constexpr int dim = 2;
  constexpr int uIdx = 0;
  constexpr int pIdx = 1;

  // command line parameters
  int refinements = getParameter(pt, "refinements", 5);
  int order = getParameter(pt, "order", 2); // order for velocity space, order for pressure space is (order-1)
  std::string empty;

  std::string s("names.type.");
  s += getParameter(pt, "solver.type", empty);
  bool direct = getParameter(pt, s, 0);

  s = "names.direct." + getParameter(pt, "solver.direct", empty);
  DirectType directType = static_cast<DirectType>(getParameter(pt, s, 4));  // 4: DirectType::UMFPACK3264

  std::cout << "original mesh shall be refined : " << refinements << " times" << std::endl;
  std::cout << "discretization order           : " << order << std::endl;
  std::cout << "direct solver                  : " << directType << std::endl;

  boost::timer::cpu_timer gridTimer;
  // grid generation
  using Grid = Dune::UGGrid<dim>;
  using H1Space = FEFunctionSpace<ContinuousLagrangeMapper<double,Grid::LeafGridView> >;
  using Spaces = vector<H1Space const*,H1Space const*>;
  using VariableDescriptions = vector<Variable<SpaceIndex<1>,Components<2>,VariableId<uIdx> >,
                                      Variable<SpaceIndex<0>,Components<1>,VariableId<pIdx> > >;
  using VariableSet = VariableSetDescription<Spaces,VariableDescriptions>;
  using CoefficientVectors = VariableSet::CoefficientVectorRepresentation<>::type;
  using Functional = StokesFunctional<double,VariableSet>;
  using Assembler = VariationalFunctionalAssembler<LinearizationAt<Functional> >;

  Dune::FieldVector<double,dim> x0(0.0), length(1.0);
  GridManager<Grid> gridManager( createRectangle<Grid>(x0,length,1.0));
  gridManager.globalRefine(refinements);
  std::cout << "computing time for generation of initial mesh: " << boost::timer::format(gridTimer.elapsed());

  // construct involved spaces.
  H1Space pressureSpace(gridManager,gridManager.grid().leafGridView(),order-1);
  H1Space velocitySpace(gridManager,gridManager.grid().leafGridView(),order);

  Spaces spaces(&pressureSpace,&velocitySpace);

  // construct variable list.
  std::string varNames[2] = { "u", "p" };

  VariableSet variableSet(spaces,varNames);

  // construct variational functional.
  Functional F;

  // construct Galerkin representation
  Assembler assembler(spaces);
  VariableSet::VariableSet x(variableSet);
  VariableSet::VariableSet dx(variableSet);

  size_t nnz = assembler.nnz(0,2,0,2,false);
  size_t dof = variableSet.degreesOfFreedom(0,2);
  std::cout << "overall degrees of freedom: " << dof << std::endl;
  std::cout << "(structurally) nonzero elements: " << nnz << std::endl;

  boost::timer::cpu_timer assembleTimer;
  assembler.assemble(linearization(F,x));
  std::cout << "computing time for assemble: " << boost::timer::format(assembleTimer.elapsed());

  if(direct)
  {
    CoefficientVectors solution(VariableSet::CoefficientVectorRepresentation<>::init(spaces));
    CoefficientVectors rhs(assembler.rhs());

    boost::timer::cpu_timer directTimer;
    // solve performing one Newton step
    directInverseOperator(AssembledGalerkinOperator<Assembler>(assembler),directType).applyscaleadd(-1.0,rhs,solution);
    std::cout << "computing time for directsolve: " << boost::timer::format(directTimer.elapsed());
    x.data = solution.data;
  }
  else // Uzawa Solver
  {
    using VectorOfU = VariableSet::CoefficientVectorRepresentation<uIdx,uIdx+1>::type;
    using VectorOfP = VariableSet::CoefficientVectorRepresentation<pIdx,pIdx+1>::type;
    using Assembler_UU = AssembledGalerkinOperator<Assembler,uIdx,uIdx+1,uIdx,uIdx+1>;
    using Assembler_PU = AssembledGalerkinOperator<Assembler,pIdx,pIdx+1,uIdx,uIdx+1>;
    using Assembler_UP = AssembledGalerkinOperator<Assembler,uIdx,uIdx+1,pIdx,pIdx+1>;
    using PreconAdapt = MatrixRepresentedOperator<MatrixAsTriplet<double>, VectorOfP, VectorOfP>;
    using UzSo = UzawaSolver<VectorOfU,VectorOfP>;

    Assembler_UU A(assembler);
    Assembler_PU B(assembler);
    Assembler_UP Bt(assembler);

    boost::timer::cpu_timer iterativeTimer;
    Dune::InverseOperatorResult res;
    const DefaultDualPairing<VectorOfU,VectorOfU> defaultScalarProduct{};
    StrakosTichyPTerminationCriterion<double> termination(1e-14,300);
    int lookAhead = getParameter(pt, "solver.lookAhead", 3);
    termination.setLookAhead(lookAhead);

    JacobiPreconditioner<Assembler_UU> jacobiPreconditioner(A);
    // inexact inner solver for upper left block
    //Dune::CGSolver<VectorOfU> cg(A,jacobiPreconditioner,1e-14,300,0);
    CG<VectorOfU,VectorOfU> cg(A,jacobiPreconditioner,defaultScalarProduct,termination,verbosity);
    TrivialPreconditioner<PreconAdapt> trivialPreconditioner;

    VectorOfU f(assembler.rhs<uIdx,uIdx+1>());
    VectorOfP g(assembler.rhs<pIdx,pIdx+1>());
    VectorOfU u(VariableSet::CoefficientVectorRepresentation<uIdx,uIdx+1>::init(spaces));
    VectorOfP p(VariableSet::CoefficientVectorRepresentation<pIdx,pIdx+1>::init(spaces));

    UzSo uzawa(A,cg,B,Bt,trivialPreconditioner,1e-4,100,2);

    UzSo::Domain solution(vector<VectorOfU,VectorOfP>(u,p));
    UzSo::Range rhs(vector<VectorOfU,VectorOfP>(f,g));
    rhs *= -1.0; // change sign as rhs is -'F while the assembler returns F'
    uzawa.apply(solution,rhs,res);

    std::cout << "computing time for iterative solve: " << boost::timer::format(iterativeTimer.elapsed());

    at_c<uIdx>(x.data).coefficients() = at_c<0>(at_c<uIdx>(solution.data).data);
    at_c<pIdx>(x.data).coefficients() = at_c<0>(at_c<pIdx>(solution.data).data);
  }

  boost::timer::cpu_timer outputTimer;
  writeVTKFile(x,"stokes", IoOptions().setOrder(order));
  std::cout << "graphical output finished, data in VTK format is written into file stokes.vtu \n";
  std::cout << "computing time for output: " << boost::timer::format(outputTimer.elapsed()) << "\n";

  std::cout << "total computing time: " << boost::timer::format(totalTimer.elapsed()) << "\n";
  std::cout << "End stokes tutorial program" << std::endl;
}


/*
    MIT License

    Copyright (c) 2021 Zhepei Wang (wangzhepei@live.com)

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

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

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

#ifndef FLATNESS_HPP
#define FLATNESS_HPP

#include <Eigen/Eigen>

#include <cmath>

namespace flatness
{
    class FlatnessMap
    {
    public:
        inline void reset(const double &vehicle_mass,
                          const double &gravitational_acceleration,
                          const double &horitonral_drag_coeff,
                          const double &vertical_drag_coeff,
                          const double &parasitic_drag_coeff,
                          const double &speed_smooth_factor)
        {
            mass = vehicle_mass;
            grav = gravitational_acceleration;
            dh = horitonral_drag_coeff;
            dv = vertical_drag_coeff;
            cp = parasitic_drag_coeff;
            veps = speed_smooth_factor;

            return;
        }
        //通过速度加速度和机动度计算角速度和旋转
        inline void forward(const Eigen::Vector3d &vel,
                            const Eigen::Vector3d &acc,
                            const Eigen::Vector3d &jer,
                            const double &psi,
                            const double &dpsi,
                            double &thr,
                            Eigen::Vector4d &quat,
                            Eigen::Vector3d &omg)
        {
            double w0, w1, w2, dw0, dw1, dw2;

            v0 = vel(0);
            v1 = vel(1);
            v2 = vel(2);
            a0 = acc(0);
            a1 = acc(1);
            a2 = acc(2);
            cp_term = sqrt(v0 * v0 + v1 * v1 + v2 * v2 + veps);
            w_term = 1.0 + cp * cp_term;
            w0 = w_term * v0;
            w1 = w_term * v1;
            w2 = w_term * v2;
            dh_over_m = dh / mass;
            zu0 = a0 + dh_over_m * w0;
            zu1 = a1 + dh_over_m * w1;
            zu2 = a2 + dh_over_m * w2 + grav;
            zu_sqr0 = zu0 * zu0;
            zu_sqr1 = zu1 * zu1;
            zu_sqr2 = zu2 * zu2;
            zu01 = zu0 * zu1;
            zu12 = zu1 * zu2;
            zu02 = zu0 * zu2;
            zu_sqr_norm = zu_sqr0 + zu_sqr1 + zu_sqr2;
            zu_norm = sqrt(zu_sqr_norm);
            z0 = zu0 / zu_norm;
            z1 = zu1 / zu_norm;
            z2 = zu2 / zu_norm;
            ng_den = zu_sqr_norm * zu_norm;
            ng00 = (zu_sqr1 + zu_sqr2) / ng_den;
            ng01 = -zu01 / ng_den;
            ng02 = -zu02 / ng_den;
            ng11 = (zu_sqr0 + zu_sqr2) / ng_den;
            ng12 = -zu12 / ng_den;
            ng22 = (zu_sqr0 + zu_sqr1) / ng_den;
            v_dot_a = v0 * a0 + v1 * a1 + v2 * a2;
            dw_term = cp * v_dot_a / cp_term;
            dw0 = w_term * a0 + dw_term * v0;
            dw1 = w_term * a1 + dw_term * v1;
            dw2 = w_term * a2 + dw_term * v2;
            dz_term0 = jer(0) + dh_over_m * dw0;
            dz_term1 = jer(1) + dh_over_m * dw1;
            dz_term2 = jer(2) + dh_over_m * dw2;
            dz0 = ng00 * dz_term0 + ng01 * dz_term1 + ng02 * dz_term2;
            dz1 = ng01 * dz_term0 + ng11 * dz_term1 + ng12 * dz_term2;
            dz2 = ng02 * dz_term0 + ng12 * dz_term1 + ng22 * dz_term2;
            f_term0 = mass * a0 + dv * w0;
            f_term1 = mass * a1 + dv * w1;
            f_term2 = mass * (a2 + grav) + dv * w2;
            thr = z0 * f_term0 + z1 * f_term1 + z2 * f_term2;
            tilt_den = sqrt(2.0 * (1.0 + z2));
            tilt0 = 0.5 * tilt_den;
            tilt1 = -z1 / tilt_den;
            tilt2 = z0 / tilt_den;
            c_half_psi = cos(0.5 * psi);
            s_half_psi = sin(0.5 * psi);
            quat(0) = tilt0 * c_half_psi;
            quat(1) = tilt1 * c_half_psi + tilt2 * s_half_psi;
            quat(2) = tilt2 * c_half_psi - tilt1 * s_half_psi;
            quat(3) = tilt0 * s_half_psi;
            c_psi = cos(psi);
            s_psi = sin(psi);
            omg_den = z2 + 1.0;
            omg_term = dz2 / omg_den;
            omg(0) = dz0 * s_psi - dz1 * c_psi -
                     (z0 * s_psi - z1 * c_psi) * omg_term;
            omg(1) = dz0 * c_psi + dz1 * s_psi -
                     (z0 * c_psi + z1 * s_psi) * omg_term;
            omg(2) = (z1 * dz0 - z0 * dz1) / omg_den + dpsi;

            return;
        }

        inline void backward(const Eigen::Vector3d &pos_grad,
                             const Eigen::Vector3d &vel_grad,
                             const double &thr_grad,
                             const Eigen::Vector4d &quat_grad,
                             const Eigen::Vector3d &omg_grad,
                             Eigen::Vector3d &pos_total_grad,
                             Eigen::Vector3d &vel_total_grad,
                             Eigen::Vector3d &acc_total_grad,
                             Eigen::Vector3d &jer_total_grad,
                             double &psi_total_grad,
                             double &dpsi_total_grad) const
        {
            double w0b, w1b, w2b, dw0b, dw1b, dw2b;
            double z0b, z1b, z2b, dz0b, dz1b, dz2b;
            double v_sqr_normb, cp_termb, w_termb;
            double zu_sqr_normb, zu_normb, zu0b, zu1b, zu2b;
            double zu_sqr0b, zu_sqr1b, zu_sqr2b, zu01b, zu12b, zu02b;
            double ng00b, ng01b, ng02b, ng11b, ng12b, ng22b, ng_denb;
            double dz_term0b, dz_term1b, dz_term2b, f_term0b, f_term1b, f_term2b;
            double tilt_denb, tilt0b, tilt1b, tilt2b, head0b, head3b;
            double cpsib, spsib, omg_denb, omg_termb;
            double tempb, tilt_den_sqr;

            tilt0b = s_half_psi * (quat_grad(3)) + c_half_psi * (quat_grad(0));
            head3b = tilt0 * (quat_grad(3)) + tilt2 * (quat_grad(1)) - tilt1 * (quat_grad(2));
            tilt2b = c_half_psi * (quat_grad(2)) + s_half_psi * (quat_grad(1));
            head0b = tilt2 * (quat_grad(2)) + tilt1 * (quat_grad(1)) + tilt0 * (quat_grad(0));
            tilt1b = c_half_psi * (quat_grad(1)) - s_half_psi * (quat_grad(2));
            tilt_den_sqr = tilt_den * tilt_den;
            tilt_denb = (z1 * tilt1b - z0 * tilt2b) / tilt_den_sqr + 0.5 * tilt0b;
            omg_termb = -((z0 * c_psi + z1 * s_psi) * (omg_grad(1))) -
                        (z0 * s_psi - z1 * c_psi) * (omg_grad(0));
            tempb = omg_grad(2) / omg_den;
            dpsi_total_grad = omg_grad(2);
            z1b = dz0 * tempb;
            dz0b = z1 * tempb + c_psi * (omg_grad(1)) + s_psi * (omg_grad(0));
            z0b = -(dz1 * tempb);
            dz1b = s_psi * (omg_grad(1)) - z0 * tempb - c_psi * (omg_grad(0));
            omg_denb = -((z1 * dz0 - z0 * dz1) * tempb / omg_den) -
                       dz2 * omg_termb / (omg_den * omg_den);
            tempb = -(omg_term * (omg_grad(1)));
            cpsib = dz0 * (omg_grad(1)) + z0 * tempb;
            spsib = dz1 * (omg_grad(1)) + z1 * tempb;
            z0b += c_psi * tempb;
            z1b += s_psi * tempb;
            tempb = -(omg_term * (omg_grad(0)));
            spsib += dz0 * (omg_grad(0)) + z0 * tempb;
            cpsib += -dz1 * (omg_grad(0)) - z1 * tempb;
            z0b += s_psi * tempb + tilt2b / tilt_den + f_term0 * (thr_grad);
            z1b += -c_psi * tempb - tilt1b / tilt_den + f_term1 * (thr_grad);
            dz2b = omg_termb / omg_den;
            z2b = omg_denb + tilt_denb / tilt_den + f_term2 * (thr_grad);
            psi_total_grad = c_psi * spsib + 0.5 * c_half_psi * head3b -
                             s_psi * cpsib - 0.5 * s_half_psi * head0b;
            f_term0b = z0 * (thr_grad);
            f_term1b = z1 * (thr_grad);
            f_term2b = z2 * (thr_grad);
            ng02b = dz_term0 * dz2b + dz_term2 * dz0b;
            dz_term0b = ng02 * dz2b + ng01 * dz1b + ng00 * dz0b;
            ng12b = dz_term1 * dz2b + dz_term2 * dz1b;
            dz_term1b = ng12 * dz2b + ng11 * dz1b + ng01 * dz0b;
            ng22b = dz_term2 * dz2b;
            dz_term2b = ng22 * dz2b + ng12 * dz1b + ng02 * dz0b;
            ng01b = dz_term0 * dz1b + dz_term1 * dz0b;
            ng11b = dz_term1 * dz1b;
            ng00b = dz_term0 * dz0b;
            jer_total_grad(2) = dz_term2b;
            dw2b = dh_over_m * dz_term2b;
            jer_total_grad(1) = dz_term1b;
            dw1b = dh_over_m * dz_term1b;
            jer_total_grad(0) = dz_term0b;
            dw0b = dh_over_m * dz_term0b;
            tempb = cp * (v2 * dw2b + v1 * dw1b + v0 * dw0b) / cp_term;
            acc_total_grad(2) = mass * f_term2b + w_term * dw2b + v2 * tempb;
            acc_total_grad(1) = mass * f_term1b + w_term * dw1b + v1 * tempb;
            acc_total_grad(0) = mass * f_term0b + w_term * dw0b + v0 * tempb;
            vel_total_grad(2) = dw_term * dw2b + a2 * tempb;
            vel_total_grad(1) = dw_term * dw1b + a1 * tempb;
            vel_total_grad(0) = dw_term * dw0b + a0 * tempb;
            cp_termb = -(v_dot_a * tempb / cp_term);
            tempb = ng22b / ng_den;
            zu_sqr0b = tempb;
            zu_sqr1b = tempb;
            ng_denb = -((zu_sqr0 + zu_sqr1) * tempb / ng_den);
            zu12b = -(ng12b / ng_den);
            tempb = ng11b / ng_den;
            ng_denb += zu12 * ng12b / (ng_den * ng_den) -
                       (zu_sqr0 + zu_sqr2) * tempb / ng_den;
            zu_sqr0b += tempb;
            zu_sqr2b = tempb;
            zu02b = -(ng02b / ng_den);
            zu01b = -(ng01b / ng_den);
            tempb = ng00b / ng_den;
            ng_denb += zu02 * ng02b / (ng_den * ng_den) +
                       zu01 * ng01b / (ng_den * ng_den) -
                       (zu_sqr1 + zu_sqr2) * tempb / ng_den;
            zu_normb = zu_sqr_norm * ng_denb -
                       (zu2 * z2b + zu1 * z1b + zu0 * z0b) / zu_sqr_norm;
            zu_sqr_normb = zu_norm * ng_denb + zu_normb / (2.0 * zu_norm);
            tempb += zu_sqr_normb;
            zu_sqr1b += tempb;
            zu_sqr2b += tempb;
            zu2b = z2b / zu_norm + zu0 * zu02b + zu1 * zu12b + 2 * zu2 * zu_sqr2b;
            w2b = dv * f_term2b + dh_over_m * zu2b;
            zu1b = z1b / zu_norm + zu2 * zu12b + zu0 * zu01b + 2 * zu1 * zu_sqr1b;
            w1b = dv * f_term1b + dh_over_m * zu1b;
            zu_sqr0b += zu_sqr_normb;
            zu0b = z0b / zu_norm + zu2 * zu02b + zu1 * zu01b + 2 * zu0 * zu_sqr0b;
            w0b = dv * f_term0b + dh_over_m * zu0b;
            w_termb = a2 * dw2b + a1 * dw1b + a0 * dw0b +
                      v2 * w2b + v1 * w1b + v0 * w0b;
            acc_total_grad(2) += zu2b;
            acc_total_grad(1) += zu1b;
            acc_total_grad(0) += zu0b;
            cp_termb += cp * w_termb;
            v_sqr_normb = cp_termb / (2.0 * cp_term);
            vel_total_grad(2) += w_term * w2b + 2 * v2 * v_sqr_normb + vel_grad(2);
            vel_total_grad(1) += w_term * w1b + 2 * v1 * v_sqr_normb + vel_grad(1);
            vel_total_grad(0) += w_term * w0b + 2 * v0 * v_sqr_normb + vel_grad(0);
            pos_total_grad(2) = pos_grad(2);
            pos_total_grad(1) = pos_grad(1);
            pos_total_grad(0) = pos_grad(0);

            return;
        }

    private:
        double mass, grav, dh, dv, cp, veps;

        double v0, v1, v2, a0, a1, a2, v_dot_a;
        double z0, z1, z2, dz0, dz1, dz2;
        double cp_term, w_term/*segma*||v||*/, dh_over_m;
        double zu_sqr_norm, zu_norm, zu0, zu1, zu2;
        double zu_sqr0, zu_sqr1, zu_sqr2, zu01, zu12, zu02;
        double ng00, ng01, ng02, ng11, ng12, ng22, ng_den;
        double dw_term, dz_term0, dz_term1, dz_term2, f_term0, f_term1, f_term2;
        double tilt_den, tilt0, tilt1, tilt2, c_half_psi, s_half_psi;
        double c_psi, s_psi, omg_den, omg_term;
    };
}

#endif


// -*- mode: c++; indent-tabs-mode: nil; -*-
//
// Paragraph
// Copyright (c) 2016-2019 Illumina, Inc.
// All rights reserved.
//
// 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
//
//

/**
 * Population statistics for genotype sets
 *
 * \author Sai Chen & Egor Dolzhenko & Peter Krusche
 * \email schen6@illumina.com & pkrusche@illumina.com & edolzhenko@illumina.com
 *
 */

#include "genotyping/PopulationStatistics.hh"

#include <algorithm>
#include <boost/math/distributions/chi_squared.hpp>
#include <numeric>

using std::vector;

namespace genotyping
{
PopulationStatistics::PopulationStatistics(GenotypeSet const& genotypes)
{
    num_valid_samples = 0;
    num_total_samples = static_cast<int>(genotypes.size());

    for (auto const& genotype : genotypes)
    {
        if (genotype.gt.empty())
        {
            continue;
        }
        num_valid_samples++;
        if (genotype_counts.find(genotype.gt) == genotype_counts.end())
        {
            genotype_counts[genotype.gt] = 1;
        }
        else
        {
            genotype_counts[genotype.gt]++;
        }

        for (auto const& gt : genotype.gt)
        {
            if (allele_counts.size() <= gt)
            {
                allele_counts.resize(gt + 1, 0);
            }
            allele_counts[gt]++;
        }
    }
}

Json::Value PopulationStatistics::toJson() const
{
    double hwe_p_chisq = getChisqPvalue();
    double hwe_p_fisher = -1;
    if (needFisherExactHWE())
    {
        hwe_p_fisher = getFisherExactPvalue();
    }
    double call_rate = getCallrate();

    Json::Value json_result;
    json_result["hwe"] = hwe_p_chisq;
    if (hwe_p_fisher == -1)
    {
        json_result["hwe_fisher"] = "";
    }
    else
    {
        json_result["hwe_fisher"] = hwe_p_fisher;
    }
    json_result["call_rate"] = call_rate;

    auto allele_frequencies = getAlleleFrequencies();
    json_result["allele_frequencies"] = Json::arrayValue;
    for (auto& a : allele_frequencies)
    {
        json_result["allele_frequencies"].append(a);
    }

    return json_result;
}

double PopulationStatistics::getChisqPvalue() const
{
    double chisq_val = 0;
    for (auto& gv : genotype_counts)
    {
        if (gv.first.size() != 2)
        {
            continue;
        }
        uint64_t h1 = gv.first[0];
        uint64_t h2 = gv.first[1];
        if (allele_counts[h1] == 0 || allele_counts[h2] == 0) // skip unobserved alleles
        {
            continue;
        }
        double e_count;
        if (h1 == h2)
        {
            e_count = ((double)allele_counts[h1] / num_valid_samples / 2)
                * ((double)allele_counts[h1] / num_valid_samples / 2) * num_valid_samples;
        }
        else
        {
            e_count = 2 * ((double)allele_counts[h1] / num_valid_samples / 2)
                * ((double)allele_counts[h2] / num_valid_samples / 2) * num_valid_samples;
        }

        double diff = e_count - gv.second;
        double norm_diff_square = (double)(diff * diff) / e_count;
        chisq_val += norm_diff_square;
    }
    boost::math::chi_squared chisq_distribution(1);
    double hwe_pval = 1 - boost::math::cdf(chisq_distribution, chisq_val);
    return hwe_pval;
}

/**
 * @return True if need to use Fisher's exact test for HWE P
 * For multi-alleles, always use chisq
 * For bi-allelic:
 *      if
 *          N<=30, or count of the rarest genotype <= 20, or rarest expected count <= 20, use fisher's exact
 *      else
 *          use chisq
 */
bool PopulationStatistics::needFisherExactHWE() const
{
    int num_observed_alleles = 0;
    for (auto a : allele_counts)
    {
        if (a > 0)
        {
            num_observed_alleles++;
        }
    }
    if (num_observed_alleles <= 1)
    {
        return false;
    }
    if (num_observed_alleles > 2)
    {
        return false;
    }

    if (num_valid_samples <= 30)
    {
        return true;
    }

    for (auto& g : genotype_counts)
    {
        if (g.second > 0 && g.second <= 20)
        {
            return true;
        }
    }

    auto minor_allele_index = minNonZeroAlleleIndex();
    double minor_allele_freq = (double)allele_counts[minor_allele_index] / 2 / num_valid_samples;
    if (minor_allele_freq * minor_allele_freq * num_valid_samples <= 20)
    {
        return true;
    }
    return false;
}

/**
 * from JE Wigginton 2005 AJHG
 */
double PopulationStatistics::getFisherExactPvalue() const
{
    auto minor_allele_index = minNonZeroAlleleIndex();
    auto p_major = std::max_element(allele_counts.begin(), allele_counts.end());
    int minor_allele_count = allele_counts[minor_allele_index];
    int major_allele_count = *p_major;

    GenotypeVector het_gv;
    het_gv.push_back(static_cast<uint64_t>(p_major - allele_counts.begin()));
    het_gv.push_back(static_cast<uint64_t>(minor_allele_index));
    std::sort(het_gv.begin(), het_gv.end());
    int observed_num_het = 0;
    for (auto& g : genotype_counts)
    {
        if (g.first.size() != 2)
        {
            continue;
        }
        if (g.first[0] == het_gv[0] && g.first[1] == het_gv[1])
        {
            observed_num_het = g.second;
            break;
        }
    }

    int num_expect_het = std::round(
        2 * ((double)minor_allele_count / num_valid_samples / 2) * ((double)major_allele_count / num_valid_samples / 2)
        * num_valid_samples);

    std::vector<double> scaled_pvals;
    double observe_scaled_pval = -1;

    // for num_het > expected het
    int prev_num_ref_hom = (minor_allele_count - num_expect_het) / 2;
    int prev_num_alt_hom = num_valid_samples - prev_num_ref_hom - num_expect_het;
    double prev_scaled_pval = 1;
    for (int num_het = num_expect_het; num_het <= minor_allele_count; num_het += 2)
    {
        if (num_het == num_expect_het)
        {
            scaled_pvals.push_back(1);
            continue;
        }
        int prev_num_het = num_het - 2;
        double iscale
            = prev_scaled_pval * (4 * prev_num_ref_hom * prev_num_alt_hom) / ((prev_num_het + 2) * (prev_num_het + 1));
        scaled_pvals.push_back(iscale);
        prev_scaled_pval = iscale;
        prev_num_ref_hom--;
        prev_num_alt_hom--;
        if (observe_scaled_pval == -1 && num_het == observed_num_het)
        {
            observe_scaled_pval = iscale;
        }
    }

    // for num_het < expected het
    prev_num_ref_hom = (minor_allele_count - num_expect_het) / 2;
    prev_num_alt_hom = num_valid_samples - prev_num_ref_hom - num_expect_het;
    prev_scaled_pval = 1;
    for (int num_het = num_expect_het; num_het >= 0; num_het -= 2)
    {
        if (num_het == num_expect_het)
        {
            continue;
        }
        int prev_num_het = num_het + 2;
        double iscale = prev_scaled_pval / 4 * prev_num_het / (prev_num_ref_hom + 1) * (prev_num_het - 1)
            / (prev_num_alt_hom + 1);
        scaled_pvals.push_back(iscale);
        prev_scaled_pval = iscale;
        prev_num_ref_hom++;
        prev_num_alt_hom++;
        if (observe_scaled_pval == -1 && num_het == observed_num_het)
        {
            observe_scaled_pval = iscale;
        }
    }

    // calculate HWE
    double hwe_scale_sum = 0;
    for (auto s : scaled_pvals)
    {
        if (s <= observe_scaled_pval)
        {
            hwe_scale_sum += s;
        }
    }
    double pval = hwe_scale_sum / std::accumulate(scaled_pvals.begin(), scaled_pvals.end(), (double)0);
    return pval;
}

/**
 * @return allele frequencies
 */
std::vector<double> PopulationStatistics::getAlleleFrequencies() const
{
    std::vector<double> result;
    const uint32_t sum = std::accumulate(allele_counts.begin(), allele_counts.end(), (uint32_t)0);
    result.reserve(allele_counts.size());
    for (auto ac : allele_counts)
    {
        if (sum > 0)
        {
            result.push_back(((double)ac) / sum);
        }
        else
        {
            result.push_back(0);
        }
    }
    return result;
}

/**
 * @return iterator to the lowest allele non-zero minor allele count
 */
size_t PopulationStatistics::minNonZeroAlleleIndex() const
{
    auto p_minor = std::min_element(allele_counts.begin(), allele_counts.end());
    if (*p_minor > 0)
    {
        return static_cast<size_t>(p_minor - allele_counts.begin());
    }
    p_minor = std::max_element(allele_counts.begin(), allele_counts.end());
    if (*p_minor == 0)
    {
        return static_cast<size_t>(0);
    }
    for (auto it = allele_counts.begin(); it != allele_counts.end(); it++)
    {
        if (*it < *p_minor)
        {
            p_minor = it;
        }
    }
    return static_cast<size_t>(p_minor - allele_counts.begin());
}
}