/* =========================================================================
   Copyright (c) 2010-2016, Institute for Microelectronics,
                            Institute for Analysis and Scientific Computing,
                            TU Wien.
   Portions of this software are copyright by UChicago Argonne, LLC.

                            -----------------
                  ViennaCL - The Vienna Computing Library
                            -----------------

   Project Head:    Karl Rupp                   rupp@iue.tuwien.ac.at

   (A list of authors and contributors can be found in the PDF manual)

   License:         MIT (X11), see file LICENSE in the base directory
============================================================================= */

/*
*   Benchmark:  Sparse matrix operations, i.e. matrix-vector products (sparse.cpp and sparse.cu are identical, the latter being required for compilation using CUDA nvcc)
*
*/

//#define VIENNACL_BUILD_INFO
#ifndef NDEBUG
 #define NDEBUG
#endif

#define VIENNACL_WITH_UBLAS 1

#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/matrix_sparse.hpp>
#include <boost/numeric/ublas/operation_sparse.hpp>
#include <boost/numeric/ublas/lu.hpp>

#include "viennacl/scalar.hpp"
#include "viennacl/vector.hpp"
#include "viennacl/coordinate_matrix.hpp"
#include "viennacl/compressed_matrix.hpp"
#include "viennacl/ell_matrix.hpp"
#include "viennacl/hyb_matrix.hpp"
#include "viennacl/sliced_ell_matrix.hpp"
#include "viennacl/linalg/prod.hpp"
#include "viennacl/linalg/norm_2.hpp"
#include "viennacl/io/matrix_market.hpp"
#include "viennacl/linalg/ilu.hpp"
#include "viennacl/tools/timer.hpp"

#include <iostream>
#include <vector>

#define BENCHMARK_RUNS          10

inline void printOps(double num_ops, double exec_time)
{
  std::cout << "GFLOPs: " << num_ops / (1000000 * exec_time * 1000) << std::endl;
}

template<typename ScalarType>
int run_benchmark()
{
  viennacl::tools::timer timer;
  double exec_time;

  ScalarType std_factor1 = ScalarType(3.1415);
  ScalarType std_factor2 = ScalarType(42.0);
  viennacl::scalar<ScalarType> vcl_factor1(std_factor1);
  viennacl::scalar<ScalarType> vcl_factor2(std_factor2);

  boost::numeric::ublas::vector<ScalarType> ublas_vec1;
  boost::numeric::ublas::vector<ScalarType> ublas_vec2;

  boost::numeric::ublas::compressed_matrix<ScalarType> ublas_matrix;
  if (!viennacl::io::read_matrix_market_file(ublas_matrix, "../examples/testdata/mat65k.mtx"))
  {
    std::cout << "Error reading Matrix file" << std::endl;
    return 0;
  }
  //unsigned int cg_mat_size = cg_mat.size();
  std::cout << "done reading matrix" << std::endl;

  ublas_vec1 = boost::numeric::ublas::scalar_vector<ScalarType>(ublas_matrix.size1(), ScalarType(1.0));
  ublas_vec2 = ublas_vec1;

  viennacl::compressed_matrix<ScalarType, 1> vcl_compressed_matrix_1;
  viennacl::compressed_matrix<ScalarType, 4> vcl_compressed_matrix_4;
  viennacl::compressed_matrix<ScalarType, 8> vcl_compressed_matrix_8;

  viennacl::coordinate_matrix<ScalarType> vcl_coordinate_matrix_128;

  viennacl::ell_matrix<ScalarType, 1> vcl_ell_matrix_1;
  viennacl::hyb_matrix<ScalarType, 1> vcl_hyb_matrix_1;
  viennacl::sliced_ell_matrix<ScalarType> vcl_sliced_ell_matrix_1;

  viennacl::vector<ScalarType> vcl_vec1(ublas_vec1.size());
  viennacl::vector<ScalarType> vcl_vec2(ublas_vec1.size());

  //cpu to gpu:
  viennacl::copy(ublas_matrix, vcl_compressed_matrix_1);
  #ifndef VIENNACL_EXPERIMENTAL_DOUBLE_PRECISION_WITH_STREAM_SDK_ON_GPU
  viennacl::copy(ublas_matrix, vcl_compressed_matrix_4);
  viennacl::copy(ublas_matrix, vcl_compressed_matrix_8);
  #endif
  viennacl::copy(ublas_matrix, vcl_coordinate_matrix_128);
  viennacl::copy(ublas_matrix, vcl_ell_matrix_1);
  viennacl::copy(ublas_matrix, vcl_hyb_matrix_1);
  viennacl::copy(ublas_matrix, vcl_sliced_ell_matrix_1);
  viennacl::copy(ublas_vec1, vcl_vec1);
  viennacl::copy(ublas_vec2, vcl_vec2);

  ///////////// Matrix operations /////////////////

  std::cout << "------- Matrix-Vector product on CPU ----------" << std::endl;
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    //ublas_vec1 = boost::numeric::ublas::prod(ublas_matrix, ublas_vec2);
    boost::numeric::ublas::axpy_prod(ublas_matrix, ublas_vec2, ublas_vec1, true);
  }
  exec_time = timer.get();
  std::cout << "CPU time: " << exec_time << std::endl;
  std::cout << "CPU "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << ublas_vec1[0] << std::endl;

  std::cout << "------- Matrix-Vector product with compressed_matrix ----------" << std::endl;

  vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_1, vcl_vec2); //startup calculation
  vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_4, vcl_vec2); //startup calculation
  vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_8, vcl_vec2); //startup calculation
  //std_result = 0.0;

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_1, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time align1: " << exec_time << std::endl;
  std::cout << "GPU align1 "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  std::cout << "Testing triangular solves: compressed_matrix" << std::endl;

  viennacl::copy(ublas_vec1, vcl_vec1);
  viennacl::linalg::inplace_solve(trans(vcl_compressed_matrix_1), vcl_vec1, viennacl::linalg::unit_lower_tag());
  viennacl::copy(ublas_vec1, vcl_vec1);
  std::cout << "ublas..." << std::endl;
  timer.start();
  boost::numeric::ublas::inplace_solve(trans(ublas_matrix), ublas_vec1, boost::numeric::ublas::unit_lower_tag());
  std::cout << "Time elapsed: " << timer.get() << std::endl;
  std::cout << "ViennaCL..." << std::endl;
  viennacl::backend::finish();
  timer.start();
  viennacl::linalg::inplace_solve(trans(vcl_compressed_matrix_1), vcl_vec1, viennacl::linalg::unit_lower_tag());
  viennacl::backend::finish();
  std::cout << "Time elapsed: " << timer.get() << std::endl;

  ublas_vec1 = boost::numeric::ublas::prod(ublas_matrix, ublas_vec2);

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_4, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time align4: " << exec_time << std::endl;
  std::cout << "GPU align4 "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_compressed_matrix_8, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time align8: " << exec_time << std::endl;
  std::cout << "GPU align8 "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  std::cout << "------- Matrix-Vector product with coordinate_matrix ----------" << std::endl;
  vcl_vec1 = viennacl::linalg::prod(vcl_coordinate_matrix_128, vcl_vec2); //startup calculation
  viennacl::backend::finish();

  viennacl::copy(vcl_vec1, ublas_vec2);
  long err_cnt = 0;
  for (std::size_t i=0; i<ublas_vec1.size(); ++i)
  {
    if ( fabs(ublas_vec1[i] - ublas_vec2[i]) / std::max(fabs(ublas_vec1[i]), fabs(ublas_vec2[i])) > 1e-2)
    {
      std::cout << "Error at index " << i << ": Should: " << ublas_vec1[i] << ", Is: " << ublas_vec2[i] << std::endl;
      ++err_cnt;
      if (err_cnt > 5)
        break;
    }
  }

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_coordinate_matrix_128, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time: " << exec_time << std::endl;
  std::cout << "GPU "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  std::cout << "------- Matrix-Vector product with ell_matrix ----------" << std::endl;
  vcl_vec1 = viennacl::linalg::prod(vcl_ell_matrix_1, vcl_vec2); //startup calculation
  viennacl::backend::finish();

  viennacl::copy(vcl_vec1, ublas_vec2);
  err_cnt = 0;
  for (std::size_t i=0; i<ublas_vec1.size(); ++i)
  {
    if ( fabs(ublas_vec1[i] - ublas_vec2[i]) / std::max(fabs(ublas_vec1[i]), fabs(ublas_vec2[i])) > 1e-2)
    {
      std::cout << "Error at index " << i << ": Should: " << ublas_vec1[i] << ", Is: " << ublas_vec2[i] << std::endl;
      ++err_cnt;
      if (err_cnt > 5)
        break;
    }
  }

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_ell_matrix_1, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time: " << exec_time << std::endl;
  std::cout << "GPU "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  std::cout << "------- Matrix-Vector product with hyb_matrix ----------" << std::endl;
  vcl_vec1 = viennacl::linalg::prod(vcl_hyb_matrix_1, vcl_vec2); //startup calculation
  viennacl::backend::finish();

  viennacl::copy(vcl_vec1, ublas_vec2);
  err_cnt = 0;
  for (std::size_t i=0; i<ublas_vec1.size(); ++i)
  {
    if ( fabs(ublas_vec1[i] - ublas_vec2[i]) / std::max(fabs(ublas_vec1[i]), fabs(ublas_vec2[i])) > 1e-2)
    {
      std::cout << "Error at index " << i << ": Should: " << ublas_vec1[i] << ", Is: " << ublas_vec2[i] << std::endl;
      ++err_cnt;
      if (err_cnt > 5)
        break;
    }
  }

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_hyb_matrix_1, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time: " << exec_time << std::endl;
  std::cout << "GPU "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  std::cout << "------- Matrix-Vector product with sliced_ell_matrix ----------" << std::endl;
  vcl_vec1 = viennacl::linalg::prod(vcl_sliced_ell_matrix_1, vcl_vec2); //startup calculation
  viennacl::backend::finish();

  viennacl::copy(vcl_vec1, ublas_vec2);
  err_cnt = 0;
  for (std::size_t i=0; i<ublas_vec1.size(); ++i)
  {
    if ( fabs(ublas_vec1[i] - ublas_vec2[i]) / std::max(fabs(ublas_vec1[i]), fabs(ublas_vec2[i])) > 1e-2)
    {
      std::cout << "Error at index " << i << ": Should: " << ublas_vec1[i] << ", Is: " << ublas_vec2[i] << std::endl;
      ++err_cnt;
      if (err_cnt > 5)
        break;
    }
  }

  viennacl::backend::finish();
  timer.start();
  for (int runs=0; runs<BENCHMARK_RUNS; ++runs)
  {
    vcl_vec1 = viennacl::linalg::prod(vcl_sliced_ell_matrix_1, vcl_vec2);
  }
  viennacl::backend::finish();
  exec_time = timer.get();
  std::cout << "GPU time: " << exec_time << std::endl;
  std::cout << "GPU "; printOps(2.0 * static_cast<double>(ublas_matrix.nnz()), static_cast<double>(exec_time) / static_cast<double>(BENCHMARK_RUNS));
  std::cout << vcl_vec1[0] << std::endl;

  return EXIT_SUCCESS;
}

int main()
{
  std::cout << std::endl;
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << "               Device Info" << std::endl;
  std::cout << "----------------------------------------------" << std::endl;

#ifdef VIENNACL_WITH_OPENCL
  std::cout << viennacl::ocl::current_device().info() << std::endl;
#endif
  std::cout << std::endl;
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << "## Benchmark :: Sparse" << std::endl;
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << std::endl;
  std::cout << "   -------------------------------" << std::endl;
  std::cout << "   # benchmarking single-precision" << std::endl;
  std::cout << "   -------------------------------" << std::endl;
  run_benchmark<float>();
#ifdef VIENNACL_WITH_OPENCL
  if ( viennacl::ocl::current_device().double_support() )
#endif
  {
    std::cout << std::endl;
    std::cout << "   -------------------------------" << std::endl;
    std::cout << "   # benchmarking double-precision" << std::endl;
    std::cout << "   -------------------------------" << std::endl;
    run_benchmark<double>();
  }
  return 0;
}


/*
 * Copyright 2013-2015 Raphael Bost
 *
 * This file is part of ciphermed.

 *  ciphermed 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.
 *
 *  ciphermed 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 ciphermed.  If not, see <http://www.gnu.org/licenses/>. 2
 *
 */

#include <assert.h>
#include <vector>
#include <crypto/paillier.hh>
#include <crypto/gm.hh>
#include <NTL/ZZ.h>
#include <gmpxx.h>
#include <math/util_gmp_rand.h>

#include <ctime>

#include<iostream>

using namespace std;
using namespace NTL;

static void test_paillier()
{
    cout << "Test Paillier ...\n" << flush;

    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);
    gmp_randseed_ui(randstate,time(NULL));

    auto sk = Paillier_priv::keygen(randstate,16,2); //600 ,256
    Paillier_priv pp(sk,randstate);

    auto pk = pp.pubkey();
    mpz_class n = pk[0];
    Paillier p(pk,randstate);

    //mpz_class pt0, pt1,m;
    //mpz_urandomm(pt0.get_mpz_t(),randstate,n.get_mpz_t());
    //mpz_urandomm(pt1.get_mpz_t(),randstate,n.get_mpz_t());
    //mpz_urandomm(m.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_class pt0 = 2, pt1 = 3, m = 5; //instead of the random values
    cout << "pt0: "<< pt0 << endl;
    cout << "pt1: "<< pt1 << endl;
    cout << "m: "<< m << endl;

    mpz_class ct0 = p.encrypt(pt0);         cout << "ct0: "<< ct0 << endl;
    mpz_class ct1 = p.encrypt(pt1);         cout << "ct1: "<< ct1 << endl;
    mpz_class sum = p.add(ct0, ct1);        cout << "sum_enc: "<< sum << endl;
    mpz_class prod = p.constMult(m,ct0);    cout << "prod_enc: "<< prod << endl;
    //    mpz_class diff = p.constMult(-1, ct0);
    mpz_class diff = p.sub(ct0, ct1);       cout << "diff_enc: "<< diff << endl;

    assert(pp.decrypt(ct0) == pt0);
    assert(pp.decrypt(ct1) == pt1);
    assert(pp.decrypt(sum) == (pt0+pt1)%n);
    mpz_class d = pt0 - pt1;
    if (d < 0) {
        d += n;
    }
    assert( pp.decrypt(diff) == d);
    assert(pp.decrypt(prod) == (m*pt0)%n);

    cout << "Test Paillier passed" << endl;

    mpz_class A[3][3] = {{1,2,3},{1,1,1},{4,5,6}};
    //int rows = sizeof A / sizeof A[0];
    //int cols = sizeof A[0] / sizeof(mpz_class);
    //cout << rows << " " << cols <<endl;
    mpz_class A_enc[3][3];
    p.encryptMatrix(A, A_enc);
    mpz_class A_dec[3][3];
    pp.decryptMatrix(A_enc, A_dec);
    cout << "Test Matrix Paillier passed" << endl;
}

static void test_paillier_fast()
{
    cout << "Test Paillier Fast..." << flush;

    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);
    gmp_randseed_ui(randstate,time(NULL));

    auto sk = Paillier_priv_fast::keygen(randstate,600);
    Paillier_priv_fast pp(sk,randstate);

    auto pk = pp.pubkey();
    mpz_class n = pk[0];
    Paillier p(pk,randstate);

    mpz_class pt0, pt1,m;
    mpz_urandomm(pt0.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(pt1.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(m.get_mpz_t(),randstate,n.get_mpz_t());

    mpz_class ct0 = pp.encrypt(pt0);
    mpz_class ct1 = pp.encrypt(pt1);
    mpz_class sum = p.add(ct0, ct1);
    mpz_class prod = p.constMult(m,ct0);
    //    mpz_class diff = p.constMult(-1, ct0);
    mpz_class diff = p.sub(ct0, ct1);

    assert(pp.decrypt(ct0) == pt0);
    assert(pp.decrypt(ct1) == pt1);
    assert(pp.decrypt(sum) == (pt0+pt1)%n);
    mpz_class d = pt0 - pt1;
    if (d < 0) {
        d += n;
    }
    assert( pp.decrypt(diff) == d);
    assert(pp.decrypt(prod) == (m*pt0)%n);

    cout << " passed" << endl;
}

static void paillier_perf(unsigned int k, unsigned int a_bits, size_t n_iteration)
{
    cout << "Test Paillier performances ..." << endl;

    cout << "k = " << k << "\na_bits = " << a_bits << "\n" << n_iteration << " iterations" << endl;

    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);
    gmp_randseed_ui(randstate,time(NULL));

    auto sk = Paillier_priv::keygen(randstate,k,a_bits);
    Paillier_priv pp(sk,randstate);

    auto pk = pp.pubkey();
    mpz_class n = pk[0];
    Paillier p(pk,randstate);

    mpz_class pt0, pt1,m;
    mpz_urandomm(pt0.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(pt1.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(m.get_mpz_t(),randstate,n.get_mpz_t());

    mpz_class ct0 = p.encrypt(pt0);
    mpz_class ct1 = p.encrypt(pt1);
    mpz_class sum = p.add(ct0, ct1);
    mpz_class prod = p.constMult(m,ct0);
    //    mpz_class diff = p.constMult(-1, ct0);
    mpz_class diff = p.sub(ct0, ct1);

    assert(pp.decrypt(ct0) == pt0);
    assert(pp.decrypt(ct1) == pt1);
    assert(pp.decrypt(sum) == (pt0+pt1)%n);
    mpz_class d = pt0 - pt1;
    if (d < 0) {
        d += n;
    }
    assert( pp.decrypt(diff) == d);
    assert(pp.decrypt(prod) == (m*pt0)%n);

    struct timespec t0,t1;

    vector<mpz_class> ct(n_iteration);
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        mpz_class pt;
        mpz_urandomm(pt.get_mpz_t(),randstate,n.get_mpz_t());
        ct[i] = p.encrypt(pt);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    uint64_t t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "public encryption: "<<  ((double)t/1000000)/n_iteration <<"ms per plaintext" << endl;

    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        mpz_class pt;
        mpz_urandomm(pt.get_mpz_t(),randstate,n.get_mpz_t());
        pp.encrypt(pt);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "private encryption: "<<  ((double)t/1000000)/n_iteration <<"ms per plaintext" << endl;

    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        mpz_class pt;
        mpz_urandomm(pt.get_mpz_t(),randstate,n.get_mpz_t());
        pp.fast_encrypt_precompute(pt);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "private encryption with precomputation: "<<  ((double)t/1000000)/n_iteration <<"ms per plaintext" << endl;

    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        pp.decrypt(ct[i]);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "decryption: "<<  ((double)t/1000000)/n_iteration <<"ms per cyphertext" << endl;

}

static void paillier_fast_perf(unsigned int k, size_t n_iteration)
{
    cout << "Test Paillier Fast performances ..." << endl;

    cout << "k = " << k << "\n" << n_iteration << " iterations" << endl;

    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);
    gmp_randseed_ui(randstate,time(NULL));

    auto sk = Paillier_priv_fast::keygen(randstate,k);
    Paillier_priv_fast pp(sk,randstate);

    auto pk = pp.pubkey();
    mpz_class n = pk[0];
//    Paillier p(pk,randstate);

    mpz_class pt0, pt1,m;
    mpz_urandomm(pt0.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(pt1.get_mpz_t(),randstate,n.get_mpz_t());
    mpz_urandomm(m.get_mpz_t(),randstate,n.get_mpz_t());

    struct timespec t0,t1;
    uint64_t t;

    vector<mpz_class> ct(n_iteration);

    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        mpz_class pt;
        mpz_urandomm(pt.get_mpz_t(),randstate,n.get_mpz_t());
        ct[i] = pp.encrypt(pt);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "private encryption with generator precomputations: "<<  ((double)t/1000000)/n_iteration <<"ms per plaintext" << endl;

    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t0);
    for (size_t i = 0; i < n_iteration; i++) {
        pp.decrypt(ct[i]);
    }
    clock_gettime(CLOCK_THREAD_CPUTIME_ID,&t1);
    t = (((uint64_t)t1.tv_sec) - ((uint64_t)t0.tv_sec) )* 1000000000 + (t1.tv_nsec - t0.tv_nsec);
    cerr << "decryption: "<<  ((double)t/1000000)/n_iteration <<"ms per cyphertext" << endl;

}

static void test_gm()
{
    cout << "Test GM ..." << flush;

    gmp_randstate_t randstate;
    gmp_randinit_default(randstate);
    gmp_randseed_ui(randstate,time(NULL));

    auto sk = GM_priv::keygen(randstate);
    GM_priv pp(sk,randstate);

    auto pk = pp.pubkey();
    GM p(pk,randstate);

    bool b0 = true; //(bool)RandomBits_long(1);
    bool b1 = false; //(bool)RandomBits_long(1);

    mpz_class ct0 = p.encrypt(b0);
    mpz_class ct1 = p.encrypt(b1);
    mpz_class XOR = p.XOR(ct0, ct1);
    mpz_class rerand = p.reRand(ct0);

    assert(pp.decrypt(pk[1]) == true);
    assert(pp.decrypt(ct0) == b0);
    assert(pp.decrypt(ct1) == b1);
    assert(pp.decrypt(XOR) == (b0 xor b1));
    assert(pp.decrypt(rerand) == b0);

    cout << " passed" << endl;
}

int main(int ac, char **av)
{
    SetSeed(to_ZZ(time(NULL)));
    //test_elgamal();
    test_paillier();
    //test_paillier_fast();
    //test_gm();

    return 0;
}


///////////////////////////////////////////////////////////////////////////////
// random::example::poisson_convergence.cpp     						     //
//                                                                           //
//  Copyright 2010 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef LIBS_RANDOM_EXAMPLE_POISSON_CONVERGENCE_CPP_ER_2010
#define LIBS_RANDOM_EXAMPLE_POISSON_CONVERGENCE_CPP_ER_2010
#include <vector>
#include <algorithm>
#include <iterator>

#include <boost/mpl/int.hpp>

#include <boost/typeof/typeof.hpp>
#include <boost/ref.hpp>
#include <boost/range.hpp>

#include <boost/assign/std/vector.hpp>
#include <boost/format.hpp>
#include <boost/foreach.hpp>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>

#include <boost/math/distributions/poisson.hpp>
#include <boost/random/poisson_distribution.hpp>

#include <boost/random/poisson_ext/devroye/detail/step4/standard.hpp>
#include <boost/random/poisson_ext/devroye/detail/step4/squeeze.hpp>
#include <boost/random/poisson_ext/devroye/sampler/meta_int_mean.hpp>

#include <boost/random/mersenne_twister.hpp>
#include <boost/random/variate_generator.hpp>

#include <boost/statistics/detail/non_parametric/kolmogorov_smirnov/check_convergence.hpp>

template<typename Int>
void example_poisson_convergence(
	std::ostream& os,
    double mean,
    std::size_t n_loops,
    std::size_t n_init,
    std::size_t n_factor
)
{

	namespace ks = boost::statistics::detail::kolmogorov_smirnov;
	namespace devroye = boost::random::poisson::devroye;

    typedef Int int_;
    typedef double 	val_;
    typedef boost::math::poisson_distribution<val_> dist_;
    typedef boost::poisson_distribution<int_> random1_;
    typedef devroye::tag::standard tag3_;
    typedef devroye::tag::squeeze tag4_;
    typedef typename devroye::sampler::template
    	meta_int_mean<tag3_,int_>::type random3_;
    typedef typename devroye::sampler::template
    	meta_int_mean<tag4_,int_>::type random4_;

	typedef boost::mt19937 urng_;
	typedef std::vector<val_> vals_;
	typedef ks::check_convergence<val_> check_;

    urng_ urng;
	urng_ urng1 = urng;
    urng_ urng2 = urng;
    urng_ urng3 = urng;
    urng_ urng4 = urng;
	check_ check;

	os  << "mean = " << mean << std::endl;
	{
    	typedef boost::variate_generator<urng_&,random1_> vg1_;
    	typedef boost::variate_generator<urng_&,random3_> vg3_;
    	typedef boost::variate_generator<urng_&,random4_> vg4_;

        dist_ dist(mean);
        os << "default : " << std::endl;
        {
			vg1_ vg1(urng,random1_(mean));
    		check(n_loops,n_init,n_factor,dist,vg1,os);
        }
        os << "devroye - standard" << std::endl;
        {
        	random3_ r(mean);
			vg3_ vg(urng,r);
            vg();
    		check(n_loops,n_init,n_factor,dist,vg,os);
            //os << vg.distribution() << std::endl;
        }
        os << "devroye - squeeze" << std::endl;
        {
        	random4_ r(mean);
			vg4_ vg(urng,r);
    		check(n_loops,n_init,n_factor,dist,vg,os);
            //os << vg.distribution() << std::endl;
        }
    }
	os << "<-" << std::endl;

}

#endif


/**
 *  @file    FDProblem.hpp
 *  @brief   The finite difference problem.
 *  @author  Francois Roy
 *  @date    12/01/2019
 */
#ifndef FDPROBLEM_H
#define FDPROBLEM_H

// #include <iostream>
#include <stdexcept>
#include <string>
#include <vector>
#include <Eigen/SparseCore>
#include "spdlog/spdlog.h"
#include "Parameters.hpp"
#include "FDMesh.hpp"

namespace numerical {

namespace fdm {

/**
 * Defines the finite difference problem over the hypercube.
 * This class only define the 1D diffusion problem with Dirichlet/Neumann
 * boundary conditions and heterogenous diffusion coefficient (for now).
 */
template <typename T>
class FDProblem {
typedef Eigen::Matrix<T, Eigen::Dynamic, 1> Vec;
typedef std::vector<Eigen::Matrix<T, 3, 1>> Coord;
// function pointer to member functions
typedef  T (FDProblem<T>::*fctptr)(T x1, T x2, T time);
private:
    int m_dim; // space dimensions of the problem
    Vec m_t;
    std::vector<T> m_dx;
    T m_dt;
protected:
    Parameters<T>* m_params;  // default parameters
    FDMesh<T>* m_mesh;  // mesh
    int m_bc_types[6];  // boundary types (default = Dirichlet)
    Coord m_coordinates;
    // array of function pointers to member functions
    fctptr m_functions[6] = {&FDProblem<T>::left, &FDProblem<T>::right,
                             &FDProblem<T>::bottom, &FDProblem<T>::top,
                             &FDProblem<T>::back, &FDProblem<T>::front};
public:
    FDProblem(Parameters<T>* params):
      m_params(params),
      m_bc_types{0, 0, 0, 0, 0, 0} {
        // define other variables variables
        m_dim = m_params->lengths.size();
        if(m_dim>1){
            // TODO implement higher dimensions
            spdlog::error("Only one-dimensional problems supported.");
        }
  	    // define mesh
  	    m_mesh = new FDMesh<T>(m_params->lengths, m_params->t0, m_params->tend,
          m_params->n, m_params->nt);
        m_coordinates = m_mesh->coordinates();  // constant
        m_t = m_mesh->t();
        m_dx = m_mesh->dx();
        m_dt = m_mesh->dt();
    }
    virtual ~FDProblem(){
        delete m_mesh;
    }

    /**
    * Diffusion coefficient value. The default is a constant obtained from
    * m_params.
    *
    * @param x The \f$x\f$-, \f$y\f$-, and \f$z\f$-coordinates of the mesh
    *    node.
    * @param t The discrete time.
    * @param u The state variable for nonlinear problems.
    * @return The diffusion coefficient value at a specified mesh location.
    */
    virtual T alpha(const Eigen::Matrix<T, 3, 1>& x, T t, T u=0){
  	    return m_params->alpha;
    }
    /**
    *  @return The mesh coordinates.
    */
    std::vector<Eigen::Matrix<T, 3, 1>> coordinates(){
        return m_coordinates;
    }
    /**
    * Set the boundary types, with the following code:
    *
    *  0 = Dirichlet
    *  1 = Neumann
    *
    * For the six boundaries in the following order:
    *
    * 0- letf
    * 1- right
    * 2- bottom
    * 3- top
    * 4- back
    * 5- front
    *
    *  @param types The type of boundary conditions for all boundaries.
    */
    void bc_types(const int (&types)[6]){
        for (int i=0; i<6; i++){
            m_bc_types[i] = types[i];
            // spdlog::info("{}", m_bc_types[i]);
        }
    }
    /**
    *
    * @param boundary The boundary selection.
    * @return The boundary type of the selected boundary,
    */
    int bc_type(int boundary){
        return m_bc_types[boundary];
    }
    /**
    * User defined function: back boundary.
    *
    * Back boundary value, i.e. at z = length[2][0]. Only for 3D models.
    *
    * @param x The \f$x\f$-coordinate.
    * @param y The \f$y\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T back(T x, T y, T t){
        return 0.0;
    }
    /**
    * User defined function: bottom boundary.
    *
    * Bottom boundary value, i.e. at y = lengths[1][0]. Only for 2D and
    * 3D models.
    *
    * @param x The \f$x\f$-coordinate.
    * @param z The \f$z\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T bottom(T x, T z, T t){
        return 0.0;
    }
    /**
    * Set the value of the diagonals of the coefficient matrix at the indices
    * of the boundary nodes.
    *
    *  @param dia The main diagonal.
    *  @param lower The lower diagonal.
    *  @param upper The upper diagonal.
    *  @param lower_a The second lower diagonal.
    *  @param upper_a The second upper diagonal.
    *  @param lower_b The third lower diagonal.
    *  @param upper_b The third upper diagonal.
    *  @param t The discrete time.
    */
    void coeffs_bc(Vec& dia, Vec& lower, Vec& upper, Vec& lower_a,
        Vec& upper_a, Vec& lower_b, Vec& upper_b, T t=0.0){
    }
    /**
    * @return The spatial dimension of the problem.
    */
    int dim(){
        return m_dim;
    }
    /**
    * @return the constant time increment.
    */
    T dt() {
        return m_dt;
    }
    /**
    * @return the spatial increment.
    */
    std::vector<T> dx() {
        return m_dx;
    }
    /**
    * User defined function: front boundary.
    *
    * Front boundary value, i.e. at z = length[2][0]. Only for 3D models.
    *
    * @param x The \f$x\f$-coordinate.
    * @param y The \f$y\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T front(T x, T y, T t){
        return 0.0;
    }
    /**
    * User defined function: initial  value.
    *
    * @param x The \f$x\f$-, \f$y\f$-, and \f$z\f$-coordinates of the mesh
    *    node.
    * @return The initial value at a specified mesh location.
    */
    virtual T initial_value(const Eigen::Matrix<T, 3, 1>& x){
        return 0.0;
    }
    /**
    * User defined function: left boundary.
    *
    * Left boundary value, i.e. for x = lengths[0][0].
    *
    * @param y The \f$y\f$-coordinate.
    * @param z The \f$z\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T left(T y, T z, T t){
        return 0.0;
    }
    /**
    * @return The total number of spatial mesh nodes.
    */
    int n() {
        int out = m_params->n[0] + 1;
        if (m_dim == 2){
            out = (m_params->n[0] + 1)*(m_params->n[1] + 1);
        }
        if (m_dim == 3){
            out = (m_params->n[0] + 1)*(m_params->n[1] + 1)*
                (m_params->n[2] + 1);
        }
        return out;
    }
    /**
    *  @return The number of division along the \f$x\f$-axis.
    */
    int n_x(){
        return m_params->n[0];
    }
    /**
    * @return The number of division along the \f$y\f$-axis.
    */
    int n_y(){
        int n = 1;
        if (m_dim != 1){
            n = m_params->n[1];
        }
        return n;
    }
    /**
    * @return The number of division along the \f$z\f$-axis.
    */
    int n_z(){
        int n = 1;
        if (m_dim == 3){
            n = m_params->n[2];
        }
        return n;
    }
    /**
    * @return The total number of discrete time steps.
    */
    int n_t(){
        return m_params->nt;
    }
    /**
    *  Get the reference solution at a specified mesh location.
    *
    * @param x The \f$x\f$-, \f$y\f$-, and \f$z\f$-coordinates of the mesh
    *    node.
    * @param t The discrete time.
    * @return The reference solution at a specified mesh location.
    */
    virtual T reference(const Eigen::Matrix<T, 3, 1>& x, T t){
        return 0.0;
    }
    /**
    *  Set the value of the RHS vector at the indices of the boundary
    *  nodes for Dirichlet and Neumann boundary conditions.
    *
    *  The user defined functions at the boundaries are time-dependent and
    *  depend on the location of the boundary nodes on the plane defining
    *  the boundaries. We use function pointers to member functions,
    *  numerical::fdm::Problem::left(), numerical::fdm::Problem::right(),
    *  numerical::fdm::Problem::bottom(), numerical::fdm::Problem::top(),
    *  numerical::fdm::Problem::back(), and numerical::fdm::Problem::front()
    *  to define the value of the Dirichlet or Neumann boundary condition at
    *  each boundary nodes. The boundary nodes are passed by address from the
    *  Mesh instance, and stored in an array of addresses of l;ength 6, where
    *  the indices respectively represent the left (0), right (1), bottom (2),
    *  top (3), back (4) and front (5) boundaries. In order to pass the right
    *  coordinates to the user defined functions, we use a two-dimensional
    *  array of size (6, 2), where the first indice correspond to the boundary
    *  and the second to the indice of the spatial coordinates, i.e. 0 for
    *  \f$x\f$, 1 for \f$y\f$ and 2 for \f$z\f$.
    *
    *  For Dirichlet boundary conditions, we only have to set the value of the
    *  RHS vector at the boundary nodes equal to the specified value obtained
    *  from the user defined function.
    *
    *  For Neumann boundary conditions, it is a little bit more complicated
    *  since we have to approximate the normal derivative at the boundary
    *  nodes. In order to define the right normal unit vector (pointing
    *  outward of the domain) and "interior node", we define
    *  the left and right side of each boundaries, such that \f$\mathbf{n}=1\f$
    *  , and side = -1 for the boundaries that have the right side outside of
    *  the domain, i.e. right, top, and front. The opposite happens for the
    *  boundaries that have the left side outside the domain, i.e., left,
    *  bottom, and back have \f$\mathbf{n}=-1\f$, and side = 1. We also use
    *  the same indices to define the space increment in the direction normal
    *  to the boundary.
    *
    *  For a boundary node that is not an edge nor a corner, the RHS vector
    *  at the position of the node is expressed as:
    *
    *  **Left-side boundaries:**
    *
    *  Using \f$\textrm{bnd}\f$ as the index of the boundary for the normal
    *  direction, for example \f$dx[\textrm{bnd}=0]=\Delta x\f$ for the left
    *  boundary, i.e. \f$\textrm{bnd}=0\f$, we have:
    *
    *  \f[
    *    \textrm{RHS}[i] = u_n[i] + \theta\left(2d\alpha_m[i]g[i]
    *        dx[\textrm{bnd}]+
    *        \Delta tf[i]\right) + \left(1-\theta\right)\left(
    *    d\left(\alpha_p[i]\left(u_n[i+1]-u_n[i]\right)-
    *           \alpha_m[i]\left(u_n[i]-u_n[i+1]\right)\right)+
    *    2d\alpha_m[i]g_n[i]dx[\textrm{bnd}]+\Delta tf_n[i]\right)
    *  \f]
    *
    *  where
    *
    *  \f[
    *    \begin{align}
    *    \alpha_m[i] &= \frac{1}{2}\left(\alpha_\textrm{out} +
    *        \alpha[i]\right)\\
    *    \alpha_p[i] &= \frac{1}{2}\left(\alpha[i] + \alpha[i+1]\right)
    *    \end{align}
    *  \f]
    *
    *  and where \f$\alpha_\textrm{out}\f$ is the value of the diffusion
    *  coefficient outside of the left side of the boundary (outside of the
    *  domain). By default we set \f$\alpha_\textrm{out}=\alpha[i]\f$
    *
    *  **Right-side boundaries:**
    *
    *  Similarly for the right-side boundaries we have:
    *
    *  \f[
    *    \textrm{RHS}[i] = u_n[i] + \theta\left(-2d\alpha_p[i]g[i]
    *        dx[\textrm{bnd}]+
    *        \Delta tf[i]\right) + \left(1-\theta\right)\left(
    *    d\left(\alpha_p[i]\left(u_n[i-1]-u_n[i]\right)-
    *           \alpha_m[i]\left(u_n[i]-u_n[i-1]\right)\right)-
    *    2d\alpha_p[i]g_n[i]dx[\textrm{bnd}]+\Delta tf_n[i]\right)
    *  \f]
    *
    *  where
    *
    *  \f[
    *    \begin{align}
    *    \alpha_m[i] &= \frac{1}{2}\left(\alpha[i-1] + \alpha[i]\right)\\
    *    \alpha_p[i] &= \frac{1}{2}\left(\alpha[i] + \alpha_\textrm{out}\right)
    *    \end{align}
    *  \f]
    *
    *  and where \f$\alpha_\textrm{out}\f$ is the value of the diffusion
    *  coefficient outside of the right side of the boundary (outside of the
    *  domain). By default we set \f$\alpha_\textrm{out}=\alpha[i]\f$
    *
    *  With the sign and side variables we can simply write:
    *
    *  \f[
    *    \textrm{RHS}[i] = u_n[i] + \theta\left(\textrm{sign}2d\alpha[i]g[i]
    *        dx[\textrm{bnd}]+
    *        \Delta tf[i]\right) + \left(1-\theta\right)\left(
    *    d\left(\alpha_p[i]\left(u_n[i+\textrm{side}]-u_n[i]\right)-
    *           \alpha_m[i]\left(u_n[i]-u_n[i+\textrm{side}]\right)\right)+
    *    \textrm{sign}2d\alpha[i]g_n[i]dx[\textrm{bnd}]+\Delta tf_n[i]\right)
    *  \f]
    *
    *  where we used the default value for the outside diffusion coefficient.
    *
    *  For 2D and 3D models we have to define the boundary corners. Boundary
    *  edges have to be defined for 3D models.
    *
    *  @param rhs The rhs vector.
    *  @param u_n The solution vector at previous time step.
    *  @param alpha The rhs vector.
    *  @param f_n The source vector at privious time step.
    *  @param f The source vector.
    *  @param dx The space increment in the \f$x\f$-direction.
    *  @param dy The space increment in the \f$y\f$-direction.
    *  @param dz The space increment in the \f$z\f$-direction.
    *  @param dt The time increment.
    *  @param theta The scheme coefficient.
    *  @param t The discrete time.
    */
    void rhs_bc(Vec& rhs, Vec& u_n, const Vec& alpha, const Vec& f_n,
                const Vec& f, const T dx, const T dy, const T dz,
                const T dt, const T theta, const T t){
        // boundary node addresses
        const std::vector<int>& left_nodes = m_mesh->left();
        const std::vector<int>& right_nodes = m_mesh->right();
        const std::vector<int>& bottom_nodes = m_mesh->bottom();
        const std::vector<int>& top_nodes = m_mesh->top();
        const std::vector<int>& back_nodes = m_mesh->back();
        const std::vector<int>& front_nodes = m_mesh->front();
        // array of addresses to specific boundry nodes
        std::vector<int> nodes[6] = {
            left_nodes, right_nodes,
            bottom_nodes, top_nodes,
            back_nodes, front_nodes
            };
        // indices of the in-plane coordinates for each boundaries
        // 0 for x, 1 for y, 2 for z
        int ind[6][2] = {{1,2},{1,2},{0,2},{0,2},{0,1},{0,1}};
        T d;
        // space increment normal to the boundaries
        T d_i[6] = {dx, dx, dy, dy, dz, dz};
        // the number of boundary to set up depend on the dimension
        int n_bnd = 2 * m_dim;
        const Coord& c = coordinates();
        for(int bnd=0; bnd<n_bnd; bnd++){
            if (m_bc_types[bnd] == 0){  // Dirichlet
                for(int i=0; i<nodes[bnd].size(); i++){
                    // the rhs value is equal to the value returned by the
                    // boundary function
                    rhs[nodes[bnd][i]] = (this->*m_functions[bnd]) (
                        c[nodes[bnd][i]][ind[bnd][0]],
                        c[nodes[bnd][i]][ind[bnd][1]],
                        t + dt);
                }
            } else if(m_bc_types[bnd] == 1){  // Neumann
                T sign = 1.0, g, g_n, alpha_m, alpha_p;
                int k, side = 1;
                for(int i=0; i<nodes[bnd].size(); i++){  // scaled by 1/2
                    k = nodes[bnd][i];
                    g = (this->*m_functions[bnd]) (
                            c[k][ind[bnd][0]],
                            c[k][ind[bnd][1]],
                            t + dt);
                    g_n = (this->*m_functions[bnd]) (
                            c[k][ind[bnd][0]],
                            c[k][ind[bnd][1]],
                            t);
                    if (bnd % 2){  // right, top, front (right side outside)
                        sign = 1.0;  // outward --> normal unit vector
                        side = -1;
                        alpha_m = 0.5 * (alpha[k] + alpha[k-1]);
                        // TODO get alpha on the right of the boundary
                        alpha_p = 0.5 * (alpha[k] + alpha[k]);
                    } else{  // left, bottom, back (left side outside)
                        sign = -1.0;  // outward --> normal unit vector
                        side = 1;
                        // TODO get alpha on the left of the boundary
                        alpha_m = 0.5 * (alpha[k] + alpha[k]);
                        alpha_p = 0.5 * (alpha[k] + alpha[k+1]);
                    }
                    d = dt/d_i[bnd]/d_i[bnd];
                    rhs[k] = 0.5 * (
                        u_n[k] +
                        theta*(sign*2.0*d*alpha[k]*g*d_i[bnd] + dt*f[k]) +
                        (1.0-theta)*(
                            d*alpha_p*(u_n[k+side]-u_n[k])-
                            d*alpha_m*(u_n[k]-u_n[k+side])+
                            sign*2.0*d*alpha[k]*g_n*d_i[bnd] + dt*f_n[k])
                        );
                }
                // TODO: fix corners for 2D and 3D and fix edges for 3D
                // if corner: we use two dirivatives in 2D and 3 derivatives
                // in 3D. If edges, use two derivatives.
            } else{
                throw std::invalid_argument(
                "Not a valid boundary condition type.");
            }
        }
    }
    /**
    * User defined function: right boundary.
    *
    * Right boundary value, i.e. for x = lengths[0][1].
    *
    * @param y The \f$y\f$-coordinate.
    * @param z The \f$z\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T right(T y, T z, T t){
        return 0.0;
    }

    /**
    *  Get the computed solution at a specified mesh location.
    *
    * @param x The \f$x\f$-, \f$y\f$-, and \f$z\f$-coordinates of the mesh
    *    node.
    * @param t The discrete time.
    * @return The computed solution at a specified mesh location.
    */
    virtual T solution(const Eigen::Matrix<T, 3, 1>& x, T t){
        return 0.0;
    }
    /**
    * User defined function: source term.
    *
    * @param x The \f$x\f$-, \f$y\f$-, and \f$z\f$-coordinates of the mesh
    *    node.
    * @param t The discrete time.
    * @param u The state variable for nonlinear problems.
    * @return The source term at a specified mesh location.
    */
    virtual T source(const Eigen::Matrix<T, 3, 1>& x, T t, T u=0){
        return 0.0;
    }
    /**
    * @return The theta parameter for the finite difference scheme.
    */
    T theta(){
        return m_params->theta;
    }
    /**
    *
    */
    Vec t(){
        return m_t;
    }
    /**
    * User defined function: top boundary.
    *
    * Top boundary value, i.e. at y = lengths[1][1]. Only for 2D and 3D models.
    *
    * @param x The \f$x\f$-coordinate.
    * @param z The \f$z\f$-coordinate.
    * @param t The discrete time.
    * @return The boundary value at a specified mesh location.
    */
    virtual T top(T x, T z, T t){
        return 0.0;
    }
    /**
    *
    *  @return The initial value vector.
    */
    Vec u_0(){
        int num = n();
        Vec u_0(num), x(num), y(num), z(num);
        Coord coords = coordinates();
        for(int i=0; i<num; i++) {
            x[i] = coords[i][0];
            y[i] = coords[i][1];
            z[i] = coords[i][2];
            u_0[i] = initial_value(coords[i]);
        }
        return u_0;
    }
};

}  // namespace fdm

}  // namespace numerical

#endif  // PROBLEM_H


#include <iostream>
#include <iomanip>
#include <fstream>

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>

#include <vector>
#include <cmath>
#include <cstring>

#include <unistd.h>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "OneChQS.hpp"

#ifndef pi
#define pi 3.141592653589793238462643383279502884197169
#endif

#ifndef _CHIN_
#define _CHIN_

double chiN(int Nsites, double Lambda)
{

  double daux[3];
  daux[0]=1.0-pow( Lambda,-((double)(Nsites)+1.0) );
  daux[1]=sqrt( 1.0-pow(Lambda,-(2.0*(double)(Nsites)+1.0)) );
  daux[2]=sqrt( 1.0-pow(Lambda,-(2.0*(double)(Nsites)+3.0)) );

  return(daux[0]/(daux[1]*daux[2]));

}

#endif

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

  // Parameters for command-line passing (GetOpt)

  //char ModelOption[]="Anderson";
  //int ModelNo=0;

#include"ModelOpt.cpp"

  // NRG objects

  CNRGarray Aeig(2);

  CNRGbasisarray AeigCut(2);

  CNRGbasisarray Abasis(2);

  CNRGbasisarray SingleSite(3);

  // STL vector

  CNRGmatrix HN;
  CNRGmatrix Qm1fNQ;

  //CNRGmatrix MQQp1;

  CNRGmatrix* MatArray;
  int NumNRGarrays=3;
  // MatArray 0 is nd
  // MatArray 1 is cd
  // MatArray 2 is a test

  int KeepSz=0;

  double U,ed,Gamma;
  double Lambda;
  double HalfLambdaFactor;
  double Dband;
  int auxIn;
  int calcdens=0;
  vector <double> Params;
  vector <int> Indexes;

  double chi_m1,chi_N,daux[4];
  int Nsites;
  int Nsitesmax=5;
  int Nsites0=0;
  double DN=0.0;

  // Thermodynamics

  CNRGthermo Suscep;
  CNRGthermo Entropy;

  double TM=0.0;
  double betabar=0.727;
  double Temp=0.0;
  double Sus=0.0;
  double SusDot=1.0/8.0;
  vector<double> SuscepChain;

  double ndot=0.0;
  double Sz2=0.0;

  int Ncutoff=700;
  int UpdateBefCut=1;

  char arqSus[32],arqname[32];

  int ii,jj,i1,i2;

  // STL iterator:

  vector<double>::iterator diter;

  // outstream
  ofstream OutFile;
  // instream
  ifstream InFile;

  /////////////////////////////////////////////////////////////////
  /////////////////////////////////////////////////////////////////
  ////                    READ PARAMETERS                      ////
  /////////////////////////////////////////////////////////////////
  /////////////////////////////////////////////////////////////////

  U=0.5;
  ed=-0.5*U;
  Gamma=0.0282691;
  Lambda=2.5;

  InFile.open("nrg_input_OneChQS.dat");
  if (InFile.is_open())
    {
      InFile >> Nsitesmax;
      InFile >> Ncutoff;
      InFile >> U;
      InFile >> Gamma;
      InFile >> ed;
      InFile >> Lambda;
      InFile >> Dband;
      InFile >> auxIn;
      InFile >> UpdateBefCut;
      InFile >> calcdens;
    }
  else
    {
      cout << "can't open nrg_input_OneChQS.dat" << endl;
      exit(0);
   }

  InFile.close();

  // NRG chain
  CNRGchain chain1(Lambda,Nsitesmax);

  // New stuff
  strcpy(Suscep.ArqName,"SuscepImp1Ch_25_726.dat");
  strcpy(Suscep.ChainArqName,"SuscepChain1Ch.dat");
  Suscep.Calc=CalcSuscep;

  strcpy(Entropy.ArqName,"EntropyImp1Ch_25_726.dat");
  strcpy(Entropy.ChainArqName,"EntropyChain1Ch.dat");
  Entropy.Calc=CalcEntropy;

  // Check for SuscepChain.dat file. If it is there, read Temp, Sus
  // If not, change to calcdens=3, U=0, ed=0 Gamma=1
  if ( (calcdens==1)&&(ModelNo==1) )
    {
      cout << "Kondo model: Can't calculate spectral function" << endl;
      exit(0);
    }
  if (ModelNo==3) calcdens=3; // Chain calculation
  if (calcdens==2)
    {
      Suscep.CalcChain=false;
      Entropy.CalcChain=false;
      Suscep.Nsite0Chain=0; // Anderson chains
      Entropy.Nsite0Chain=0; // Anderson chains

      if ( (!FileExists(Suscep.ChainArqName))||
	   (!FileExists(Entropy.ChainArqName)) )
	{
	  cout << " Can't find chain files: " << endl
	  << "   " << Suscep.ChainArqName << endl
	  << "   " << Entropy.ChainArqName << endl;
	  cout << " Exiting... " << endl;
	  exit(0);
	}
      else
	  cout << " Found files "
	       << Suscep.ChainArqName << ", "
	       << Entropy.ChainArqName
	       << endl;

      //int ReadChainStatus=Suscep.ReadChain();
      //int ReadChainStatus2=Entropy.ReadChain();
      //if (ReadChainStatus==-1)calcdens=3;

    }
  if (calcdens==3)
    {
      Suscep.CalcChain=true;
      Entropy.CalcChain=true;
      U=0.0;ed=0.0;Gamma=0.0;
    }

  HalfLambdaFactor=0.5*(1.0+(1.0/Lambda));
  chi_m1=sqrt(2.0*Gamma/pi)/(sqrt(Lambda)*HalfLambdaFactor);

  double U_tilde=ed/(Lambda*HalfLambdaFactor);
  cout << "e1 = " << 0.5*U/(Lambda*HalfLambdaFactor) << endl;
  cout << "e2 = " << (ed+0.5*U)/(Lambda*HalfLambdaFactor) << endl;
  cout << "e3 = " << (2.0*ed+1.5*U)/(Lambda*HalfLambdaFactor) << endl;
  cout << "chi_m1  = " << chi_m1 << endl;

  cout << " Nsitesmax    = " << Nsitesmax-1 << endl;
  cout << " Ncutoff      = " << Ncutoff << endl;
  if (strcmp(ModelOption,"Anderson")==0){
    cout << " U            = " << U << endl;
    cout << " Gamma        = " << Gamma << endl;
    cout << " ed           = " << ed << endl;
  }
  if (strcmp(ModelOption,"Kondo")==0){
    cout << " JK            = " << Gamma << endl;
  }
  cout << " Lambda       = " << Lambda << endl;
  cout << " Dband        = " << Dband << endl;
  cout << " UpdateBefCut = " << UpdateBefCut << endl;
  cout << " calcdens     = " << calcdens << endl;

  /////////////////////////////////////////////////
  OutFile.open("NRG_in.txt");
  OutFile << "Begin NRG calculation. Model :" << ModelOption << endl;
  OutFile << " Nsitesmax    = " << Nsitesmax-1 << endl;
  OutFile << " Ncutoff      = " << Ncutoff << endl;
  if (strcmp(ModelOption,"Anderson")==0){
    OutFile << " U            = " << U << endl;
    OutFile << " Gamma        = " << Gamma << endl;
    OutFile << " ed           = " << ed << endl;
  }
  if (strcmp(ModelOption,"Kondo")==0){
    OutFile << " JK            = " << Gamma << endl;
  }
  OutFile << " Lambda       = " << Lambda << endl;
  OutFile << " Dband        = " << Dband << endl;
  OutFile << " UpdateBefCut = " << UpdateBefCut << endl;
  OutFile << " calcdens     = " << calcdens << endl;
  OutFile.close();

  /////////////////////////////////////////////////////////////////
  /////////////////////////////////////////////////////////////////
  ////                   END READ PARAMETERS                   ////
  /////////////////////////////////////////////////////////////////
  /////////////////////////////////////////////////////////////////

  // Set SingleSite

  OneChQS_SetSingleSite(SingleSite);

  // Allocate MatArray (hope this works!)
  MatArray = new CNRGmatrix [NumNRGarrays];

  // Test Kondo model

  Params.clear();

  switch (ModelNo)
    {
    case 0 :
      // Set initial CNRG array (N=-1)
      Params.push_back(U);
      Params.push_back(ed);
      Params.push_back(Lambda);
      Params.push_back(HalfLambdaFactor);
      //OneChQS_SetAndersonHm1(Params,&Aeig, &Qm1fNQ, &MQQp1);
      OneChQS_SetAndersonHm1(Params,&Aeig, &Qm1fNQ, MatArray);

      Nsites0=0; // Adding N=0 site
      chi_N=chi_m1; // tunneling to site N=0 (hi_N connects sites N and N+1)
      SusDot=1.0/8.0;
      // Local Susceptibility,Entropy
      Suscep.dImpValue=SusDot;
      Entropy.dImpValue=2.0*log(2.0);
      break;
    case 1 :
      Params.push_back(Gamma); // JK
      OneChQS_SetKondoH0(Params,&Aeig,&Qm1fNQ,&SingleSite );
      Nsites0=1; // Adding N=1 site
      chi_N=chiN(0,Lambda);
      SusDot=1.0/4.0; // check this difference
      // Local Susceptibility,Entropy
      Suscep.dImpValue=SusDot;
      Entropy.dImpValue=log(2.0);
      break;
    case 2 :
      cout << " Not implemented yet..." << endl;
      exit(0);
      break;
    case 3 :
      OneChQS_SetH0Chain(&Aeig, &Qm1fNQ);
      Nsites0=1; // Adding N=1 site
      chi_N=chiN(0,Lambda); // tunneling to site N=1
                            //(chi_N connects sites N and N+1)
      // Local Susceptibility,Entropy
      Suscep.CalcChain=true;
      Entropy.CalcChain=true;
      Suscep.dImpValue=0.0;
      Entropy.dImpValue=0.0;
      break;
    case 5 :
      OneChQS_SetSMM_Hm1(Params,&Aeig, &Qm1fNQ,MatArray);
      Nsites0=0; // N=0 site only
      chi_N=chiN(0,Lambda); // tunneling to site N=1
                            //(chi_N connects sites N and N+1)
      KeepSz=1;
      calcdens=0;
      break;
    default :
      cout << " Model not implemented. Exiting... " << endl;
      exit(0);
      break;
    }
  //end switch models

  // Set AeigCut
  AeigCut.ClearAll();
  AeigCut=CutStates(&Aeig, Ncutoff);

 // Set operators to update

  // Will this work??? Yes!!

//   MatArray[0].NeedOld=false;
//   MatArray[0].CheckForMatEl=OneChQS_cd_check;
//   MatArray[0].CalcMatEl=OneChQS_fN_MatEl;

  //////////////
  Nsites=Nsites0-1; //Nsites=-1 or 0
  DN=HalfLambdaFactor*pow(Lambda,(-(Nsites-1)/2.0) );
  TM=DN/betabar;
  cout << "DN = " << DN << "TM = " << TM << endl;

  if (calcdens==1)
    {
      MatArray[0].NeedOld=true;
      MatArray[0].CheckForMatEl=OneChQS_nd_check;
      MatArray[0].CalcMatEl=OneChQS_nd_MatEl;

      MatArray[1].NeedOld=true;
      MatArray[1].CheckForMatEl=OneChQS_cd_check;
      MatArray[1].CalcMatEl=OneChQS_cd_MatEl;

      // Sep 08: calculating <Sz2>
      MatArray[2].NeedOld=true;
      MatArray[2].CheckForMatEl=OneChQS_nd_check;
      MatArray[2].CalcMatEl=OneChQS_nd_MatEl;

      Params.clear();
      Params.push_back(betabar);
      ndot=CalcOpAvg(Params,&AeigCut,&MatArray[0],true,1);
      cout << " ed = " << ed << " Initial ndot = " << ndot << endl;
      Sz2=CalcOpAvg(Params,&AeigCut,&MatArray[2],true,1);
      cout << " Sz2 = " << Sz2 << " Initial Sz2 = " << Sz2 << endl;

      for (int ibl=0;ibl<MatArray[1].NumMatBlocks();ibl++)
	MatArray[1].PrintMatBlock(ibl);

      // Interface with old functions.
//       AeigCut.SaveQSParameters();
//       MatArray[1].SaveInOldFormat();

    }
  if ( (ModelNo==3)&&(calcdens==3) ){
    TM=DN/betabar;
    Params.clear();
    Params.push_back(betabar);
    // New Stuff
    //Suscep.ReadNChainValue(Nsites,Nsites0); // Chain starts at N=0
    // Wrong: Nsites=Nsites0-1 AND it's a chain calculation!
    Suscep.AddValue(Params,&Aeig,1,true,TM);
    Suscep.SaveNValue(Nsites,Nsites0);
    //Entropy.ReadNChainValue(Nsites,Nsites0);
    // Wrong: Nsites=Nsites0-1 AND it's a chain calculation!
    Entropy.AddValue(Params,&Aeig,1,true,TM);
    Entropy.SaveNValue(Nsites,Nsites0);
  }

  // Jul 09: testing CNRGchain

  //chain1.SetChainWilson(Nsitesmax);

  // Loop on Nsites

  // Entering calculation of H_Nsites0
  Nsites=Nsites0;
  while (Nsites<Nsitesmax)
    {
      cout << "Nsites = " << Nsites << endl;
      cout << "Old Nshell = " << Aeig.Nshell << endl;
      cout << "BEG Nshell = " << Aeig.Nshell+1 << endl;
      cout << "chi_(N-1) = " << chi_N << endl;
      DN=HalfLambdaFactor*pow(Lambda,(-(Nsites-1)/2.0) );
      TM=DN/betabar;
      cout << "DN = " << DN << "TM = " << TM << endl;

      // Build Abasis

      /////QS_BuildBasis(&AeigCut,&Abasis,&SingleSite,UpdateBefCut);
      QS_BuildBasis(&AeigCut,&Abasis,&SingleSite,UpdateBefCut,KeepSz);

      //Abasis.PrintAll();

      cout << "Basis Nstates = " << Abasis.Nstates() << endl;

      // Build and diagonalize H_N+1 (general): Build Aeigv
      // 1 - Get old matrix elements
      // 2 - Build and diagonalize H_N+1

      cout << "Diagonalizing HN... " << endl;

      OneChQS_DiagHN(Qm1fNQ,Abasis,SingleSite,Aeig,0.0,chi_N,Lambda);

      cout << "... done diagonalizing HN. " << endl;

      Aeig.PrintEn();
//       for (int ibl=0;ibl<Aeig.NumBlocks();ibl++)
// 	Aeig.PrintBlock(ibl);

      if (UpdateBefCut==1)
	{
	  cout << "Updating matrices before cutting... " << endl;
	  OneChQS_UpdateQm1fQ(&Qm1fNQ,&Aeig,&Abasis,&SingleSite);
	  cout << "... done updating matrices. " << endl;
	}

      // Calculate Susceptibility

      if ( (calcdens==2)||(calcdens==3) ){
	TM=DN/betabar;
	Params.clear();
	Params.push_back(betabar);
	// New Stuff
	//Suscep.ReadNChainValue(Nsites,Nsites0); // Chain starts at N=0 WRONG
	Suscep.ReadNChainValue(Nsites-Nsites0,0);// Chain starts at N=1!
	Suscep.AddValue(Params,&Aeig,1,true,TM);
	Suscep.SaveNValue(Nsites,Nsites0);
	//Entropy.ReadNChainValue(Nsites,Nsites0); // WRONG
	Entropy.ReadNChainValue(Nsites-Nsites0,0);
	Entropy.AddValue(Params,&Aeig,1,true,TM);
	Entropy.SaveNValue(Nsites,Nsites0);
      }

      // Eliminate states, update matrices, calculate stuff

      // Try this here!
      cout << "Cutting states..." << endl;

      AeigCut.ClearAll();
      AeigCut=CutStates(&Aeig, Ncutoff);

      cout << "... done cutting states." << endl;

      // Calculate new matrix elements using the CUT basis:
      // update Qm1fNQ, Qm1cdQ, etc.

      if (UpdateBefCut==0)
	{
	  cout << "Updating matrices after cutting... " << endl;
	  //OneChQS_UpdateMatrixAfterCutting(&Qm1fNQ,&MQQp1,
	  //          &AeigCut,&Abasis,&SingleSite);
	  OneChQS_UpdateMatrixAfterCutting(&Qm1fNQ,&MatArray[1],
					   &AeigCut,&Abasis,&SingleSite);
	  cout << "... done updating matrices. " << endl;

	  // New Update

	  // Calculate Spectral density, <ndot>
	  if (calcdens==1)
	    {
// 	      UpdateMatrices(&SingleSite,&AeigCut,
// 			     &Abasis,MatArray, 2);
// Calc Sz2 too now: update 3 mats
	      UpdateMatrices(&SingleSite,&AeigCut,
			     &Abasis,MatArray, 3);

// 	      for (int ibl=0;ibl<MatArray[1].NumMatBlocks();ibl++)
// 		MatArray[1].PrintMatBlock(ibl);

	      if (Nsites==0)
		for (int ist=0;ist<MatArray[1].Nstates();ist++)
		  {
		    for (int jst=0;jst<MatArray[1].Nstates();jst++)
		      cout << MatArray[1].GetMatEl(ist,jst) << "  ";
		    cout << endl;
		  }

	      // Changed to NumMats to 2
	      TM=DN/betabar;
	      Params.clear();
	      Params.push_back(betabar);
	      ndot=CalcOpAvg(Params,&AeigCut,&MatArray[0],true,1);

	      cout << " ed = " << ed;
	      cout.precision(10);
	      cout << scientific << " Temp = " << TM;
	      cout << resetiosflags (ios_base::floatfield);
	      cout << " ndot = " << ndot << endl;

	      Sz2=CalcOpAvg(Params,&AeigCut,&MatArray[2],true,1);
	      cout << " ed = " << ed;
	      cout.precision(10);
	      cout << scientific << " Temp = " << TM;
	      cout << resetiosflags (ios_base::floatfield);
	      cout << " Sz2 = " << Sz2 << endl;

	      // Interface with old functions.
	      AeigCut.SaveQSParameters();
	      MatArray[1].SaveInOldFormat();
	    }
	  // end if calcdens==1

	}
      // end if UpdateCefCut==0

      cout << "END Nshell = " << Aeig.Nshell << endl;

     // Update chi_N

//       daux[0]=(double)( 1.0-pow(Lambda,(-(Nsites+1))) );
//       daux[1]=(double)sqrt( 1.0-pow(Lambda,-(2*(Nsites+1)-1)) );
//       daux[2]=(double)sqrt( 1.0-pow(Lambda,-(2*(Nsites+1)+1)) );
//       daux[3]=0.5*(1.0+(1.0/Lambda))*(double)sqrt(Lambda);

//       chi_N=daux[0]/(daux[1]*daux[2]);
// new (Jul 09)
      chi_N=chain1.GetChin(Nsites);

      cout << "chi_N = " << chi_N << endl;
//       cout << "Chi(N="<< Nsites <<") = " << chiN(Nsites,Lambda) << endl;

      // Update sites
      Nsites++;
    }
  // end NRG loop

  // De-allocate MatArray
  delete[] MatArray;

  cout << "=== Calculation Finished! ==== "<< endl;
  OutFile.open("NRG_end.txt");
  OutFile << "END NRG calculation" << endl;
  OutFile.close();

}
// END code

////////////////////////////////
///                          ///
///       Trash can          ///
///                          ///
////////////////////////////////


#include <omp.h>
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <iostream>
#include <limits>
#include <vector>
#include "Python.h"
#include "arrayobject.h"
#include "kdtree.h"
#include "progress_bar.h"
#include "python_util.h"
#include "timer.h"

using namespace Eigen;
using namespace std;

template <typename T>
void estimate_normals(vector<T>* eigenvectors, vector<T>* eigenvalues,
                      vector<int>* neighborhood_sizes, const vector<T>& points,
                      const std::size_t k, const T d_max,
                      const std::vector<int>* subsample_indices,
                      int num_eigen = 1, bool verbose = true,
                      int num_procs = -1) {
  size_t num_points = points.size() / 3;
  Map<const Matrix<T, Dynamic, 3, RowMajor> > P(&points[0], num_points, 3);

  if (num_procs < 0) num_procs = omp_get_num_procs();
  if (verbose)
    cout << "(estimate_normals) using " << num_procs << " threads" << endl;

  // organize points into k-d tree
  pointkd::BuildParams build_params;
  build_params.num_proc = num_procs;
  double build_time = getTime();
  pointkd::KdTree<T, 3> tree(points, build_params);
  build_time = getTime() - build_time;
  if (verbose) {
    cout << "(estimate_normals) ";
    cout << "k-d tree build time (s): " << build_time << endl;
  }

  int num_normals = num_points;
  if (subsample_indices != NULL)
    num_normals = subsample_indices->size();

  // pre-allocate space for results (assume num_eigen either 1 or 3)
  if (eigenvectors) eigenvectors->resize(num_normals * num_eigen * 3);
  if (eigenvalues) eigenvalues->resize(num_normals * num_eigen);
  if (neighborhood_sizes) neighborhood_sizes->resize(num_normals);

  tbb::task_scheduler_init(1);  // use just 1 thread for k-d tree queries
  omp_set_num_threads(num_procs);

  if (verbose) {
    cout << "(estimate_normals) neighborhood parameters: " << endl;
    cout << "  k = " << k << endl;
    cout << "  r = " << d_max << endl;
  }

  ProgressBar<int> bar((int)num_normals);

  double pca_time = getTime();

#pragma omp parallel for schedule(static, 1000)
  for (int i = 0; i < (int)num_normals; i++) {
    if (verbose && omp_get_thread_num() == 0 && i % 1000 == 0) {
      bar.update(i);
      cout << "\r" << bar.get_string();
    }

    int i_ = i;
    if (subsample_indices != NULL)
      i_ = (*subsample_indices)[i];

    // not using KNearestNeighborsSelf, because we want to include the
    // current point in the normal estimation calculation
    pointkd::Indices indices;
    vector<T> q(P.data() + i_ * 3, P.data() + (i_ + 1) * 3);
    if (k == -1)
      tree.RNearNeighbors(indices, &q[0], d_max);
    else
      tree.KNearestNeighbors(indices, &q[0], k, d_max);

    Matrix<T, Dynamic, 3, RowMajor> X(indices.size(), 3);
    for (size_t j = 0; j < indices.size(); j++) X.row(j) = P.row(indices[j]);
    X.rowwise() -= X.colwise().mean();
    Matrix<T, 3, 3> C = X.transpose() * X;
    C /= (T)indices.size();
    SelfAdjointEigenSolver<Matrix<T, 3, 3> > es(C);

    // record results of PCA
    if (eigenvectors) {
      if (num_eigen == 1) {
        Map<Matrix<T, 3, 1> > temp(&(*eigenvectors)[i * 3]);
        temp = es.eigenvectors().col(0);
      } else {  // num_eigen == 3
        Map<Matrix<T, 3, 3> > temp(&(*eigenvectors)[i * 9]);
        temp = es.eigenvectors();
      }
    }
    if (eigenvalues) {
      if (num_eigen == 1) {
        (*eigenvalues)[i] = es.eigenvalues()(0);
      } else {  // num_eigen == 1
        Map<Matrix<T, 3, 1> > temp(&(*eigenvalues)[i * 3]);
        temp = es.eigenvalues();
      }
    }
    if (neighborhood_sizes)
      (*neighborhood_sizes)[i] = indices.size();
  }

  pca_time = getTime() - pca_time;

  if (verbose) {
    bar.update((int)num_normals);
    cout << "\r" << bar.get_string() << endl;
    cout << "(estimate_normals) PCA time (s): " << pca_time << std::endl;
  }
}

template <typename T>
struct NumpyTypeNumber {};

template <>
struct NumpyTypeNumber<float> {
  static const int value = NPY_FLOAT32;
};

template <>
struct NumpyTypeNumber<double> {
  static const int value = NPY_FLOAT64;
};

template <typename T>
void estimate_normals(PyObject*& out1, PyObject*& out2, PyObject*& out3,
                      const Array2D& arr, int k, float r,
                      std::vector<int>* ptr_subsample_indices,
                      bool output_eigenvalues, bool output_all_eigenvectors,
                      bool output_neighborhood_sizes, bool verbose,
                      int num_procs) {
  vector<T> points;
  VectorFromArray2D(points, arr);

  int num_normals = arr.m;
  if (ptr_subsample_indices != NULL)
    num_normals = ptr_subsample_indices->size();

  int num_eigen = 1;
  int out1_ndim = 2;
  int out2_ndim = 1;
  npy_intp out1_dims[3] = {num_normals, 3, -1};
  npy_intp out2_dims[2] = {num_normals, -1};
  if (output_all_eigenvectors) {
    num_eigen = 3;
    out1_ndim = 3;
    out2_ndim = 2;
    out1_dims[2] = 3;
    out2_dims[1] = 3;
  }

  int out3_ndim = 1;
  int out3_typenum = NPY_INT;
  npy_intp out3_dims[2] = {num_normals, -1};

  int typenum = NumpyTypeNumber<T>::value;

  out1 = NULL;
  out2 = NULL;
  out3 = NULL;
  vector<T> evecs;
  vector<T> evals;
  vector<int> nbhd_sizes;
  estimate_normals<T>(&evecs, output_eigenvalues == 1 ? &evals : NULL,
                      output_neighborhood_sizes == 1 ? &nbhd_sizes : NULL,
                      points, k, r, ptr_subsample_indices, num_eigen, verbose,
                      num_procs);
  out1 = PyArray_EMPTY(out1_ndim, out1_dims, typenum, false);
  copy(evecs.begin(), evecs.end(), (T*)PyArray_DATA((PyArrayObject*)out1));
  if (output_eigenvalues) {
    out2 = PyArray_EMPTY(out2_ndim, out2_dims, typenum, false);
    copy(evals.begin(), evals.end(), (T*)PyArray_DATA((PyArrayObject*)out2));
  }
  if (output_neighborhood_sizes) {
    out3 = PyArray_EMPTY(out3_ndim, out3_dims, out3_typenum, false);
    copy(nbhd_sizes.begin(), nbhd_sizes.end(),
         (int*)PyArray_DATA((PyArrayObject*)out3));
  }
}

static char estimate_normals_usage[] =
    "Estimates normals at all points using principal component analysis\n"
    "(PCA). Specifically, computes eigenvectors of the covariance matrix C\n"
    "based on the local neighborhoods of each point.\n"
    "\n"
    ".. math::\n"
    "   C = \\frac{1}{N}\\sum_{i=1}^{N}{(p_i-\\mu)(p_i-\\mu)^T}\n"
    "\n"
    "where :math:`p_1 ... p_N` are the points in a given neighborhood and\n"
    ":math:`\\mu` is the centroid of the points.\n"
    "\n"
    "\n"
    "Parameters\n"
    "----------\n"
    "points : 3-column numpy array of type float32 or float64\n"
    "    Input point cloud.\n"
    "k : int\n"
    "    Number of neighbors to use.\n"
    "    For pure r-near neighborhoods, set this to -1.\n"
    "r : float\n"
    "    Use neighbors within r of query point.\n"
    "    For pure k-nearest neighborhoods, set this to np.inf.\n"
    "subsample : 1-d array of bool or int, optional\n"
    "    Optionally estimate normals at subset of points specified by a\n"
    "    boolean mask having length equal to the number of points, or by\n"
    "    integer indices into the array of points. Default: None.\n"
    "output_eigenvalues : bool, optional\n"
    "    Default: False.\n"
    "output_all_eigenvectors : bool, optional\n"
    "    Default: False.\n"
    "output_neighborhood_sizes : bool, optional\n"
    "    Default: False.\n"
    "verbose : bool, optional\n"
    "    Default: True.\n"
    "num_procs : int, optional\n"
    "    Default: use all processors.\n"
    "\n"
    "Returns\n"
    "-------\n"
    "results : ndarray or tuple of ndarray's\n"
    "    PCA results. The tuple return type is triggered if any of\n"
    "    output_eigenvalues or output_neighborhood_sizes is True.\n"
    "\n"
    "\n"
    "The following table summarizes all potential return types. Here m is the\n"
    "number of PCAs performed. If subsample is not None, then m is the size\n"
    "of subsample. Otherwise, m is just the number of input points.\n"
    "\n"
    "+---------------------------+------------------------------------------+\n"
    "|                           |output_all_eigenvectors                   |\n"
    "|                           +----------------+-------------------------+\n"
    "|                           |F               |T                        |\n"
    "+======================+====+================+=========================+\n"
    "|**output_eigenvalues**|F, F|m x 3           |m x 3 x 3                |\n"
    "|                      +----+----------------+-------------------------+\n"
    "|**output_nbhd_sizes** |F, T|(m x 3, None, m)|(m x 3 x 3, None, m)     |\n"
    "|                      +----+----------------+-------------------------+\n"
    "|                      |T, F|(m x 3, m, None)|(m x 3 x 3, m x 3, None) |\n"
    "|                      +----+----------------+-------------------------+\n"
    "|                      |T, T|(m x 3, m, m)   |(m x 3 x 3, m x 3, m)    |\n"
    "+----------------------+----+----------------+-------------------------+\n"
    "\n"
    "Note:\n"
    "    * The j-th eigenvector of the i-th PCA result is given by\n"
    "      indices [i, j, :].\n"
    "    * eigenvectors are sorted in order of increasing eigenvalue\n"
    "\n";

static PyObject* estimate_normals_wrapper(PyObject* self, PyObject* args,
                                          PyObject* kwargs) {
  PyObject* p = NULL;
  PyObject* subsample = Py_None;
  int k;
  float r;
  int output_eigenvalues = 0;
  int output_all_eigenvectors = 0;
  int output_neighborhood_sizes = 0;
  int verbose = 1;
  int num_procs = -1;
  static char* keywords[] = {
      "points", "k", "r", "subsample", "output_eigenvalues",
      "output_all_eigenvectors", "output_neighborhood_sizes", "verbose",
      "num_procs", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Oif|Oiiiii", keywords, &p, &k,
                                   &r, &subsample, &output_eigenvalues,
                                   &output_all_eigenvectors,
                                   &output_neighborhood_sizes, &verbose,
                                   &num_procs)) {
    PyErr_SetString(PyExc_RuntimeError, "Failed to parse inputs");
    return NULL;
  }

  // check k and r
  if (r <= 0.0f) {
    PyErr_SetString(PyExc_ValueError, "r must be positive");
    return NULL;
  } else if (k == 0) {
    PyErr_SetString(PyExc_ValueError, "k cannot be zero");
    return NULL;
  } else if (k < 0 && r == std::numeric_limits<float>::infinity()) {
    PyErr_SetString(PyExc_ValueError, "invalid combo: r == inf and k < 0");
    return NULL;
  }

  Array2D arr;
  if (!CheckAndExtractArray2D(arr, p)) {
    if (!PyErr_Occurred())
      PyErr_SetString(PyExc_TypeError,
                      "points must be interpretable as 0, 1 or 2-d array");
    return NULL;
  }

  std::vector<int> subsample_indices;
  std::vector<int>* ptr_subsample_indices = NULL;
  if (subsample != Py_None) {
    if (!CheckAndExtractIndices(subsample_indices, subsample, arr.m)) {
      if (!PyErr_Occurred())
        PyErr_SetString(PyExc_TypeError,
                        "subsample must be interpretable as specifying a "
                        "subset of points");
      return NULL;
    }
    ptr_subsample_indices = &subsample_indices;
  }

  PyObject* out1 = NULL;
  PyObject* out2 = NULL;
  PyObject* out3 = NULL;
  if (arr.type_num == NPY_FLOAT32) {
    estimate_normals<float>(out1, out2, out3, arr, k, r, ptr_subsample_indices,
                            (bool)output_eigenvalues,
                            (bool)output_all_eigenvectors,
                            (bool)output_neighborhood_sizes, (bool)verbose,
                            num_procs);
  } else if (arr.type_num == NPY_FLOAT64) {
    estimate_normals<double>(out1, out2, out3, arr, k, r, ptr_subsample_indices,
                             (bool)output_eigenvalues,
                             (bool)output_all_eigenvectors,
                             (bool)output_neighborhood_sizes, (bool)verbose,
                             num_procs);
  } else {
    PyErr_SetString(PyExc_TypeError, "points must be float32 or float64");
    return NULL;
  }

  if (out2 == NULL && out3 == NULL) {
    return out1;
  } else {
    PyObject* out = PyTuple_New(3);
    PyTuple_SetItem(out, 0, out1);
    PyTuple_SetItem(out, 1, out2 != NULL ? out2 : Py_None);
    PyTuple_SetItem(out, 2, out3 != NULL ? out3 : Py_None);
    return out;
  }
}

static PyMethodDef methods[] = {
    {"estimate_normals", (PyCFunction)estimate_normals_wrapper,
     METH_VARARGS | METH_KEYWORDS, estimate_normals_usage},
    {NULL, NULL, 0, NULL}};

#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef module_def = {PyModuleDef_HEAD_INIT,
                                        "estimate_normals",
                                        NULL,
                                        -1,
                                        methods,
                                        NULL,
                                        NULL,
                                        NULL,
                                        NULL};

PyMODINIT_FUNC PyInit_estimate_normals(void) {
  PyObject* module = PyModule_Create(&module_def);
#else
PyMODINIT_FUNC initestimate_normals(void) {
  (void)Py_InitModule("estimate_normals", methods);
#endif

  import_array();

#if PY_MAJOR_VERSION >= 3
  return module;
#endif
}


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

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

// This file was modified by Oracle on 2014.
// Modifications copyright (c) 2014, Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

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

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_AZIMUTH_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_AZIMUTH_HPP

#include <boost/geometry/core/cs.hpp>
#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/radian_access.hpp>
#include <boost/geometry/core/tags.hpp>

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

#include <boost/geometry/algorithms/not_implemented.hpp>
#include <boost/geometry/algorithms/detail/vincenty_inverse.hpp>

namespace boost { namespace geometry
{

// An azimuth is an angle between a vector/segment from origin to a point of
// interest and a reference vector. Typically north-based azimuth is used.
// North direction is used as a reference, angle is measured clockwise
// (North - 0deg, East - 90deg). For consistency in 2d cartesian CS
// the reference vector is Y axis, angle is measured clockwise.
// http://en.wikipedia.org/wiki/Azimuth

#ifndef DOXYGEN_NO_DISPATCH
namespace detail_dispatch
{

template <typename ReturnType, typename Tag>
struct azimuth
    : not_implemented<Tag>
{};

template <typename ReturnType>
struct azimuth<ReturnType, geographic_tag>
{
    template <typename P1, typename P2, typename Spheroid>
    static inline ReturnType apply(P1 const& p1, P2 const& p2, Spheroid const& spheroid)
    {
        return geometry::detail::vincenty_inverse<ReturnType>
                    ( get_as_radian<0>(p1), get_as_radian<1>(p1),
                      get_as_radian<0>(p2), get_as_radian<1>(p2),
                      spheroid ).azimuth();
    }

    template <typename P1, typename P2>
    static inline ReturnType apply(P1 const& p1, P2 const& p2)
    {
        return apply(p1, p2, srs::spheroid<ReturnType>());
    }
};

template <typename ReturnType>
struct azimuth<ReturnType, spherical_equatorial_tag>
{
    template <typename P1, typename P2, typename Sphere>
    static inline ReturnType apply(P1 const& p1, P2 const& p2, Sphere const& /*unused*/)
    {
        // http://williams.best.vwh.net/avform.htm#Crs
        ReturnType dlon = get_as_radian<0>(p2) - get_as_radian<0>(p1);
        ReturnType cos_p2lat = cos(get_as_radian<1>(p2));

        // An optimization which should kick in often for Boxes
        //if ( math::equals(dlon, ReturnType(0)) )
        //if ( get<0>(p1) == get<0>(p2) )
        //{
        //    return - sin(get_as_radian<1>(p1)) * cos_p2lat);
        //}

        // "An alternative formula, not requiring the pre-computation of d"
        // In the formula below dlon is used as "d"
        return atan2(sin(dlon) * cos_p2lat,
            cos(get_as_radian<1>(p1)) * sin(get_as_radian<1>(p2))
            - sin(get_as_radian<1>(p1)) * cos_p2lat * cos(dlon));
    }

    template <typename P1, typename P2>
    static inline ReturnType apply(P1 const& p1, P2 const& p2)
    {
        return apply(p1, p2, 0); // dummy model
    }
};

template <typename ReturnType>
struct azimuth<ReturnType, spherical_polar_tag>
    : azimuth<ReturnType, spherical_equatorial_tag>
{};

template <typename ReturnType>
struct azimuth<ReturnType, cartesian_tag>
{
    template <typename P1, typename P2, typename Plane>
    static inline ReturnType apply(P1 const& p1, P2 const& p2, Plane const& /*unused*/)
    {
        ReturnType x = get<0>(p2) - get<0>(p1);
        ReturnType y = get<1>(p2) - get<1>(p1);

        // NOTE: azimuth 0 is at Y axis, increasing right
        // as in spherical/geographic where 0 is at North axis
        return atan2(x, y);
    }

    template <typename P1, typename P2>
    static inline ReturnType apply(P1 const& p1, P2 const& p2)
    {
        return apply(p1, p2, 0); // dummy model
    }
};

} // detail_dispatch
#endif // DOXYGEN_NO_DISPATCH

#ifndef DOXYGEN_NO_DETAIL
namespace detail
{

/// Calculate azimuth between two points.
/// The result is in radians.
template <typename ReturnType, typename Point1, typename Point2>
inline ReturnType azimuth(Point1 const& p1, Point2 const& p2)
{
    return detail_dispatch::azimuth
            <
                ReturnType,
                typename geometry::cs_tag<Point1>::type
            >::apply(p1, p2);
}

/// Calculate azimuth between two points.
/// The result is in radians.
template <typename ReturnType, typename Point1, typename Point2, typename Model>
inline ReturnType azimuth(Point1 const& p1, Point2 const& p2, Model const& model)
{
    return detail_dispatch::azimuth
            <
                ReturnType,
                typename geometry::cs_tag<Point1>::type
            >::apply(p1, p2, model);
}

} // namespace detail
#endif // DOXYGEN_NO_DETAIL

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_AZIMUTH_HPP


/// General Matrix Multiplication
#include <HElib/FHE.h>
#include <HElib/FHEContext.h>
#include <HElib/EncryptedArray.h>
#include <HElib/NumbTh.h>

#include "SMP/Matrix.hpp"
#include "SMP/Timer.hpp"
#include "SMP/HElib.hpp"
#include "SMP/literal.hpp"
#include "SMP/network/net_io.hpp"
#include "SMP/SMPServer.hpp"

#include <boost/asio.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <iostream>
#include <numeric>
#include <list>
using boost::asio::ip::tcp;
constexpr int REPEAT = 1;
std::atomic<int> global_counter(0);

inline long round_div(long a, long b) {
    return (a + b - 1) / b;
}

void zero(Matrix &mat) {
    for (long i = 0; i < mat.NumRows(); i++)
        for (long j = 0; j < mat.NumCols(); j++)
            mat[i][j] = 0;
}

void randomize(Matrix &mat, long p = 3) {
    for (long i = 0; i < mat.NumRows(); i++)
        for (long j = 0; j < mat.NumCols(); j++)
            mat[i][j] = NTL::RandomBnd(p);
}

void fill_compute(Matrix& mat,
				  long row_blk,
				  long col,
                  const std::vector<long> &inner_prod,
                  const EncryptedArray *ea)
{
    const long l = ea->size();
    assert(inner_prod.size() == l);
	const bool is_vec = mat.NumRows() == 1;
	const long row_start = is_vec ? 0 : row_blk * l;
	const long col_start = is_vec ? row_blk * l : col;

    for (long ll = 0; ll < l; ll++) {
        long computed = inner_prod[ll];
		if (!is_vec) {
			long row = row_start + ll;
			if (row < mat.NumRows())
				mat.put(row, col, computed);
			else
				break;
		} else {
			long col = col_start + ll;
			if (col < mat.NumCols())
				mat.put(0, col, computed);
			else
				break;
		}
    }
}

struct ClientBenchmark {
    std::vector<double> pack_times;
    std::vector<double> enc_times;
    std::vector<double> dec_times;
    std::vector<double> unpack_times;
    std::vector<double> total_times;
    int ctx_sent, ctx_recv;
};
ClientBenchmark clt_ben;

struct ServerBenchmark {
    std::vector<double> eval_times;
};
ServerBenchmark srv_ben;

void play_client(tcp::iostream &conn,
                 FHESecKey &sk,
                 FHEcontext &context,
                 const long n1,
                 const long n2,
                 const long n3) {
	//* Convert to evalution key.
	//* This function is not provided by the origin HElib. Checkout our fork.
    sk.convertToSymmetric();
    FHEPubKey ek(sk);
    conn << ek;
    const EncryptedArray *ea = context.ea;
    const long l = ea->size();
    const long d = ea->getDegree();

    Matrix A, B, ground_truth;
    A.SetDims(n1, n2);
    B.SetDims(n2, n3);
    NTL::SetSeed(NTL::to_ZZ(123));
    randomize(A, ek.getPtxtSpace());
    randomize(B, ek.getPtxtSpace());
    ground_truth = mul(A, B);
    /// print grouth truth for debugging
    const long MAX_X1 = round_div(A.NumRows(), l);
    const long MAX_Y1 = round_div(A.NumCols(), d);
    const long MAX_X2 = round_div(B.NumCols(), l);

    std::vector<std::vector<Ctxt>> uploading;
    uploading.resize(MAX_X1, std::vector<Ctxt>(MAX_Y1, Ctxt(sk)));
	double enc_time = 0.;
    double pack_time = 0.;
	/// encrypt matrix
	NTL::ZZX packed_poly;
	for (int x = 0; x < MAX_X1; x++) {
		for (int k = 0; k < MAX_Y1; k++) {
			internal::BlockId blk = {x, k};
			double one_pack_time = 0.;
			double one_enc_time = 0.;
			auto block = internal::partition(A, blk, ea, false);
			{/// packing
				AutoTimer timer(&one_pack_time);
				rawEncode(packed_poly, block.polys, context);
			}
			{/// encryption
				AutoTimer timer(&one_enc_time);
				sk.Encrypt(uploading[x][k], packed_poly);
			}
			pack_time += one_pack_time;
			enc_time += one_enc_time;
		}
	}
    clt_ben.pack_times.push_back(pack_time);
    clt_ben.enc_times.push_back(enc_time);

    /// send ciphertexts of matrix
    for (auto const& row : uploading) {
        for (auto const& ctx : row) {
            conn << ctx;
            //_ctx_sent++;
        }
    }
	conn.flush();
    clt_ben.ctx_sent = MAX_X1 * MAX_Y1;
	/// we convert DoubleCRT to poly form when send ciphertexts through, and thus, we
	/// count this cost as a part of encryption.
	clt_ben.enc_times.back() += getTimerByName("TO_POLY_OUTPUT")->getTime() * 1000.;

    std::vector<GMMPrecompTable> tbls = precompute_gmm_tables(context);
    /// waiting results
    long rows_of_A = A.NumRows();
    long rows_of_Bt = B.NumCols(); // Bt::Rows = B::Cols
	int64_t ctx_cnt = 0;
	conn >> ctx_cnt;
    clt_ben.ctx_recv = ctx_cnt;
    std::vector<Ctxt> ret_ctxs(ctx_cnt, Ctxt(ek));
	for (size_t k = 0; k < ctx_cnt; k++) {
		conn >> ret_ctxs.at(k);
        //_ctx_received++;
    }
    double eval_time = 0.;
    conn >> eval_time;
    srv_ben.eval_times.push_back(eval_time);
    /// decrypt
    Matrix computed;
    computed.SetDims(A.NumRows(), B.NumCols());
    zero(computed);
    int x = 0;
    int y = 0;
    std::vector<long> slots;
    std::vector<NTL::zz_pX> _slots;
    //NTL::Vec<long> decrypted;
	NTL::ZZX decrypted;
	double decrypt_time = 0.;
    double unpack_time = 0.;
	long ctx_idx = 0;
	bool dec_pass = true;
	for (const auto &ctx : ret_ctxs) {
		double one_dec_time = 0.;
		double one_unpack_time = 0.;
		do {
			AutoTimer timer(&one_dec_time);
			dec_pass &= ctx.isCorrect();
			//faster_decrypt(decrypted, sk, ctx);
			sk.Decrypt(decrypted, ctx);
		} while(0);
		do {
			AutoTimer timer(&one_unpack_time);
            extract_inner_products(slots, decrypted, tbls, context);
		} while(0);
        decrypt_time += one_dec_time;
        unpack_time += one_unpack_time;

		long row_blk = ctx_idx / B.NumCols();
		long column = ctx_idx % B.NumCols();
		ctx_idx += 1;
        fill_compute(computed, row_blk, column, slots, ea);
    }
	/// we convert poly to DoubleCRT when receiving ciphertexts.
	decrypt_time += getTimerByName("FROM_POLY_OUTPUT")->getTime() * 1000.;
    clt_ben.dec_times.push_back(decrypt_time);
    clt_ben.unpack_times.push_back(unpack_time);
	if (!::is_same(ground_truth, computed, NTL::zz_p::modulus()))
		std::cerr << "The computation seems wrong " << std::endl;
	if (!dec_pass)
		std::cerr << "Decryption might fail" << std::endl;
    global_counter++;
}

int run_client(std::string const& addr, long port,
               long n1, long n2, long n3) {
    const long m = 8192;
    const long p = 70913;
    const long r = 1;
    const long L = 2;
    NTL::zz_p::init(p);
    FHEcontext context(m, p, r);
    context.bitsPerLevel = 60;
    buildModChain(context, L);
    FHESecKey sk(context);
    sk.GenSecKey(64);
    auto start_time_stamp = std::clock();
    auto last_time_stamp = std::clock();
    int done = 0;
    while (true) {
        tcp::iostream conn(addr, std::to_string(port));
        if (!conn) {
            std::cerr << "Can not connect to server!" << std::endl;
            return -1;
        }

        /// send FHEcontext obj
        double all_time = 0.;
        do {
            send_context(conn, context);
            AutoTimer time(&all_time);
            // send the evaluation key
            play_client(conn, sk, context, n1, n2, n3);
        } while(0);
        clt_ben.total_times.push_back(all_time);
        conn.close();
		resetAllTimers(); // reset timers in HElib

        ++done;
        auto now_time = std::clock();
        if (now_time - last_time_stamp >= 60 * CLOCKS_PER_SEC) {
            std::cout << done << "\n";
            last_time_stamp = now_time;
        }

        if (now_time - start_time_stamp >= 3600 * CLOCKS_PER_SEC) {
            /// one hour
            break;
        }
    }
    std::cout << "one hour finished: " << done << "\n";
    return 0;
}

int run_server(long port, long n1, long n2, long n3) {
    boost::asio::io_service ios;
    tcp::endpoint endpoint(tcp::v4(), port);
    tcp::acceptor acceptor(ios, endpoint);
    for (long run = 0; run < REPEAT; run++) {
        tcp::iostream conn;
        boost::system::error_code err;
        acceptor.accept(*conn.rdbuf(), err);

        if (!err) {
			SMPServer server;
			server.run(conn, n1, n2, n3);
			resetAllTimers(); // reset timers in HElib
        }
    }
	SMPServer::print_statistics();
    return 0;
}

int main(int argc, char *argv[]) {
    ArgMapping argmap;
    long role = -1;
    long n1 = 8;
    long n2 = 8;
    long n3 = 8;
    std::string addr = "127.0.0.1";
	long port = 12345;
    argmap.arg("N", n1, "n1");
    argmap.arg("M", n2, "n2");
    argmap.arg("D", n3, "n3");
    argmap.arg("R", role, "role. 0 for server and 1 for client");
	argmap.arg("a", addr, "server address");
	argmap.arg("p", port, "port");
    argmap.parse(argc, argv);
    if (role == 0) {
        run_server(port, n1, n2, n3);
    } else if (role == 1) {
        run_client(addr, port, n1, n2, n3);
    } else {
		argmap.usage("General Matrix Multiplication for |N*M| * |M*D|");
		return -1;
	}
}


#pragma once

#include <boost/numeric/odeint.hpp>

#include "eigenIntegration.hpp"

#include "activeModel.hpp"

namespace scpp::discretization
{

template <bool INTERPOLATE_INPUT, bool VARIABLE_TIME>
class ODE
{
private:
    Model::input_vector_t u_t0, u_t1;
    double time;
    double dt;
    Model::ptr_t model;

public:
    using ode_matrix_t = typename Eigen::Matrix<double, Model::state_dim,
                                                1 + Model::state_dim + Model::input_dim +
                                                    INTERPOLATE_INPUT * Model::input_dim +
                                                    VARIABLE_TIME + 1>;

    ODE(const Model::input_vector_t &u_t0,
        const Model::input_vector_t &u_t1,
        const double &time,
        double dt,
        Model::ptr_t model)
        : u_t0(u_t0), u_t1(u_t1), time(time), dt(dt), model(model) {}

    void operator()(const ode_matrix_t &V, ode_matrix_t &dVdt, const double t);
};

template <bool INTERPOLATE_INPUT, bool VARIABLE_TIME>
void ODE<INTERPOLATE_INPUT, VARIABLE_TIME>::operator()(const ode_matrix_t &V, ode_matrix_t &dVdt, const double t)
{
    const Model::state_vector_t x = V.col(0);

    Model::input_vector_t u;
    if constexpr (INTERPOLATE_INPUT)
    {
        u = u_t0 + t / dt * (u_t1 - u_t0);
    }
    else
    {
        u = u_t0;
    }

    Model::state_vector_t f;
    Model::state_matrix_t A;
    Model::control_matrix_t B;
    model->computef(x, u, f);
    model->computeJacobians(x, u, A, B);

    if constexpr (VARIABLE_TIME)
    {
        A *= time;
        B *= time;
    }

    const Model::state_matrix_t Phi_A_xi = V.template block<Model::state_dim, Model::state_dim>(0, 1);
    const Model::state_matrix_t Phi_A_xi_inverse = Phi_A_xi.inverse();

    size_t cols = 0;

    // state
    if constexpr (VARIABLE_TIME)
    {
        dVdt.template block<Model::state_dim, 1>(0, cols) = time * f;
    }
    else
    {
        dVdt.template block<Model::state_dim, 1>(0, cols) = f;
    }
    cols += 1;

    // A
    dVdt.template block<Model::state_dim, Model::state_dim>(0, cols).noalias() = A * Phi_A_xi;
    cols += Model::state_dim;

    if constexpr (INTERPOLATE_INPUT)
    {
        // B
        const double alpha = (dt - t) / dt;
        dVdt.template block<Model::state_dim, Model::input_dim>(0, cols).noalias() = Phi_A_xi_inverse * B * alpha;
        cols += Model::input_dim;

        // C
        const double beta = t / dt;
        dVdt.template block<Model::state_dim, Model::input_dim>(0, cols).noalias() = Phi_A_xi_inverse * B * beta;
        cols += Model::input_dim;
    }
    else
    {
        // B
        dVdt.template block<Model::state_dim, Model::input_dim>(0, cols).noalias() = Phi_A_xi_inverse * B;
        cols += Model::input_dim;
    }

    if constexpr (VARIABLE_TIME)
    {
        // s
        dVdt.template block<Model::state_dim, 1>(0, cols).noalias() = Phi_A_xi_inverse * f;
        cols += 1;
        // z
        dVdt.template block<Model::state_dim, 1>(0, cols).noalias() = Phi_A_xi_inverse * (-A * x - B * u);
        cols += 1;
    }
    else
    {
        // z
        dVdt.template block<Model::state_dim, 1>(0, cols).noalias() = Phi_A_xi_inverse * (f - A * x - B * u);
        cols += 1;
    }

    assert(cols == ode_matrix_t::ColsAtCompileTime);
}

template <bool INTERPOLATE_INPUT, bool VARIABLE_TIME>
void multipleShootingImplementation(
    Model::ptr_t model,
    trajectory_data_t &td,
    discretization_data_t &dd)
{
    const size_t K = td.n_X();

    using ODEFun = ODE<INTERPOLATE_INPUT, VARIABLE_TIME>;
    using ode_matrix_t = typename ODEFun::ode_matrix_t;

    double dt = 1. / double(K - 1);

    if constexpr (not VARIABLE_TIME)
    {
        dt *= td.t;
    }

    using namespace boost::numeric::odeint;
    runge_kutta_fehlberg78<ode_matrix_t, double, ode_matrix_t, double, vector_space_algebra> stepper;

    for (size_t k = 0; k < K - 1; k++)
    {
        ode_matrix_t V;
        V.col(0) = td.X.at(k);
        V.template block<Model::state_dim, Model::state_dim>(0, 1).setIdentity();
        V.template rightCols<ode_matrix_t::ColsAtCompileTime - 1 - Model::state_dim>().setZero();

        const Model::input_vector_t u0 = td.U[k];
        const Model::input_vector_t u1 = INTERPOLATE_INPUT ? td.U[k + 1] : u0;
        ODEFun odeMultipleShooting(u0, u1, td.t, dt, model);

        integrate_adaptive(stepper, odeMultipleShooting, V, 0., dt, dt / 5.);

        size_t cols = 1;

        dd.A[k] = V.template block<Model::state_dim, Model::state_dim>(0, cols);
        cols += Model::state_dim;

        dd.B[k].noalias() = dd.A[k] * V.template block<Model::state_dim, Model::input_dim>(0, cols);
        cols += Model::input_dim;

        if constexpr (INTERPOLATE_INPUT)
        {
            dd.C[k].noalias() = dd.A[k] * V.template block<Model::state_dim, Model::input_dim>(0, cols);
            cols += Model::input_dim;
        }

        if constexpr (VARIABLE_TIME)
        {
            dd.s[k].noalias() = dd.A[k] * V.template block<Model::state_dim, 1>(0, cols);
            cols += 1;
        }

        dd.z[k].noalias() = dd.A[k] * V.template block<Model::state_dim, 1>(0, cols);
        cols += 1;

        assert(cols == ode_matrix_t::ColsAtCompileTime);
    }
}

} // namespace scpp::discretization