#include <cstdio>
#include <iostream>

#include "gflags/gflags.h"
#include "sndfile.h"
#include "CImg.h"
#include "ipp.h"
#include "tbb/tbb.h"
#include "tbb/pipeline.h"
#include "tbb/scalable_allocator.h"
#include "tbb/cache_aligned_allocator.h"
#include <boost/filesystem.hpp>
#include "bakuage/sndfile_wrapper.h"
#include "bakuage/dft.h"
#include "bakuage/memory.h"
#include "bakuage/utils.h"
#include "bakuage/window_func.h"
#include "bakuage/vector_math.h"

DEFINE_bool(quick_exit, true, "quick exit");
DEFINE_int32(worker_count, 0, "thread count");
DEFINE_string(input, "", "input wav file path");
DEFINE_string(output, "", "output directory or - for stdout");
DEFINE_string(background, "", "background image path");
DEFINE_string(foreground, "", "foreground image path");
DEFINE_int32(min_hz, 50, "min freq in hz");
DEFINE_int32(max_hz, 8000, "max freq in hz");
DEFINE_int32(band_count, 30, "band count");
DEFINE_int32(width, 640, "width of output video");
DEFINE_int32(height, 480, "height of output video");
DEFINE_double(fps, 30, "frame per sec");
DEFINE_double(window_sec, 0.04, "dft window size in sec");

namespace {

typedef float Float;
typedef float PixelType;
    typedef std::shared_ptr<bakuage::AlignedPodVector<char>> BufferPtr;

    void PrintMemoryUsage() {
        std::cerr << "Peak RSS(MB)\t" << bakuage::GetPeakRss() / (1024 * 1024)
        << "\tCurrent RSS(MB)\t" << bakuage::GetCurrentRss() / (1024 * 1024)
        << std::endl;
    }

// 動画の長さが音源の長さ以上になるようにする
void CalculateSpectrogram(Float *input, int channels, int samples, int sample_freq, std::vector<bakuage::AlignedPodVector<float>> *output) {
    using namespace bakuage;

    const double min_mel = bakuage::HzToMel(FLAGS_min_hz);
    const double max_mel = bakuage::HzToMel(FLAGS_max_hz);

    const int width = 2 * (sample_freq * FLAGS_window_sec / 2); // even
    const int spec_len = width / 2 + 1;
    bakuage::AlignedPodVector<float> window(width);
    bakuage::CopyHanning(width, window.begin());
    bakuage::AlignedPodVector<float> fft_input(width);
    bakuage::AlignedPodVector<std::complex<float>> fft_output(spec_len);
    bakuage::RealDft<float> dft(width);

    std::vector<bakuage::AlignedPodVector<float>> spectrogram;

    for (int pos_i = 0; pos_i < samples * FLAGS_fps / sample_freq; pos_i++) {
        const int pos = std::floor((double)pos_i / FLAGS_fps * sample_freq) - width / 2;

        bakuage::AlignedPodVector<std::complex<float>> complex_spec_mid(spec_len);
        bakuage::AlignedPodVector<std::complex<float>> complex_spec_side(spec_len);

        // window and fft
        for (int i = 0; i < channels; i++) {
            for (int j = 0; j < width; j++) {
                int k = pos + j;
                fft_input[j] = (0 <= k && k < samples) ? input[channels * k + i] * window[j] : 0;
            }
            dft.Forward(fft_input.data(), (float *)fft_output.data());
            for (int j = 0; j < spec_len; j++) {
                const auto spec = fft_output[j];
                complex_spec_mid[j] += spec;
                complex_spec_side[j] += spec * (2.0f * i - 1);
            }
        }

        // 3dB/oct スロープ補正 + エネルギー正規化
        // mean モードを使うので、ピンクノイズは-3dB/octになることに注意
        for (int j = 0; j < spec_len; j++) {
            const auto freq = 1.0 * j / width * sample_freq;
            const auto slope_scale = std::sqrt(freq); // linear空間なのでsqrt
            const auto normalize_scale = 1.0 / std::sqrt(width);
            const auto scale = slope_scale * normalize_scale;
            complex_spec_mid[j] *= scale;
            complex_spec_side[j] *= scale;
        }

        bakuage::AlignedPodVector<float> row(FLAGS_band_count);
        bakuage::AlignedPodVector<float> row_count(FLAGS_band_count);
        for (int j = 0; j < spec_len; j++) {
            const double compensation = (j == 0 || j == spec_len - 1) ? 1 : 2;
            const double freq = 1.0 * j / width * sample_freq;
            const int k = FLAGS_band_count * (bakuage::HzToMel(freq) - min_mel) / (max_mel - min_mel);
            if (0 <= k && k < row.size()) {
                row[k] += std::norm(complex_spec_mid[j]) * compensation;
                row_count[k] += 1;
            }
        }
        for (int j = 0; j < row.size(); j++) {
            row[j] /= 1e-37 + row_count[j];
        }
        spectrogram.emplace_back(std::move(row));
    }

    *output = std::move(spectrogram);
}

void overlay_4ch_on_3ch(const cimg_library::CImg<PixelType> &src, cimg_library::CImg<PixelType> *dest) {
    // CImgのメモリ配置
    /*
     T& operator()(const unsigned int x, const unsigned int y, const unsigned int z, const unsigned int c) {
     return _data[x + y*(ulongT)_width + z*(ulongT)_width*_height + c*(ulongT)_width*_height*_depth];
     }
     */

    bakuage::AlignedPodVector<float> one_minus_alpha(src.width());

    const double scale = 1.0 / 255.0;
    for (int c = 0; c < 3; c++) {
        for (int j = 0; j < src.height(); j++) {
#if 1
            const PixelType *src_row = &src(0, j, 0, c);
            const PixelType *src_alpha_row = &src(0, j, 0, 3);
            PixelType *dest_row = &(*dest)(0, j, 0, c);

            bakuage::VectorSubConstantRev(src_alpha_row, 255.0f, one_minus_alpha.data(), src.width());
            bakuage::VectorMulInplace(one_minus_alpha.data(), dest_row, src.width());
            bakuage::VectorMadInplace(src_alpha_row, src_row, dest_row, src.width());
            bakuage::VectorMulConstantInplace(scale, dest_row, src.width());
#else
#if 1
            // for integer
            for (int i = 0; i < src.width(); i++) {
                const PixelType alpha = src(i, j, 0, 3);
                (*dest)(i, j, 0, c) = ((*dest)(i, j, 0, c) * (255 - alpha) + src(i, j, 0, c) * alpha) >> 8;
            }
#else
            for (int i = 0; i < src.width(); i++) {
                const double alpha = src(i, j, 0, 3) * scale;
                (*dest)(i, j, 0, c) = (*dest)(i, j, 0, c) * (1 - alpha) + src(i, j, 0, c) * alpha;
            }
#endif
#endif
        }
    }
}

}

// wavを受け取って、標準出力にrawvideoを出力
int main(int argc, char* argv[]) {
    gflags::SetVersionString("1.0.0-oss");
    gflags::ParseCommandLineFlags(&argc, &argv, true);

    ippInit();
    const IppLibraryVersion *lib = ippGetLibVersion();
    std::cerr << "Ipp initialized " << lib->Name << " " << lib->Version << std::endl;
    PrintMemoryUsage();

    // TBBの初期化とか (ここで初期化しておくと、毎回初期化しなくても良いらしい)
    // https://www.xlsoft.com/jp/products/intel/perflib/tbb/41/tbb_userguide_lnx/reference/task_scheduler/task_scheduler_init_cls.htm
    tbb::task_scheduler_init tbb_init(FLAGS_worker_count ? FLAGS_worker_count : tbb::task_scheduler_init::default_num_threads());
    std::cerr << "TBB default_num_threads:" << tbb::task_scheduler_init::default_num_threads() << std::endl;
    PrintMemoryUsage();

    bakuage::SndfileWrapper infile;
    SF_INFO sfinfo = { 0 };

    const auto input_file_path = FLAGS_input;

    if ((infile.set(sf_open (input_file_path.c_str(), SFM_READ, &sfinfo))) == NULL) {
        fprintf(stderr, "Not able to open input file %s.\n", input_file_path.c_str());
        fprintf(stderr, "%s\n", sf_strerror(NULL));
        return 1;
    }

    // check format
    fprintf(stderr, "sfinfo.format 0x%08x.\n", sfinfo.format);
    switch (sfinfo.format & SF_FORMAT_TYPEMASK) {
        case SF_FORMAT_WAV:
        case SF_FORMAT_WAVEX:
            break;
        default:
            fprintf(stderr, "Not supported sfinfo.format 0x%08x.\n", sfinfo.format);
            return 2;
    }

    bakuage::AlignedPodVector<float> buffer(sfinfo.channels * sfinfo.frames);
    int read_size = sf_readf_float(infile.get(), buffer.data(), sfinfo.frames);
    fprintf(stderr, "%d samples read.\n", read_size);
    if (read_size != sfinfo.frames) {
        fprintf(stderr, "sf_readf_float error: %d %d\n", read_size, (int)sfinfo.frames);
        return 3;
    }

    // calculate spectrogram (energy)
    std::vector<bakuage::AlignedPodVector<float>> spectrogram;
    CalculateSpectrogram(buffer.data(), sfinfo.channels, sfinfo.frames, sfinfo.samplerate, &spectrogram);

    // log and normalize spectrogram
    {
        double min_x = 1e100;
        double max_x = -1e100;
        for (auto &spectrum: spectrogram) {
            for (auto &x: spectrum) {
                x = std::log(1e-7 + x);
                min_x = std::min<double>(min_x, x);
                max_x = std::max<double>(max_x, x);
            }
        }
        for (auto &spectrum: spectrogram) {
            for (auto &x: spectrum) {
                x = (x - min_x) / (max_x - min_x);
            }
        }
    }

    cimg_library::CImg<PixelType> background(FLAGS_width, FLAGS_height, 1, 3, 0);
    if (!FLAGS_background.empty()) {
        cimg_library::CImg<PixelType> tmp(FLAGS_background.c_str());
        tmp.resize(FLAGS_width, FLAGS_height);
        for (int i = 0; i < FLAGS_width; i++) {
            for (int j = 0; j < FLAGS_height; j++) {
                background(i, j, 0, 0) = tmp(i, j, 0, 0);
                background(i, j, 0, 1) = tmp(i, j, 0, 1);
                background(i, j, 0, 2) = tmp(i, j, 0, 2);
            }
        }
    }

    cimg_library::CImg<PixelType> foreground(FLAGS_width, FLAGS_height, 1, 4, 0);
    if (!FLAGS_foreground.empty()) {
        cimg_library::CImg<PixelType> tmp(FLAGS_foreground.c_str());
        tmp.resize(FLAGS_width, FLAGS_height);
        for (int i = 0; i < FLAGS_width; i++) {
            for (int j = 0; j < FLAGS_height; j++) {
                foreground(i, j, 0, 0) = tmp(i, j, 0, 0);
                foreground(i, j, 0, 1) = tmp(i, j, 0, 1);
                foreground(i, j, 0, 2) = tmp(i, j, 0, 2);
                foreground(i, j, 0, 3) = tmp.spectrum() < 4 ? 255 : tmp(i, j, 0, 3);
            }
        }
    }

    auto spectrogram_it = spectrogram.begin();
    const auto spectrogram_end = spectrogram.end();
    const auto filter1_func = [&spectrogram_it, spectrogram_end](tbb::flow_control& fc) -> bakuage::AlignedPodVector<float> * {
        if (spectrogram_it != spectrogram_end) {
            return &(*(spectrogram_it++));
        } else {
            fc.stop();
            return nullptr;
        }
    };
    const auto filter2_func = [&foreground, &background](bakuage::AlignedPodVector<float> *spectrum){
        cimg_library::CImg<PixelType> img(FLAGS_width, FLAGS_height, 1, 3, 0);

        // draw background image
        img = background;

        // draw spectrum
        cimg_library::CImg<PixelType> spec_img(FLAGS_width, FLAGS_height, 1, 4, 0);
        spec_img.fill((PixelType)0, 0, 0, 0);
        for (int i = 0; i < spectrum->size(); i++) {
            const int block_div = 24;
            const double center_x = FLAGS_width / 2;
            const double center_y = FLAGS_height * 0.7;
            const double spectrum_width = FLAGS_width * 0.5;
            const double spectrum_height = FLAGS_height * 0.3;
            const double spectrum_left = center_x - spectrum_width / 2;
            const double spectrum_bottom = center_y + spectrum_height / 2;
            const double x1 = spectrum_left + spectrum_width * i / spectrum->size();
            const double x2 = spectrum_left + spectrum_width * (i + 1) / spectrum->size();
            const double space_x = (x2 - x1) * 0.5;
            const int block_count = std::floor((*spectrum)[i] * block_div + 0.5);
            for (int j = 0; j < block_count; j++) {
                const double y1 = spectrum_bottom - spectrum_height * (j + 1) / block_div;
                const double y2 = spectrum_bottom - spectrum_height * j / block_div;
                const double space_y = (y2 - y1) * 0.5;
                const PixelType color[4] = { 255, 255, 255, 200 };
                spec_img.draw_rectangle(x1 + space_x / 2, y1 + space_y / 2, 0, x2 - space_x / 2, y2 - space_y / 2, 1, color);
            }
        }
        overlay_4ch_on_3ch(spec_img, &img);

        // alpha blend foreground image
        overlay_4ch_on_3ch(foreground, &img);

        const auto temp_path = (boost::filesystem::temp_directory_path() / boost::filesystem::unique_path()).native();
        // img.save_bmp(output_path.str().c_str()); // slow
        img.save_png(temp_path.c_str());
        BufferPtr buffer;
        bakuage::LoadDataFromFile(temp_path.c_str(), [&buffer](const char *data, size_t size) {
            buffer = std::make_shared<bakuage::AlignedPodVector<char>>(data, data + size);
        });
        boost::filesystem::remove(temp_path);
        return buffer;
    };
    int spectrogram_i = 0;
    const auto filter3_func = [&spectrogram_i](BufferPtr buffer) {
        if (FLAGS_output == "-") {
            const size_t n = std::fwrite(buffer->data(), 1, buffer->size(), stdout);
            if (n != buffer->size()) {
                throw std::logic_error("failed to write stdout");
            }
        } else {
            std::stringstream output_path;
            output_path << FLAGS_output << "/" << spectrogram_i++ << ".png";

            FILE* pf = std::fopen(output_path.str().c_str(), "wb");
            std::fwrite(buffer->data(), 1, buffer->size(), pf);
            std::fclose(pf);
        }
    };
    tbb::parallel_pipeline(256,
                           tbb::make_filter<void, bakuage::AlignedPodVector<float> *>(tbb::filter::serial, filter1_func)
                           & tbb::make_filter<bakuage::AlignedPodVector<float> *,BufferPtr>(tbb::filter::parallel, filter2_func)
                           & tbb::make_filter<BufferPtr, void>(tbb::filter::serial, filter3_func)
                           );
    std::fflush(stdout);

    PrintMemoryUsage();

    if (FLAGS_quick_exit) {
        // 普通に終了するとクラッシュする。多分thread_localとかのデストラクタ周り
        // CircleCI上で再現したので要注意
        // クラッシュ回避ついでに効率的に終了できるので使うが、根本的にバグも直したい
        // https://stackoverflow.com/questions/24821265/exiting-a-c-app-immediately
        std::cerr << "quick exiting: " << 0 << std::endl;
        std::_Exit(0);
    } else {
        return 0;
    }
}


/**************************************************************************
 *
 * (C) Copyright VMware, Inc 2010.
 * (C) Copyright John Maddock 2006.
 * Use, modification and distribution are subject to the
 * Boost Software License, Version 1.0. (See accompanying file
 * LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
 *
 **************************************************************************/

/*
 * This file allows to compute the minimax polynomial coefficients we use
 * for fast exp2/log2.
 *
 * How to use this source:
 *
 * - Download and build the NTL library from
 *   http://shoup.net/ntl/download.html , or install libntl-dev package if on
 *   Debian.
 *
 * - Download boost source code matching to your distro.
 *
 * - Goto libs/math/minimax and replace f.cpp with this file.
 *
 * - Build as
 *
 *   g++ -o minimax -I /path/to/ntl/include main.cpp f.cpp /path/to/ntl/src/ntl.a
 *
 * - Run as
 *
 *    ./minimax
 *
 * - For example, to compute exp2 5th order polynomial between [0, 1] do:
 *
 *    variant 0
 *    range 0 1
 *    order 5 0
 *    step 200
 *    info
 *
 *  and take the coefficients from the P = { ... } array.
 *
 * - To compute log2 4th order polynomial between [0, 1/9] do:
 *
 *    variant 1
 *    range 0 0.111111112
 *    order 4 0
 *    step 200
 *    info
 *
 * - For more info see
 * http://www.boost.org/doc/libs/1_47_0/libs/math/doc/sf_and_dist/html/math_toolkit/toolkit/internals2/minimax.html
 */

#define L22
#include <boost/math/bindings/rr.hpp>
#include <boost/math/tools/polynomial.hpp>

#include <cmath>

boost::math::ntl::RR exp2(const boost::math::ntl::RR& x)
{
      return exp(x*log(2.0));
}

boost::math::ntl::RR log2(const boost::math::ntl::RR& x)
{
      return log(x)/log(2.0);
}

boost::math::ntl::RR f(const boost::math::ntl::RR& x, int variant)
{
   switch(variant)
   {
   case 0:
      return exp2(x);

   case 1:
      return log2((1.0 + sqrt(x))/(1.0 - sqrt(x)))/sqrt(x);
   }

   return 0;
}

void show_extra(
   const boost::math::tools::polynomial<boost::math::ntl::RR>& n,
   const boost::math::tools::polynomial<boost::math::ntl::RR>& d,
   const boost::math::ntl::RR& x_offset,
   const boost::math::ntl::RR& y_offset,
   int variant)
{
   switch(variant)
   {
   default:
      // do nothing here...
      ;
   }
}


#include <CGAL/Real_timer.h>
#include <CGAL/Random.h>
#include <CGAL/Simple_cartesian.h>

#include <CGAL/Point_set_3.h>
#include <CGAL/Point_set_3/IO.h>

#include <CGAL/Shape_detection/Region_growing/Region_growing.h>
#include <CGAL/Shape_detection/Region_growing/Region_growing_on_point_set.h>

#include <boost/iterator/function_output_iterator.hpp>

#include <fstream>

using Kernel = CGAL::Simple_cartesian<double>;
using Point_3 = Kernel::Point_3;
using Vector_3 = Kernel::Vector_3;
using Point_2 = Kernel::Point_2;
using Vector_2 = Kernel::Vector_2;

using Point_set_3 = CGAL::Point_set_3<Point_3, Vector_3>;
using Point_set_2 = CGAL::Point_set_3<Point_2, Vector_2>;
using Point_map = Point_set_2::Point_map;
using Normal_map = Point_set_2::Vector_map;

namespace Shape_detection = CGAL::Shape_detection::Point_set;

using Neighbor_query = Shape_detection::K_neighbor_query
  <Kernel, Point_set_2, Point_map>;
using Region_type = Shape_detection::Least_squares_circle_fit_region
  <Kernel, Point_set_2, Point_map, Normal_map>;
using Sorting = Shape_detection::Least_squares_circle_fit_sorting
  <Kernel, Point_set_2, Neighbor_query, Point_map>;
using Region_growing = CGAL::Shape_detection::Region_growing
  <Point_set_2, Neighbor_query, Region_type, typename Sorting::Seed_map>;

int main (int argc, char** argv)
{
  std::ifstream ifile (argc > 1 ? argv[1] : "data/circles.ply");
  Point_set_3 points3;
  ifile >> points3;

  std::cerr << points3.size() << " points read" << std::endl;

  // Input should have normals
  assert (points3.has_normal_map());

  Point_set_2 points;
  points.add_normal_map();
  for (Point_set_3::Index idx : points3)
  {
    const Point_3& p = points3.point(idx);
    const Vector_3& n = points3.normal(idx);
    points.insert (Point_2 (p.x(), p.y()), Vector_2 (n.x(), n.y()));
  }

  // Default parameters for data/circles.ply
  const std::size_t k = 12;
  const double tolerance = 0.01;
  const double max_angle = 10.;
  const std::size_t min_region_size = 20;

  // No constraint on radius
  const double min_radius = 0.;
  const double max_radius = std::numeric_limits<double>::infinity();

  Neighbor_query neighbor_query(points, k, points.point_map());
  Region_type region_type(points, tolerance, max_angle, min_region_size,
                          min_radius, max_radius,
                          points.point_map(), points.normal_map());

  // Sort indices
  Sorting sorting(points, neighbor_query, points.point_map());
  sorting.sort();

  Region_growing region_growing(points, neighbor_query, region_type, sorting.seed_map());

  // Add maps to get colored output
  Point_set_3::Property_map<unsigned char>
    red = points3.add_property_map<unsigned char>("red", 0).first,
    green = points3.add_property_map<unsigned char>("green", 0).first,
    blue = points3.add_property_map<unsigned char>("blue", 0).first;

  CGAL::Random random;

  std::size_t nb_circles = 0;
  CGAL::Real_timer timer;
  timer.start();
  region_growing.detect
    (boost::make_function_output_iterator
     ([&](const std::vector<std::size_t>& region)
      {
        // Assign a random color to each region
        unsigned char r = static_cast<unsigned char>(random.get_int(64, 192));
        unsigned char g = static_cast<unsigned char>(random.get_int(64, 192));
        unsigned char b = static_cast<unsigned char>(random.get_int(64, 192));
        for (const std::size_t& idx : region)
        {
          red[idx] = r;
          green[idx] = g;
          blue[idx] = b;
        }
        ++ nb_circles;
      }));
  timer.stop();

  std::cerr << nb_circles << " circles detected in "
            << timer.time() << " seconds" << std::endl;

  // Save in colored_circles.ply
  std::ofstream out ("colored_circles.ply");
  CGAL::IO::set_binary_mode (out);
  out << points3;

  return EXIT_SUCCESS;
}


#include <Python.h>
#include <iostream>
#include <Eigen/Dense>

using namespace Eigen;

typedef struct
{
    PyObject_HEAD
    MatrixXd *matrix=nullptr;
} PyMatrixObject;

static PyObject *
PyMatrix_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    PyMatrixObject *self;
    self = (PyMatrixObject *)type->tp_alloc(type, 0);

    char *kwlist[] = {"width", "height", NULL};
    int width = 0;
    int height = 0;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "ii", kwlist,
                                     &width, &height))
    {
        Py_DECREF(self);
        return NULL;
    }
    if (width <= 0 or height <= 0)
    {
        PyErr_SetString(PyExc_ValueError, "The height and width must be greater than 0.");
        return NULL;
    }

    self->matrix = new MatrixXd(width, height);
    return (PyObject *)self;
}

static void
*PyMatrix_dealloc(PyObject *obj)
{
    delete ((PyMatrixObject *)obj)->matrix;
    Py_TYPE(obj)->tp_free(obj);
}

inline MatrixXd *ParseMatrix(PyObject *obj){
    return ((PyMatrixObject *)obj)->matrix;
}

inline PyObject *ReturnMatrix(MatrixXd *m, PyTypeObject *type){
    PyMatrixObject *obj = PyObject_NEW(PyMatrixObject, type);
    obj->matrix = m;
    return (PyObject *)obj;
}

static PyObject *
PyMatrix_add(PyObject *a, PyObject *b)
{
    MatrixXd *matrix_a = ParseMatrix(a);
    MatrixXd *matrix_b = ParseMatrix(b);

    if (matrix_a->cols() != matrix_b->cols() or matrix_a->rows() != matrix_b->rows()){
        PyErr_SetString(PyExc_ValueError, "The input matrix must be the same shape.");
        return NULL;
    }

    MatrixXd *matrix_c = new MatrixXd(matrix_a->cols(), matrix_b->rows());
    *matrix_c = *matrix_a + *matrix_b;

    return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *
PyMatrix_minus(PyObject *a, PyObject *b)
{
    MatrixXd *matrix_a = ParseMatrix(a);
    MatrixXd *matrix_b = ParseMatrix(b);

    if (matrix_a->cols() != matrix_b->cols() or matrix_a->rows() != matrix_b->rows()){
        PyErr_SetString(PyExc_ValueError, "The input matrix must be the same shape.");
        return NULL;
    }

    MatrixXd *matrix_c = new MatrixXd(matrix_a->cols(), matrix_b->rows());
    *matrix_c = *matrix_a + *matrix_b;
    return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *
PyMatrix_multiply(PyObject *a, PyObject *b)
{
    MatrixXd *matrix_a = ParseMatrix(a);
    MatrixXd *matrix_b = ParseMatrix(b);

    if (matrix_a->cols() != matrix_b->rows()){
        PyErr_SetString(PyExc_ValueError, "The colonm rank of matrix A must be the same as the row rank of matrix B.");
        return NULL;
    }
    MatrixXd *matrix_c = new MatrixXd(matrix_a->rows(), matrix_b->cols());
    *matrix_c = (*matrix_a) * (*matrix_b);
    return ReturnMatrix(matrix_c, a->ob_type);
}

static PyObject *PyMatrix_str(PyObject *a)
{
    MatrixXd *matrix = ParseMatrix(a);
    std::stringstream ss;
    ss << *matrix;
    return Py_BuildValue("s", ss.str().c_str());
}

static PyNumberMethods numberMethods = {
    PyMatrix_add,      //nb_add
    PyMatrix_minus,    //nb_subtract;
    PyMatrix_multiply, //nb_multiply
    nullptr,           //nb_remainder;
    nullptr,           //nb_divmod;
    nullptr,           // nb_power;
    nullptr,           // nb_negative;
    nullptr,           // nb_positive;
    nullptr,           // nb_absolute;
    nullptr,           // nb_bool;
    nullptr,           // nb_invert;
    nullptr,           // nb_lshift;
    nullptr,           // nb_rshift;
    nullptr,           // nb_and;
    nullptr,           // nb_xor;
    nullptr,           // nb_or;
    nullptr,           // nb_int;
    nullptr,           // nb_reserved;
    nullptr,           // nb_float;

    nullptr, // nb_inplace_add;
    nullptr, // nb_inplace_subtract;
    nullptr, // nb_inplace_multiply;
    nullptr, // nb_inplace_remainder;
    nullptr, // nb_inplace_power;
    nullptr, // nb_inplace_lshift;
    nullptr, // nb_inplace_rshift;
    nullptr, // nb_inplace_and;
    nullptr, // nb_inplace_xor;
    nullptr, // nb_inplace_or;

    nullptr, // nb_floor_divide;
    nullptr, // nb_true_divide;
    nullptr, // nb_inplace_floor_divide;
    nullptr, // nb_inplace_true_divide;

    nullptr, // nb_index;

    nullptr, //nb_matrix_multiply;
    nullptr  //nb_inplace_matrix_multiply;

};

PyObject *PyMatrix_data(PyObject *self, void *closure)
{

    PyMatrixObject *obj = (PyMatrixObject *)self;
    Py_ssize_t width = obj->matrix->cols();
    Py_ssize_t height = obj->matrix->rows();

    PyObject *list = PyList_New(height);
    for (int i = 0; i < height; i++)
    {
        PyObject *internal = PyList_New(width);

        for (int j = 0; j < width; j++)
        {
            PyObject *value = PyFloat_FromDouble((*obj->matrix)(i, j));
            PyList_SetItem(internal, j, value);
        }

        PyList_SetItem(list, i, internal);
    }
    return list;
}

PyObject *PyMatrix_rows(PyObject *self, void *closure)
{
    PyMatrixObject *obj = (PyMatrixObject *)self;
    return Py_BuildValue("i", obj->matrix->rows());
}

PyObject *PyMatrix_cols(PyObject *self, void *closure)
{
    PyMatrixObject *obj = (PyMatrixObject *)self;
    return Py_BuildValue("i", obj->matrix->cols());
}

static PyGetSetDef MatrixGetSet[] = {
    {"data", (getter)PyMatrix_data, nullptr, nullptr},
    {"row", (getter)PyMatrix_rows, nullptr, nullptr},
    {"colunm", (getter)PyMatrix_cols, nullptr, nullptr},
    {nullptr}};

PyObject *PyMatrix_tolist(PyObject *self, PyObject *args)
{
    return PyMatrix_data(self, nullptr);
}

static PyMethodDef MatrixMethods[] = {
    {"to_list", (PyCFunction)PyMatrix_tolist, METH_VARARGS, "Return the matrix data to a list object."},
    {nullptr}};

static PyTypeObject MatrixType = {
    PyVarObject_HEAD_INIT(nullptr, 0) "matrix.Matrix", /* tp_name */
    sizeof(PyMatrixObject),                            /* tp_basicsize */
    0,                                                 /* tp_itemsize */
    (destructor)PyMatrix_dealloc,                      /* tp_dealloc */
    nullptr,                                           /* tp_print */
    nullptr,                                           /* tp_getattr */
    nullptr,                                           /* tp_setattr */
    nullptr,                                           /* tp_reserved */
    nullptr,                                           /* tp_repr */
    &numberMethods,                                    /* tp_as_number */
    nullptr,                                           /* tp_as_sequence */
    nullptr,                                           /* tp_as_mapping */
    nullptr,                                           /* tp_hash  */
    nullptr,                                           /* tp_call */
    PyMatrix_str,                                      /* tp_str */
    nullptr,                                           /* tp_getattro */
    nullptr,                                           /* tp_setattro */
    nullptr,                                           /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,          /* tp_flags */
    "Coustom matrix class.",                           /* tp_doc */
    nullptr,                                           /* tp_traverse */
    nullptr,                                           /* tp_clear */
    nullptr,                                           /* tp_richcompare */
    0,                                                 /* tp_weaklistoffset */
    nullptr,                                           /* tp_iter */
    nullptr,                                           /* tp_iternext */
    MatrixMethods,                                     /* tp_methods */
    nullptr,                                           /* tp_members */
    MatrixGetSet,                                      /* tp_getset */
    nullptr,                                           /* tp_base */
    nullptr,                                           /* tp_dict */
    nullptr,                                           /* tp_descr_get */
    nullptr,                                           /* tp_descr_set */
    0,                                                 /* tp_dictoffset */
    nullptr,                                           /* tp_init */
    nullptr,                                           /* tp_alloc */
    PyMatrix_new                                       /* tp_new */
};

static PyObject *PyMatrix_ones(PyObject *self, PyObject *args, PyObject *kwargs)
{
    PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
    m->matrix->setOnes();
    return (PyObject *)m;
}

static PyObject *PyMatrix_zeros(PyObject *self, PyObject *args, PyObject *kwargs)
{
    PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
    m->matrix->setZero();
    return (PyObject *)m;
}

static PyObject *PyMatrix_random(PyObject *self, PyObject *args, PyObject *kwargs)
{
    PyMatrixObject *m = (PyMatrixObject *)PyMatrix_new(&MatrixType, args, kwargs);
    m->matrix->setRandom();
    return (PyObject *)m;
}

static PyObject *PyMatrix_matrix(PyObject *self, PyObject *args)
{
    PyObject *data = nullptr;
    if (!PyArg_ParseTuple(args, "O", &data))
    {
        PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 1");
        return nullptr;
    }
    if (!PyList_Check(data))
    {
        PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 2");
        return nullptr;
    }
    int height = PyList_GET_SIZE(data);
    if (height <= 0)
    {
        PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 2");
        return nullptr;
    }
    PyObject *list = PyList_GET_ITEM(data, 0);
    if (!PyList_Check(list))
    {
        PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 3");
        return nullptr;
    }
    int width = PyList_GET_SIZE(list);
    MatrixXd *p_mat = new MatrixXd(width, height);
    for (int i = 0; i < height; i++)
    {
        PyObject *list = PyList_GET_ITEM(data, i);
        if (!PyList_Check(list))
        {
            PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. 3");
            return nullptr;
        }
        int tmp = PyList_GET_SIZE(list);
        if (width != tmp)
        {
            PyErr_SetString(PyExc_ValueError, "Please pass a 2 dimensions list object. Each elements of it must be the same length.");
            return nullptr;
        }
        width = tmp;

        for (int j = 0; j < width; j++)
        {
            PyObject *num = PyList_GET_ITEM(list, j);
            if (!PyFloat_Check(num))
            {
                PyErr_SetString(PyExc_ValueError, "Every elements of the matrix must float.");
                return nullptr;
            }
            (*p_mat)(i, j) = ((PyFloatObject *)num)->ob_fval;
        }
    }

    return ReturnMatrix(p_mat, &MatrixType);
}

static PyMethodDef matrixMethods[] = {
    {"ones", (PyCFunction)PyMatrix_ones, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with initial values one."},
    {"zeros", (PyCFunction)PyMatrix_zeros, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with initial values zero."},
    {"random", (PyCFunction)PyMatrix_random, METH_VARARGS | METH_KEYWORDS, "Return a new matrix with random values"},
    {"matrix", (PyCFunction)PyMatrix_matrix, METH_VARARGS, "Return a new matrix with given values"},
    {nullptr}};

static struct PyModuleDef module = {
    PyModuleDef_HEAD_INIT,
    "matrix",
    "Python interface for Matrix calculation",
    -1,
    matrixMethods};

PyObject *initModule(void)
{
    PyObject *m;
    if (PyType_Ready(&MatrixType) < 0)
        return NULL;

    m = PyModule_Create(&module);
    if (m == NULL)
        return NULL;

    Py_INCREF(&MatrixType);
    if (PyModule_AddObject(m, "Matrix", (PyObject *)&MatrixType) < 0)
    {
        Py_DECREF(&MatrixType);
        Py_DECREF(m);
        return NULL;
    }

    return m;
}


/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010-2020, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 *  @file  TimeOfArrivalExample.cpp
 *  @brief Track a moving object "Time of Arrival" measurements at 4
 * microphones.
 *  @author Frank Dellaert
 *  @author Jay Chakravarty
 *  @date March 2020
 */

#include <gtsam/nonlinear/LevenbergMarquardtOptimizer.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/nonlinear/expressions.h>
#include <gtsam_unstable/geometry/Event.h>
#include <gtsam_unstable/slam/TOAFactor.h>

#include <boost/format.hpp>

#include <vector>

using namespace std;
using namespace gtsam;

// units
static const double ms = 1e-3;
static const double cm = 1e-2;

// Instantiate functor with speed of sound value
static const TimeOfArrival kTimeOfArrival(330);

/* ************************************************************************* */
// Create microphones
vector<Point3> defineMicrophones() {
  const double height = 0.5;
  vector<Point3> microphones;
  microphones.push_back(Point3(0, 0, height));
  microphones.push_back(Point3(403 * cm, 0, height));
  microphones.push_back(Point3(403 * cm, 403 * cm, height));
  microphones.push_back(Point3(0, 403 * cm, 2 * height));
  return microphones;
}

/* ************************************************************************* */
// Create ground truth trajectory
vector<Event> createTrajectory(size_t n) {
  vector<Event> trajectory;
  double timeOfEvent = 10;
  // simulate emitting a sound every second while moving on straight line
  for (size_t key = 0; key < n; key++) {
    trajectory.push_back(
        Event(timeOfEvent, 245 * cm + key * 1.0, 201.5 * cm, (212 - 45) * cm));
    timeOfEvent += 1;
  }
  return trajectory;
}

/* ************************************************************************* */
// Simulate time-of-arrival measurements for a single event
vector<double> simulateTOA(const vector<Point3>& microphones,
                           const Event& event) {
  size_t K = microphones.size();
  vector<double> simulatedTOA(K);
  for (size_t i = 0; i < K; i++) {
    simulatedTOA[i] = kTimeOfArrival(event, microphones[i]);
  }
  return simulatedTOA;
}

/* ************************************************************************* */
// Simulate time-of-arrival measurements for an entire trajectory
vector<vector<double>> simulateTOA(const vector<Point3>& microphones,
                                   const vector<Event>& trajectory) {
  vector<vector<double>> simulatedTOA;
  for (auto event : trajectory) {
    simulatedTOA.push_back(simulateTOA(microphones, event));
  }
  return simulatedTOA;
}

/* ************************************************************************* */
// create factor graph
NonlinearFactorGraph createGraph(const vector<Point3>& microphones,
                                 const vector<vector<double>>& simulatedTOA) {
  NonlinearFactorGraph graph;

  // Create a noise model for the TOA error
  auto model = noiseModel::Isotropic::Sigma(1, 0.5 * ms);

  size_t K = microphones.size();
  size_t key = 0;
  for (auto toa : simulatedTOA) {
    for (size_t i = 0; i < K; i++) {
      graph.emplace_shared<TOAFactor>(key, microphones[i], toa[i], model);
    }
    key += 1;
  }
  return graph;
}

/* ************************************************************************* */
// create initial estimate for n events
Values createInitialEstimate(size_t n) {
  Values initial;

  Event zero;
  for (size_t key = 0; key < n; key++) {
    initial.insert(key, zero);
  }
  return initial;
}

/* ************************************************************************* */
int main(int argc, char* argv[]) {
  // Create microphones
  auto microphones = defineMicrophones();
  size_t K = microphones.size();
  for (size_t i = 0; i < K; i++) {
    cout << "mic" << i << " = " << microphones[i] << endl;
  }

  // Create a ground truth trajectory
  const size_t n = 5;
  auto groundTruth = createTrajectory(n);

  // Simulate time-of-arrival measurements
  auto simulatedTOA = simulateTOA(microphones, groundTruth);
  for (size_t key = 0; key < n; key++) {
    for (size_t i = 0; i < K; i++) {
      cout << "z_" << key << i << " = " << simulatedTOA[key][i] / ms << " ms"
           << endl;
    }
  }

  // Create factor graph
  auto graph = createGraph(microphones, simulatedTOA);

  // Create initial estimate
  auto initialEstimate = createInitialEstimate(n);
  initialEstimate.print("Initial Estimate:\n");

  // Optimize using Levenberg-Marquardt optimization.
  LevenbergMarquardtParams params;
  params.setAbsoluteErrorTol(1e-10);
  params.setVerbosityLM("SUMMARY");
  LevenbergMarquardtOptimizer optimizer(graph, initialEstimate, params);
  Values result = optimizer.optimize();
  result.print("Final Result:\n");
}
/* ************************************************************************* */


/*
 * Copyright 2009-2020 The VOTCA Development Team (http://www.votca.org)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
#include <libint2/initialize.h>
#define BOOST_TEST_MAIN

#define BOOST_TEST_MODULE density_integration_test

// Standard includes
#include <fstream>

// Third party includes
#include <boost/test/unit_test.hpp>

// Local VOTCA includes
#include "votca/tools/eigenio_matrixmarket.h"
#include "votca/xtp/density_integration.h"
#include "votca/xtp/orbitals.h"
#include "votca/xtp/vxc_grid.h"

using namespace votca::xtp;
using namespace std;

BOOST_AUTO_TEST_SUITE(density_integration_test)

AOBasis CreateBasis(const QMMolecule& mol) {

  BasisSet basis;
  basis.Load(std::string(XTP_TEST_DATA_FOLDER) +
             "/densityintegration/3-21G.xml");
  AOBasis aobasis;
  aobasis.Fill(basis, mol);
  return aobasis;
}

BOOST_AUTO_TEST_CASE(density_test) {
  libint2::initialize();
  QMMolecule mol("none", 0);

  mol.LoadFromFile(std::string(XTP_TEST_DATA_FOLDER) +
                   "/densityintegration/molecule.xyz");
  AOBasis aobasis = CreateBasis(mol);

  Eigen::MatrixXd dmat = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/densityintegration/dmat.mm");

  Vxc_Grid grid;
  grid.GridSetup("medium", mol, aobasis);
  DensityIntegration<Vxc_Grid> num(grid);

  double ntot = num.IntegrateDensity(dmat);
  BOOST_CHECK_CLOSE(ntot, 8.000000, 1e-5);

  Eigen::Vector3d pos = {3, 3, 3};

  BOOST_CHECK_CLOSE(num.IntegratePotential(pos), -1.543242, 1e-4);

  Eigen::Vector3d field = num.IntegrateField(pos);
  Eigen::Vector3d field_ref = {0.172802, 0.172802, 0.172802};
  bool field_check = field.isApprox(field_ref, 1e-5);
  if (!field_check) {
    std::cout << "field" << std::endl;
    std::cout << field.transpose() << std::endl;
    std::cout << "ref" << std::endl;
    std::cout << field_ref.transpose() << std::endl;
  }
  libint2::finalize();
}

BOOST_AUTO_TEST_CASE(gyration_test) {
  libint2::initialize();
  ofstream xyzfile("molecule.xyz");
  xyzfile << " 5" << endl;
  xyzfile << " methane" << endl;
  xyzfile << " C            1.000000     1.000000     1.000000" << endl;
  xyzfile << " H            1.629118     1.629118     1.629118" << endl;
  xyzfile << " H           0.370882    0.370882     1.629118" << endl;
  xyzfile << " H            1.629118    0.370882    0.370882" << endl;
  xyzfile << " H           0.370882     1.629118   0.370882" << endl;
  xyzfile.close();

  QMMolecule mol("none", 0);

  mol.LoadFromFile("molecule.xyz");
  AOBasis aobasis = CreateBasis(mol);

  Eigen::MatrixXd dmat = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/densityintegration/dmat.mm");
  Vxc_Grid grid;
  grid.GridSetup("medium", mol, aobasis);
  DensityIntegration<Vxc_Grid> num(grid);

  Gyrationtensor tensor = num.IntegrateGyrationTensor(dmat);
  BOOST_CHECK_CLOSE(tensor.mass, 8.0000005, 1e-5);

  Eigen::Vector3d dip_ref = Eigen::Vector3d::Zero();
  dip_ref << 1.88973, 1.88973, 1.88973;
  bool centroid_check = dip_ref.isApprox(tensor.centroid, 1e-5);
  BOOST_CHECK_EQUAL(centroid_check, true);
  if (!centroid_check) {
    std::cout << "centroid" << std::endl;
    std::cout << tensor.centroid.transpose() << std::endl;
    std::cout << "ref" << std::endl;
    std::cout << dip_ref.transpose() << std::endl;
  }
  Eigen::Matrix3d gyro_ref = Eigen::Matrix3d::Zero();
  gyro_ref << 0.596158, 2.85288e-12, 2.86873e-12, 2.85289e-12, 0.596158,
      2.87163e-12, 2.86874e-12, 2.87161e-12, 0.596158;
  bool gyro_check = gyro_ref.isApprox(tensor.gyration, 1e-5);
  BOOST_CHECK_EQUAL(gyro_check, true);
  if (!gyro_check) {
    std::cout << "gyro" << std::endl;
    std::cout << tensor.gyration << std::endl;
    std::cout << "ref" << std::endl;
    std::cout << gyro_ref << std::endl;
  }

  libint2::finalize();
}

BOOST_AUTO_TEST_SUITE_END()


#include<iostream>
#include <vector>
#include <tuple>
#include <float.h>
#include <numeric>
#include <Eigen/Dense>
#include "../nanoflann.hpp"
#include "gsec.hpp"
#include "utils.hpp"
#include "icp.hpp"
#include "../SO3.hpp"

using namespace nanoflann;
using namespace Geometry;

//
// Robust ICP
//

icp::Icp::Icp(icp::PointCloud &P, icp::PointCloud &M):
    resultSet(1),
    m_index(dim, std::cref(M), 10),
    p_index(dim, P, 10)
{
    // std::cout << "T: " << this->T << std::endl;
    // std::cout<<"&P = "<<&P<<std::endl;
    this->P = &P;
    // std::cout<<"&this->P = "<<(this->P)<<std::endl;
    this->Np = P.rows();
    this->M = &M;
    //
    this->constructKDTree();
    //
    this->resultSet.init(&(this->m_ret_index), &(this->m_out_dist_sqr));
    this->closestPtsFinder = &icp::Icp::closestpoints;
}

/**
 * if reciprocal is true, reject points based on eps,
 * otherwise based on a mutual closeness
*/
icp::Icp::Icp(icp::PointCloud &P, icp::PointCloud &M, bool reciprocal): icp::Icp(P, M)
{
    this->reciprocal = reciprocal;
    if(this->reciprocal) this->closestPtsFinder = &icp::Icp::closestpointsReciprocal;
}

/**
 * update iterations and distances vector
*/
void icp::Icp::init()
{
    this->iter++;
    this->initdists();
}

/**
 * clear ditances vector content
*/
void icp::Icp::initdists()
{
    // initialize heap
    this->distances.clear();
    // std::cout<< std::distance(distances.begin(), distances.begin()+Npo) << std::endl;
    // std::cout<< mse_cost(distances, xi) << std::endl;
    std::make_heap(this->distances.begin(), this->distances.end(), icp::comparePairs);
}

/**
 * find corresponces in the two pointclouds and compute distances.
 * kdtree is used
*/
void icp::Icp::closestpoints()
{
    // Step 1: Closest Point
    for (size_t i = 0; i < this->Np; i++)
    {
        this->resultSet.init(&(this->m_ret_index), &(this->m_out_dist_sqr));
        this->p_query_pt = {(*this->P)(i,0), (*this->P)(i,1), (*this->P)(i,2)};//{-15.7116, -10.9587, 0.650531};
        //
        this->m_index.index->findNeighbors(
            this->resultSet,
            &this->p_query_pt[0],
            nanoflann::SearchParams(10));

        // wrong pair rejection
        if(icp::rejectWrongPair(*this->M, this->m_ret_index, this->p_index, i)) continue;

        // update heap(distance)
        std::tuple<float, int, int> tmpDist = std::make_tuple(this->m_out_dist_sqr, i, this->m_ret_index);
        this->distances.push_back(tmpDist);
        std::push_heap(this->distances.begin(), this->distances.end(), icp::comparePairs);

    }
}

/**
 * using Reciprocal: Geometric criterion
*/
void icp::Icp::closestpointsReciprocal()
{
    // Step 1: Closest Point
    for (size_t i = 0; i < this->Np; i++)
    {

        this->resultSet.init(&(this->m_ret_index), &(this->m_out_dist_sqr));
        //auto query_ptt = P.row(i);
        this->p_query_pt = {(*this->P)(i,0), (*this->P)(i,1), (*this->P)(i,2)};
        //
        this->m_index.index->findNeighbors(
            this->resultSet,
            &this->p_query_pt[0],
            nanoflann::SearchParams(10));

        std::vector<float> m_query_pt = {(*this->M)(this->m_ret_index,0), (*this->M)(this->m_ret_index,1), (*this->M)(this->m_ret_index,2)};
        // wrong pair rejection
        if(icp::rejectWrongPair(
            m_query_pt, *this->P, this->p_index, i,
            this->eps))
            continue;

        std::tuple<float, int, int> tmpDist = std::make_tuple(this->m_out_dist_sqr, i, this->m_ret_index);
        this->distances.push_back(tmpDist);
        std::push_heap(this->distances.begin(), this->distances.end(), icp::comparePairs);

    }
    std::sort_heap(this->distances.begin(), this->distances.end(),  icp::comparePairs);
    // auto [ dist, pidx, midx ] = this->distances[0];
    // std::cout << "dist: " << dist << std::endl;
}

void icp::Icp::updateNpo()
{
    this->Npo = this->distances.size();//< Npo? distances.size():Npo;
    // std::cout << "Npo = " << Npo << std::endl;
}

/**
 * calculate squared distances and mean squared error (mse)
*/
void icp::Icp::calculateErrors()
{
    this->updateNpo();
    // this->Npo = this->distances.size();
    // Step 2: Squares
    this->Sts = std::accumulate(this->distances.begin(), this->distances.begin()+this->Npo, 0.0, icp::tupleAccumulateOP);
    // MSE
    this->e = this->Sts / this->Npo;

}

bool icp::Icp::convergencetest()
{
    // Step 3: Convergence test
    return (
        this->e <= this->mse_tol ||
        std::abs(this->e - this->e_prev)/this->e <= this->tol ||
        this->iter > this->Niter);// ||
        //std::abs(Sts-Stsprime) <= stsTolerance)

}

void icp::Icp::constructKDTree()
{
    // this->m_index(this->dim, std::cref(this->M), 10);
    // this->m_index.index->buildIndex();
    // // for Geometric criteria: rejecting wrong point pairs
    // this->p_index.index->buildIndex();
}

void icp::Icp::calculateMotion()
{
    // Step 4: Motion Calculation
    // https://igl.ethz.ch/projects/ARAP/svd_rot.pdf
    // select Npo points
    icp::PointCloud A(this->Npo, this->dim), B(this->Npo, this->dim);
    for (size_t i = 0; i < (this->Npo); i++)
    {
        auto [ dist, pidx, midx ] = this->distances[i];
        // std::cout<<"dist = "<<dist<<std::endl;
        A.row(i) = this->P->row(pidx);
        B.row(i) = this->M->row(midx);
    }

    Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic> umeyama = Eigen::umeyama(A.transpose(), B.transpose(), false);
    // std::cout << "The matrix umeyama is of size "
    //         << umeyama.rows() << "x" << umeyama.cols() << std::endl;
    // std::cout << "Transformation: " << std::endl << umeyama << std::endl;
    this->R = umeyama.block(0,0,3,3);
    this->t  = umeyama.block(0,3,3,1);

    // std::cout << "rotation error = " <<
    //             std::acos(((
    //                 this->R.transpose().array() * this->Rtotal.array()).matrix().trace() - 1) / 2.0
    //                 ) * 180.f/M_PI <<
    //             std::endl;

    this->Rtotal = this->R * this->Rtotal;
    this->T = umeyama * this->T;
    this->ttotal = this->t + this->ttotal;
    // auto centerA = A.colwise().mean();
    // auto centerB = B.colwise().mean();

    // A.rowwise() -= centerA;
    // B.rowwise() -= centerB;

    // icp::PointCloud H = A.transpose()*B;
    // Eigen::BDCSVD<Eigen::Matrix<float, -1, -1> > svd(H,Eigen::ComputeFullU | Eigen::ComputeFullV);

    // // Rotation
    // this->R = svd.matrixV() * svd.matrixU().transpose();
    // // Translation
    // this->t = centerB.transpose() - this->R * centerA.transpose();
}

void icp::Icp::transform()
{
    // affine transformation
    (*this->P) = (R * (this->P->transpose())).transpose();
    (*this->P).rowwise() += this->t.transpose();
    // t.col(0)
}

void icp::Icp::updateParameters()
{
    // update parameters
    // float* rng = gss(mse_cost, distances, 0.4, 1.0);
    // xi = (rng[0] + rng[1])/ 2.f; //0.5; // 1 -> ICP
    // Npo = xi*Np;

    (this->p_index).index->buildIndex();
    (this->e_prev) = this->e;
    (this->Stsprime) = this->Sts;
}

/**
 * calculate rotaion error using Log(R_gt*R).norm()
*/
float icp::Icp::rotationError(Eigen::Matrix3f R)
{
    // Eigen::Matrix3f Rt = this->Rtotal;
    Eigen::Matrix3f Rt = this->T.block(0,0,3,3);
    Eigen::Matrix3f rotE = R * (Rt).transpose();
    // std::cout << "Log(rotE): " << SO3::Log(rotE).transpose() << std::endl;
    // std::cout << "||Log(rotE)||_2: " << SO3::Log(rotE).norm() << std::endl;
    return SO3::Log(rotE).norm() * 180.f / M_PI;

}

void icp::Icp::run()
{
    do
    {
        // initialize heap
        this->init();

        // Step 1: Closest Point
        // if(this->reciprocal)
        //     this->closestpoints(this->eps);
        // else this->closestpoints();
        (this->*closestPtsFinder)();

        // Step 2: Trimmed Squares
        this->calculateErrors();

        // Step 3: Convergence test
        if(this->convergencetest()) break;

        // Step 4: Motion Calculation
        this->calculateMotion();

        // Step 5: Transformation
        this->transform();

        // Step 6: update parameters
        this->updateParameters();

    } while (true);
}

// getters
size_t icp::Icp::getIter()
{
    return this->iter;
}
size_t icp::Icp::getNpo()
{
    return this->Npo;
}
float icp::Icp::getSts()
{
    return this->Sts;
}
float icp::Icp::getMse()
{
    return this->e;
}

/**
 * return the accumulated rotation matrix
*/
icp::PointCloud icp::Icp::getR()
{
    return this->T.block(0,0,3,3);
    // return this->Rtotal;
}

/**
 * return the accumulated translation vector
*/
Eigen::Vector3f icp::Icp::getT()
{
    return this->T.block(0,3,3,1);
    // return   this->ttotal;
}
//
icp::Icp::~Icp()
{
    this->P = NULL;
    delete this->P;
    this->M = NULL;
    delete this->M;
    // this->closestPtsFinder = NULL;
    // delete this->closestPtsFinder;
}

//
// Trimmed Icp
//
// Constructors
icp::TrIcp::TrIcp(icp::PointCloud& P, icp::PointCloud& M): icp::Icp(P, M)
{
    this->Npo = this->xi * this->Np;
}

icp::TrIcp::TrIcp(icp::PointCloud& P, icp::PointCloud& M, float xi): icp::Icp(P, M)
{
    this->xi = xi;
    this->Npo = xi * this->Np;
}

// public methods
void icp::TrIcp::closestpoints()
{
    // Step 1: Closest Point
    for (size_t i = 0; i < this->Np; i++)
    {
        this->resultSet.init(&(this->m_ret_index), &(this->m_out_dist_sqr));
        this->p_query_pt = {(*this->P)(i,0), (*this->P)(i,1), (*this->P)(i,2)};//{-15.7116, -10.9587, 0.650531};
        //
        this->m_index.index->findNeighbors(
            this->resultSet,
            &this->p_query_pt[0],
            nanoflann::SearchParams(10));

        // update heap(distance)
        std::tuple<float, int, int> tmpDist = std::make_tuple(this->m_out_dist_sqr, i, this->m_ret_index);
        this->distances.push_back(tmpDist);
        std::push_heap(this->distances.begin(), this->distances.end(), icp::comparePairs);

    }
    std::sort_heap(this->distances.begin(), this->distances.end(),  icp::comparePairs);
    // auto [ dist, pidx, midx ] = this->distances[0];
    // std::cout << "dist: " << dist << std::endl;
}

void icp::TrIcp::updateNpo()
{
    if(!this->isUpdatedNpo)
    {
        float* rng = gss(mse_cost, this->distances, 0.4, 1.0);
        this->xi = (rng[0] + rng[1])/ 2.f; //0.5; // 1 -> ICP
        this->Npo = this->xi * this->Np;
        //
        // this->isUpdatedNpo = true;
    }
}

void icp::TrIcp::updateParameters()
{
    (this->e_prev) = this->e;
    (this->Stsprime) = this->Sts;
}

// destructor
icp::TrIcp::~TrIcp()
{
}


/*------------------------------------------------------*/
/* Biological Engineering: Problem Set 03               */
/* main program                                         */
/*                                                      */
/* file name    :   BiologicalEngineering_set3.cpp      */
/* compiler     :   gcc                                 */
/* Student ID   :   1526084                             */
/* author       :   Takaharu Nakajima                   */
/* date         :   2016.11.11                          */
/* memo         :   $B@8BN9)3XBh;02sL\$N(B1.1$B$+$i(B1.3$B$N(B      */
/*                  $B%W%m%0%i%`$r$^$H$a$?$b$N$K$J$j$^$9(B  */
/*                                                      */
/*                  $B0lItJQ?t$NDj5A$N$?$aJQ99$7$?ItJ,(B    */
/*                  ($BNc(B)                                */
/*                  RT0 -> RTa                          */
/*                  $B$3$N$h$&$K(B0, 1...E$B$r(Ba, b...E$B$K(B      */
/*                  $BJQ99$7$^$7$?(B                        */
/*                                                      */
/*------------------------------------------------------*/

#include <iostream>
#include <Eigen/Core>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

#define PI 3.1415926

double deg2rad(double degree)
{
    double radian = degree * PI / 180.0f;
    return radian;
}

int main()
{
    double deg1, deg2, deg3, deg4, deg5;
    double rad1, rad2, rad3, rad4, rad5;
    Matrix<double,4,4> RTa, aTb, bTc, cTd, dTe, eTf, fTg;
    Matrix<double,4,4> RTb, RTc, RTd, RTe, RTf, RTg;
    Vector3d RZa, RZb, RZc, RZd, RZe, RZf, RZg;
    Vector3d Rpa, Rpb, Rpc, Rpd, Rpe, Rpf, Rpg;
    Vector3d RpEa, RpEb, RpEc, RpEd, RpEe, RpEf, RpEg;

    /* Input */
    cout << "Angle input" << endl;
    cout << "deg1= ";
    cin >> deg1;
    cout << "deg2= ";
    cin >> deg2;
    cout << "deg3= ";
    cin >> deg3;
    cout << "deg4= ";
    cin >> deg4;
    cout << "deg5= ";
    cin >> deg5;

    /* Convert to radians */
    rad1 = deg2rad(deg1);
    rad2 = deg2rad(deg2);
    rad3 = deg2rad(deg3);
    rad4 = deg2rad(deg4);
    rad5 = deg2rad(deg5);

    /* Link parameters */
    RTa << 1, 0, 0, 0,
           0, 1, 0, 0,
           0, 0, 1, 1,
           0, 0, 0, 1;
    aTb << cos(rad1), -sin(rad1), 0, 0,
           sin(rad1),  cos(rad1), 0, 0,
                   0,          0, 1, 0,
                   0,          0, 0, 1;
    bTc << cos(rad2), -sin(rad2), 0, 0,
                   0,          0, -1, 0,
           sin(rad2),  cos(rad2), 0, 0,
                   0,          0, 0, 1;
    cTd << cos(rad3), -sin(rad3), 0, 1,
           sin(rad3),  cos(rad3), 0, 0,
                   0,          0, 1, 0,
                   0,          0, 0, 1;
    dTe << cos(rad4), -sin(rad4), 0, 1,
           sin(rad4),  cos(rad4), 0, 0,
                   0,          0, 1, 0,
                   0,          0, 0, 1;
    eTf << cos(rad5), -sin(rad5), 0, 0,
                   0,          0, -1, 0,
           sin(rad5),  cos(rad5), 0, 0,
                   0,          0, 0, 1;
    fTg << 1, 0, 0, 0,
           0, 1, 0, 0,
           0, 0, 1, 1,
           0, 0, 0, 1;

    /* Matrix calculation */
    RTb = RTa * aTb;
    RTc = RTb * bTc;
    RTd = RTc * cTd;
    RTe = RTd * dTe;
    RTf = RTe * eTf;
    RTg = RTf * fTg;

    RZa = RTa.block<3,1>(0,2);
    RZb = RTb.block<3,1>(0,2);
    RZc = RTc.block<3,1>(0,2);
    RZd = RTd.block<3,1>(0,2);
    RZe = RTe.block<3,1>(0,2);
    RZf = RTf.block<3,1>(0,2);
    RZg = RTg.block<3,1>(0,2);

    Rpa = RTa.block<3,1>(0,3);
    Rpb = RTb.block<3,1>(0,3);
    Rpc = RTc.block<3,1>(0,3);
    Rpd = RTd.block<3,1>(0,3);
    Rpe = RTe.block<3,1>(0,3);
    Rpf = RTf.block<3,1>(0,3);
    Rpg = RTg.block<3,1>(0,3);

    RpEa = Rpg - Rpa;
    RpEb = Rpg - Rpb;
    RpEc = Rpg - Rpc;
    RpEd = Rpg - Rpd;
    RpEe = Rpg - Rpe;
    RpEf = Rpg - Rpf;
    RpEg = Rpg - Rpg;

    /* Calculation result display */
    /* RT0 ~ 3TE */
    cout << "RT0=" << endl;
    cout << RTa << endl;
    cout << "0T1=" << endl;
    cout << aTb << endl;
    cout << "1T2=" << endl;
    cout << bTc << endl;
    cout << "2T3=" << endl;
    cout << cTd << endl;
    cout << "3T4=" << endl;
    cout << dTe << endl;
    cout << "4T5=" << endl;
    cout << eTf << endl;
    cout << "5T6=" << endl;
    cout << fTg << endl << endl;

    /* RZ0 ~ RZE */
    cout << "RZ0=" << endl;
    cout << RZa << endl;
    cout << "RZ1=" << endl;
    cout << RZb << endl;
    cout << "RZ2=" << endl;
    cout << RZc << endl;
    cout << "RZ3=" << endl;
    cout << RZd << endl;
    cout << "RZ4=" << endl;
    cout << RZe << endl;
    cout << "RZ5=" << endl;
    cout << RZf << endl;
    cout << "RZ6=" << endl;
    cout << RZg << endl;

    /* Rp0 ~ Rp1 */
    cout << "Rp0=" << endl;
    cout << Rpa << endl;
    cout << "Rp1=" << endl;
    cout << Rpb << endl;
    cout << "Rp2=" << endl;
    cout << Rpc << endl;
    cout << "Rp3=" << endl;
    cout << Rpd << endl;
    cout << "Rp4=" << endl;
    cout << Rpe << endl;
    cout << "Rp5=" << endl;
    cout << Rpf << endl;
    cout << "Rp6=" << endl;
    cout << Rpg << endl;

    /* RZ0 x RpE,0 ~ RZE x RpE,E */
    cout << "RZ0 x RpE,0=" << endl;
    cout << RZa.cross(RpEa) << endl;
    cout << "RZ1 x RpE,1=" << endl;
    cout << RZb.cross(RpEb) << endl;
    cout << "RZ2 x RpE,2=" << endl;
    cout << RZc.cross(RpEc) << endl;
    cout << "RZ3 x RpE,3=" << endl;
    cout << RZd.cross(RpEd) << endl;
    cout << "RZ4 x RpE,4=" << endl;
    cout << RZe.cross(RpEe) << endl;
    cout << "RZ5 x RpE,5=" << endl;
    cout << RZf.cross(RpEf) << endl;
    cout << "RZ6 x RpE,6=" << endl;
    cout << RZg.cross(RpEg) << endl;

    return 0;
}


/*
* Copyright (c) by CryptoLab inc.
* This program is licensed under a
* Creative Commons Attribution-NonCommercial 3.0 Unported License.
* You should have received a copy of the license along with this
* work.  If not, see <http://creativecommons.org/licenses/by-nc/3.0/>.
*/

#include "RingMultiplier.h"

#include <NTL/BasicThreadPool.h>
#include <NTL/tools.h>
#include <cmath>
#include <cstdlib>
#include <iterator>

RingMultiplier::RingMultiplier() {

	uint64_t primetest = (1ULL << pbnd) + 1;
	for (long i = 0; i < nprimes; ++i) {
		while(true) {
			primetest += M;
			if(primeTest(primetest)) {
				pVec[i] = primetest;
				break;
			}
		}
	}

	for (long i = 0; i < nprimes; ++i) {
		red_ss_array[i] = _ntl_general_rem_one_struct_build(pVec[i]);
		pInvVec[i] = inv(pVec[i]);
		prVec[i] = (static_cast<unsigned __int128>(1) << kbar2) / pVec[i];
		uint64_t root = findMthRootOfUnity(M, pVec[i]);
		uint64_t rootinv = invMod(root, pVec[i]);
		uint64_t NInv = invMod(N, pVec[i]);
		mulMod(scaledNInv[i], NInv, (1ULL << 32), pVec[i]);
		mulMod(scaledNInv[i], scaledNInv[i], (1ULL << 32), pVec[i]);
		scaledRootPows[i] = new uint64_t[N]();
		scaledRootInvPows[i] = new uint64_t[N]();
		uint64_t power = 1;
		uint64_t powerInv = 1;
		for (long j = 0; j < N; ++j) {
			uint32_t jprime = bitReverse(static_cast<uint32_t>(j)) >> (32 - logN);
			uint64_t rootpow = power;
			mulMod(scaledRootPows[i][jprime], rootpow,(1ULL << 32), pVec[i]);
			mulMod(scaledRootPows[i][jprime], scaledRootPows[i][jprime], (1ULL << 32), pVec[i]);
			uint64_t rootpowInv = powerInv;
			mulMod(scaledRootInvPows[i][jprime], rootpowInv, (1ULL << 32), pVec[i]);
			mulMod(scaledRootInvPows[i][jprime], scaledRootInvPows[i][jprime], (1ULL << 32), pVec[i]);
			mulMod(power, power, root, pVec[i]);
			mulMod(powerInv, powerInv, rootinv, pVec[i]);
		}
	}

	for (long i = 0; i < nprimes; ++i) {
		coeffpinv_array[i] = new mulmod_precon_t[i + 1];
		pProd[i] = (i == 0) ? to_ZZ((long) pVec[i]) : pProd[i - 1] * (long) pVec[i];
		pProdh[i] = pProd[i] / 2;
		pHat[i] = new ZZ[i + 1];
		pHatInvModp[i] = new uint64_t[i + 1];
		for (long j = 0; j < i + 1; ++j) {
			pHat[i][j] = ZZ(1);
			for (long k = 0; k < j; ++k) {
				pHat[i][j] *= (long) pVec[k];
			}
			for (long k = j + 1; k < i + 1; ++k) {
				pHat[i][j] *= (long) pVec[k];
			}
			pHatInvModp[i][j] = to_long(pHat[i][j] % (long) pVec[j]);
			pHatInvModp[i][j] = invMod(pHatInvModp[i][j], pVec[j]);
			coeffpinv_array[i][j] = PrepMulModPrecon(pHatInvModp[i][j], pVec[j]);
		}
	}
}

bool RingMultiplier::primeTest(uint64_t p) {
	if(p < 2) return false;
	if(p != 2 && p % 2 == 0) return false;
	uint64_t s = p - 1;
	while(s % 2 == 0) {
		s /= 2;
	}
	for(long i = 0; i < 200; i++) {
		uint64_t temp1 = rand();
		temp1  = (temp1 << 32) | rand();
		temp1 = temp1 % (p - 1) + 1;
		uint64_t temp2 = s;
		uint64_t mod = powMod(temp1,temp2,p);
		while (temp2 != p - 1 && mod != 1 && mod != p - 1) {
			mulMod(mod, mod, mod, p);
		    temp2 *= 2;
		}
		if (mod != p - 1 && temp2 % 2 == 0) return false;
	}
	return true;
}

void RingMultiplier::NTT(uint64_t* a, long index) {
	long t = N;
	long logt1 = logN + 1;
	uint64_t p = pVec[index];
	uint64_t pInv = pInvVec[index];
	for (long m = 1; m < N; m <<= 1) {
		t >>= 1;
		logt1 -= 1;
		for (long i = 0; i < m; i++) {
			long j1 = i << logt1;
			long j2 = j1 + t - 1;
			uint64_t W = scaledRootPows[index][m + i];
			for (long j = j1; j <= j2; j++) {
				butt(a[j], a[j+t], W, p, pInv);
			}
		}
	}
}

void RingMultiplier::INTT(uint64_t* a, long index) {
	uint64_t p = pVec[index];
	uint64_t pInv = pInvVec[index];
	long t = 1;
	for (long m = N; m > 1; m >>= 1) {
		long j1 = 0;
		long h = m >> 1;
		for (long i = 0; i < h; i++) {
			long j2 = j1 + t - 1;
			uint64_t W = scaledRootInvPows[index][h + i];
			for (long j = j1; j <= j2; j++) {
				ibutt(a[j], a[j+t], W, p, pInv);
			}
			j1 += (t << 1);
		}
		t <<= 1;
	}

	uint64_t NScale = scaledNInv[index];
	for (long i = 0; i < N; i++) {
		idivN(a[i], NScale, p, pInv);
	}
}

//----------------------------------------------------------------------------------
//   FFT
//----------------------------------------------------------------------------------

void RingMultiplier::CRT(uint64_t* rx, ZZ* x, const long np) {
	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rxi[n] = _ntl_general_rem_one_struct_apply(x[n].rep, pi, red_ss);
		}
		NTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;
}

void RingMultiplier::addNTTAndEqual(uint64_t* ra, uint64_t* rb, const long np) {
	for (long i = 0; i < np; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t pi = pVec[i];
		for (long n = 0; n < N; ++n) {
			rai[n] += rbi[n];
			if(rai[n] > pi) rai[n] -= pi;
		}
	}
}

void RingMultiplier::reconstruct(ZZ* x, uint64_t* rx, long np, const ZZ& q) {
	ZZ* pHatnp = pHat[np - 1];
	uint64_t* pHatInvModpnp = pHatInvModp[np - 1];
	mulmod_precon_t* coeffpinv_arraynp = coeffpinv_array[np - 1];
	ZZ& pProdnp = pProd[np - 1];
	ZZ& pProdhnp = pProdh[np - 1];
	NTL_EXEC_RANGE(N, first, last);
	for (long n = first; n < last; ++n) {
		ZZ& acc = x[n];
		QuickAccumBegin(acc, pProdnp.size());
		for (long i = 0; i < np; i++) {
			long p = pVec[i];
			long tt = pHatInvModpnp[i];
			mulmod_precon_t ttpinv = coeffpinv_arraynp[i];
			long s = MulModPrecon(rx[n + (i << logN)], tt, p, ttpinv);
			QuickAccumMulAdd(acc, pHatnp[i], s);
		}
		QuickAccumEnd(acc);
		rem(x[n], x[n], pProdnp);
		if (x[n] > pProdhnp) x[n] -= pProdnp;
		x[n] %= q;
	}
	NTL_EXEC_RANGE_END;
}

void RingMultiplier::mult(ZZ* x, ZZ* a, ZZ* b, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();
	uint64_t* rb = new uint64_t[np << logN]();
	uint64_t* rx = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
			rbi[n] = _ntl_general_rem_one_struct_apply(b[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		NTT(rbi, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rxi[n], rai[n], rbi[n], pi, pri);
		}
		INTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;

	reconstruct(x, rx, np, mod);

	delete[] ra;
	delete[] rb;
	delete[] rx;
}

void RingMultiplier::multNTT(ZZ* x, ZZ* a, uint64_t* rb, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();
	uint64_t* rx = new uint64_t[np << logN]();
	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rxi[n], rai[n], rbi[n], pi, pri);
		}
		INTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;

	reconstruct(x, rx, np, mod);

	delete[] ra;
	delete[] rx;
}

void RingMultiplier::multDNTT(ZZ* x, uint64_t* ra, uint64_t* rb, long np, const ZZ& mod) {
	uint64_t* rx = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rxi[n], rai[n], rbi[n], pi, pri);
		}
		INTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;

	reconstruct(x, rx, np, mod);

	delete[] rx;
}

void RingMultiplier::multAndEqual(ZZ* a, ZZ* b, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();
	uint64_t* rb = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
			rbi[n] = _ntl_general_rem_one_struct_apply(b[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		NTT(rbi, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rai[n], rai[n], rbi[n], pi, pri);
		}
		INTT(rai, i);
	}
	NTL_EXEC_RANGE_END;

	ZZ* pHatnp = pHat[np - 1];
	uint64_t* pHatInvModpnp = pHatInvModp[np - 1];

	reconstruct(a, ra, np, mod);

	delete[] ra;
	delete[] rb;
}

void RingMultiplier::multNTTAndEqual(ZZ* a, uint64_t* rb, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rbi = rb + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rai[n], rai[n], rbi[n], pi, pri);
		}
		INTT(rai, i);
	}
	NTL_EXEC_RANGE_END;

	ZZ* pHatnp = pHat[np - 1];
	uint64_t* pHatInvModpnp = pHatInvModp[np - 1];

	reconstruct(a, ra, np, mod);

	delete[] ra;
}

void RingMultiplier::square(ZZ* x, ZZ* a, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();
	uint64_t* rx = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rxi[n], rai[n], rai[n], pi, pri);
		}
		INTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;

	ZZ* pHatnp = pHat[np - 1];
	uint64_t* pHatInvModpnp = pHatInvModp[np - 1];

	reconstruct(x, rx, np, mod);

	delete[] ra;
	delete[] rx;
}

void RingMultiplier::squareNTT(ZZ* x, uint64_t* ra, long np, const ZZ& mod) {
	uint64_t* rx = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t* rxi = rx + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rxi[n], rai[n], rai[n], pi, pri);
		}
		INTT(rxi, i);
	}
	NTL_EXEC_RANGE_END;

	reconstruct(x, rx, np, mod);

	delete[] rx;
}

void RingMultiplier::squareAndEqual(ZZ* a, long np, const ZZ& mod) {
	uint64_t* ra = new uint64_t[np << logN]();

	NTL_EXEC_RANGE(np, first, last);
	for (long i = first; i < last; ++i) {
		uint64_t* rai = ra + (i << logN);
		uint64_t pi = pVec[i];
		uint64_t pri = prVec[i];
		_ntl_general_rem_one_struct* red_ss = red_ss_array[i];
		for (long n = 0; n < N; ++n) {
			rai[n] = _ntl_general_rem_one_struct_apply(a[n].rep, pi, red_ss);
		}
		NTT(rai, i);
		for (long n = 0; n < N; ++n) {
			mulModBarrett(rai[n], rai[n], rai[n], pi, pri);
		}
		INTT(rai, i);
	}
	NTL_EXEC_RANGE_END;

	reconstruct(a, ra, np, mod);

	delete[] ra;
}

void RingMultiplier::mulMod(uint64_t &r, uint64_t a, uint64_t b, uint64_t m) {
	unsigned __int128 mul = static_cast<unsigned __int128>(a) * b;
	mul %= static_cast<unsigned __int128>(m);
	r = static_cast<uint64_t>(mul);
}

void RingMultiplier::mulModBarrett(uint64_t& r, uint64_t a, uint64_t b, uint64_t p, uint64_t pr) {
	unsigned __int128 mul = static_cast<unsigned __int128>(a) * b;
	uint64_t abot = static_cast<uint64_t>(mul);
	uint64_t atop = static_cast<uint64_t>(mul >> 64);
	unsigned __int128 tmp = static_cast<unsigned __int128>(abot) * pr;
	tmp >>= 64;
	tmp += static_cast<unsigned __int128>(atop) * pr;
	tmp >>= kbar2 - 64;
	tmp *= p;
	tmp = mul - tmp;
	r = static_cast<uint64_t>(tmp);
	if(r >= p) r -= p;
}

void RingMultiplier::butt(uint64_t& a, uint64_t& b, uint64_t W, uint64_t p, uint64_t pInv) {
	unsigned __int128 U = static_cast<unsigned __int128>(b) * W;
	uint64_t U0 = static_cast<uint64_t>(U);
	uint64_t U1 = U >> 64;
	uint64_t Q = U0 * pInv;
	unsigned __int128 Hx = static_cast<unsigned __int128>(Q) * p;
	uint64_t H = Hx >> 64;
	uint64_t V = U1 < H ? U1 + p - H : U1 - H;
	b = a < V ? a + p - V : a - V;
	a += V;
	if (a > p) a -= p;
}

void RingMultiplier::ibutt(uint64_t& a, uint64_t& b, uint64_t W, uint64_t p, uint64_t pInv) {
	uint64_t T = a < b ? a + p - b : a - b;
	a += b;
	if (a > p) a -= p;
	unsigned __int128 UU = static_cast<unsigned __int128>(T) * W;
	uint64_t U0 = static_cast<uint64_t>(UU);
	uint64_t U1 = UU >> 64;
	uint64_t Q = U0 * pInv;
	unsigned __int128 Hx = static_cast<unsigned __int128>(Q) * p;
	uint64_t H = Hx >> 64;
	b = (U1 < H) ? U1 + p - H : U1 - H;
}

void RingMultiplier::idivN(uint64_t& a, uint64_t NScale, uint64_t p, uint64_t pInv) {
	unsigned __int128 U = static_cast<unsigned __int128>(a) * NScale;
	uint64_t U0 = static_cast<uint64_t>(U);
	uint64_t U1 = U >> 64;
	uint64_t Q = U0 * pInv;
	unsigned __int128 Hx = static_cast<unsigned __int128>(Q) * p;
	uint64_t H = Hx >> 64;
	a = (U1 < H) ? U1 + p - H : U1 - H;
}

uint64_t RingMultiplier::invMod(uint64_t x, uint64_t m) {
	return powMod(x, m - 2, m);
}

uint64_t RingMultiplier::powMod(uint64_t x, uint64_t y, uint64_t modulus) {
	uint64_t res = 1;
	while (y > 0) {
		if (y & 1) {
			mulMod(res, res, x, modulus);
		}
		y = y >> 1;
		mulMod(x, x, x, modulus);
	}
	return res;
}

uint64_t RingMultiplier::inv(uint64_t x) {
	return pow(x, static_cast<uint64_t>(-1));
}

uint64_t RingMultiplier::pow(uint64_t x, uint64_t y) {
	uint64_t res = 1;
	while (y > 0) {
		if (y & 1) {
			res *= x;
		}
		y = y >> 1;
		x *= x;
	}
	return res;
}

uint32_t RingMultiplier::bitReverse(uint32_t x) {
	x = (((x & 0xaaaaaaaa) >> 1) | ((x & 0x55555555) << 1));
	x = (((x & 0xcccccccc) >> 2) | ((x & 0x33333333) << 2));
	x = (((x & 0xf0f0f0f0) >> 4) | ((x & 0x0f0f0f0f) << 4));
	x = (((x & 0xff00ff00) >> 8) | ((x & 0x00ff00ff) << 8));
	return ((x >> 16) | (x << 16));
}

void RingMultiplier::findPrimeFactors(vector<uint64_t> &s, uint64_t number) {
	while (number % 2 == 0) {
		s.push_back(2);
		number /= 2;
	}
	for (uint64_t i = 3; i < sqrt(number); i++) {
		while (number % i == 0) {
			s.push_back(i);
			number /= i;
		}
	}
	if (number > 2) {
		s.push_back(number);
	}
}

uint64_t RingMultiplier::findPrimitiveRoot(uint64_t modulus) {
	vector<uint64_t> s;
	uint64_t phi = modulus - 1;
	findPrimeFactors(s, phi);
	for (uint64_t r = 2; r <= phi; r++) {
		bool flag = false;
		for (auto it = s.begin(); it != s.end(); it++) {
			if (powMod(r, phi / (*it), modulus) == 1) {
				flag = true;
				break;
			}
		}
		if (flag == false) {
			return r;
		}
	}
	return -1;
}

uint64_t RingMultiplier::findMthRootOfUnity(uint64_t M, uint64_t mod) {
    uint64_t res;
    res = findPrimitiveRoot(mod);
    if((mod - 1) % M == 0) {
        uint64_t factor = (mod - 1) / M;
        res = powMod(res, factor, mod);
        return res;
    }
    else {
        return -1;
    }
}


#include <boost_adaptbx/graph/graph_type.hpp>
#include <boost_adaptbx/graph/graph_export_adaptor.hpp>
#include <boost_adaptbx/graph/vertex_map.hpp>

#include <boost_adaptbx/exporting.hpp>

#include <boost/python/module.hpp>
#include <boost/python/list.hpp>
#include <boost/python/tuple.hpp>
#include <boost/python/dict.hpp>
#include <boost/python/def.hpp>
#include <boost/python/extract.hpp>
#include <boost/python/stl_iterator.hpp>

#include <boost/graph/graph_traits.hpp>
#include <boost/graph/stoer_wagner_min_cut.hpp>
#include <boost/graph/one_bit_color_map.hpp>
#include <boost/graph/boykov_kolmogorov_max_flow.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/properties.hpp>

#include <boost/mpl/if.hpp>
#include <boost/type_traits.hpp>

#include <iostream>

namespace boost_adaptbx
{

template< typename Graph >
struct minimum_cut_export
{
  typedef boost::graph_traits< Graph > graph_traits;
  typedef typename graph_traits::vertex_iterator vertex_iterator;
  typedef typename graph_traits::edge_iterator edge_iterator;
  typedef vertex_map::index_map< Graph > index_map_type;
  typedef typename index_map_type::property_map_type index_property_map_type;
  typedef boost::one_bit_color_map< index_property_map_type > parity_map_type;
  typedef edge_map::generic_edge_map< Graph, double > edge_map_type;
  typedef typename edge_map_type::property_map_type edge_property_map_type;

  static boost::python::tuple stoer_wagner_minimum_cut(Graph const& graph)
  {
    index_map_type index_map( graph );
    parity_map_type parities = boost::make_one_bit_color_map(
      boost::num_vertices( graph ),
      index_map.get()
      );

    edge_map_type edge_map( graph );
    edge_property_map_type epropmap( edge_map.get() );
    edge_iterator ei, ej;
    double w;

    for( boost::tie( ei, ej ) = boost::edges( graph ); ei != ej; ++ei )
    {
      w = boost::python::extract< double >( boost::get( boost::edge_weight, graph, *ei ) );
      boost::put( epropmap, *ei, w);
    }

    w = boost::stoer_wagner_min_cut(
      graph,
      epropmap,
      boost::parity_map( parities ).
      vertex_index_map( index_map.get() )
      );

    vertex_iterator di, dj;
    boost::python::list result;

    for( boost::tie( di, dj ) = boost::vertices( graph ); di != dj; ++di )
    {
      result.append( bool( boost::get( parities, *di ) ) );
    }

    return boost::python::make_tuple( w, result );
  }

  static void process()
  {
    using namespace boost::python;

    def( "stoer_wagner_min_cut", stoer_wagner_minimum_cut, arg( "graph" ) );
  }
};

struct min_cut_exporter
{
  template< typename Export >
  void operator ()(boost::mpl::identity< Export > myexport) const
  {
    typedef typename Export::first graph_type;
    typedef typename Export::second name_type;

    typedef boost::graph_traits< graph_type > graph_traits;
    typedef typename boost::mpl::if_<
      boost::is_same<
        typename  graph_traits::directed_category,
        boost::undirected_tag
        >,
      minimum_cut_export< graph_type >,
      graph_export_adaptor::no_export< graph_type >
      >::type exporter_type;
    exporter_type::process();
  }
};

template< typename Graph >
struct maximum_flow_export
{
  typedef boost::graph_traits< Graph > graph_traits;
  typedef typename graph_traits::vertex_iterator vertex_iterator;
  typedef typename graph_traits::vertex_descriptor vertex_descriptor;
  typedef typename graph_traits::edge_iterator edge_iterator;
  typedef typename graph_traits::edge_descriptor edge_descriptor;

  typedef vertex_map::index_map< Graph > index_map_type;
  typedef typename index_map_type::property_map_type index_property_map_type;

  typedef edge_map::generic_edge_map< Graph, double > capacity_map_type;
  typedef typename capacity_map_type::property_map_type capacity_property_map_type;

  typedef edge_map::generic_edge_map< Graph, edge_descriptor > reverse_edge_map_type;
  typedef typename reverse_edge_map_type::property_map_type reverse_edge_property_map_type;

  typedef vertex_map::generic_vertex_map< Graph, boost::default_color_type > color_map_type;
  typedef typename color_map_type::property_map_type color_property_map_type;

  typedef graph_export_adaptor::vertex_descriptor_converter< vertex_descriptor > converter;

  static boost::python::tuple boykov_kolmogorov_max_flow(
    Graph const& graph,
    boost::python::dict reverse_edge_map,
    typename converter::type source,
    typename converter::type sink
    )
  {
    using namespace boost;
    using namespace boost::python;

    index_map_type index_map( graph );

    capacity_map_type capmap( graph );
    capacity_property_map_type cappropmap( capmap.get() );
    edge_iterator ei, ej;

    for( boost::tuples::tie( ei, ej ) = edges( graph ); ei != ej; ++ei )
    {
      double weight( extract< double >( get( edge_weight, graph, *ei ) ) );
      boost::put( cappropmap, *ei, weight );
    }

    reverse_edge_map_type revmap( graph );
    reverse_edge_property_map_type revpropmap( revmap.get() );

    stl_input_iterator< object > end;

    for( stl_input_iterator< object > it( reverse_edge_map.iteritems() ); it!=end; ++it)
    {
      edge_descriptor ed = extract< edge_descriptor >( (*it)[0] );
      edge_descriptor red = extract< edge_descriptor >( (*it)[1] );
      put( revpropmap, ed, red );
    }

    color_map_type colormap( graph );
    color_property_map_type colorpropmap( colormap.get() );
    capacity_map_type resicapmap( graph );

    double w = boost::boykov_kolmogorov_max_flow(
        graph,
        cappropmap,
        resicapmap.get(),
        revpropmap,
        colorpropmap,
        index_map.get(),
        converter::backward( source ),
        converter::backward( sink )
        );

    vertex_iterator di, dj;
    boost::python::list result;
    default_color_type const black( color_traits< default_color_type >::black() );
    default_color_type current;

    for( boost::tie( di, dj ) = boost::vertices( graph ); di != dj; ++di )
    {
      current = get( colorpropmap, *di );
      result.append( current == black );
    }

    return boost::python::make_tuple( w, result );
  }

  static void process()
  {
    using namespace boost::python;

    def(
      "boykov_kolmogorov_max_flow",
      boykov_kolmogorov_max_flow,
      ( arg( "graph" ), arg( "reverse_edge_map" ), arg( "source" ), arg( "sink" ) )
      );
  }
};

struct max_flow_exporter
{
  template< typename Export >
  void operator ()(boost::mpl::identity< Export > myexport) const
  {
    typedef typename Export::first graph_type;
    typedef typename Export::second name_type;

    typedef boost::graph_traits< graph_type > graph_traits;
    typedef typename boost::mpl::if_<
      boost::is_same<
        typename  graph_traits::directed_category,
        boost::directed_tag
        >,
      maximum_flow_export< graph_type >,
      graph_export_adaptor::no_export< graph_type >
      >::type exporter_type;
    exporter_type::process();
  }
};

} // namespace boost_adaptbx

BOOST_PYTHON_MODULE(boost_adaptbx_graph_min_cut_max_flow_ext)
{
  boost_adaptbx::exporting::class_list< boost_adaptbx::graph_type::exports >::process(
    boost_adaptbx::min_cut_exporter()
    );
  boost_adaptbx::exporting::class_list< boost_adaptbx::graph_type::exports >::process(
    boost_adaptbx::max_flow_exporter()
    );
}