#include <iostream>
#include <fstream>
#include <string>
#include <Eigen/Dense>
#include <iomanip>
#include <vector>
#include <algorithm>
#include <cassert>
#include <math.h>
#include <boost/algorithm/string/split.hpp>
#include <boost/algorithm/string/classification.hpp>
#include <experimental/filesystem>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include "opencv2/opencv.hpp"
#include <limits>
#include <chrono>
#include <omp.h>
#include <array>

using namespace std;
using namespace Eigen;
using namespace std::chrono;
namespace fs = std::experimental::filesystem;
namespace ba = boost::algorithm;

string LIDAR_FILE = "";

class projection{
public:
  projection(){
    Eigen::initParallel();
    Eigen::setNbThreads(12);
    MatrixXf data = readbinfile(LIDAR_FILE);

    MatrixXf P;
    P.resize(3, 4);
    P << 609.6954, -721.4216, -1.2513,   -123.0418,
         180.3842,  7.6448,   -719.6515, -101.0167,
         0.9999,    1.2437e-4, 0.0105,   -0.2694;
    auto start = high_resolution_clock::now();

    data = points_filter(P, data);
    cv::Mat result = DenseMap(data, 4);
    auto stop = high_resolution_clock::now();
    auto duration = duration_cast<microseconds>(stop - start);
    cout << duration.count() << " microseconds" << endl;

    double minVal, maxVal;
    cv::Point minLoc, maxLoc;
    cv::minMaxLoc(result, &minVal, &maxVal, &minLoc, &maxLoc);
    result = 255 * (result - minVal) / (maxVal - minVal);
    result.convertTo(result, CV_8UC1);
    cv::imshow("result", result);
    cv::waitKey(0);
  }

  MatrixXf readbinfile(const string dir){

    ifstream fin(dir.c_str(), ios::binary);
    assert(fin);

    fin.seekg(0, ios::end);
    const size_t num_elements = fin.tellg() / sizeof(float);
    fin.seekg(0, ios::beg);

    vector<float> l_data(num_elements);
    fin.read(reinterpret_cast<char*>(&l_data[0]), num_elements*sizeof(float));

    MatrixXf data = Map<MatrixXf>(l_data.data(), 4, l_data.size()/4);

    return data;
  }

  MatrixXf points_filter(MatrixXf &P, MatrixXf &data){
    data = P * data;
    vector<int> v;
    omp_set_num_threads(16);
    #pragma omp parallel
    {
      vector<int> v1;
      #pragma omp for nowait
      for (int j = 0; j < data.cols(); j++){
        if (data(2, j) > 0){
          float x = data(0, j) / data(2, j);
          float y = data(1, j) / data(2, j);
          if ( (x > 0 && x < COL - 0.5) && (y > 0 && y < ROW - 0.5) )
            v1.push_back(j);
        }
      }
      #pragma omp critical
      v.insert(v.end(), v1.begin(), v1.end());
    }
    MatrixXf result;
    result.resize(3, v.size());
    result.fill(0.);
    cout << data.rows() << " " << data.cols() << endl;
    omp_set_num_threads(16);
    #pragma omp parallel
    {
      MatrixXf res_private;
      res_private.resize(3, v.size());
      res_private.fill(0.);
      #pragma omp for nowait
      for (auto i = 0; i < v.size(); i++){
        res_private(0, i) = data(0, v[i]) / data(2, v[i]);
        res_private(1, i) = data(1, v[i]) / data(2, v[i]);
        res_private(2, i) = data(2, v[i]);

      }
      #pragma omp critical
      result += res_private;
    }
    return result;
  }

  cv::Mat DenseMap(MatrixXf &data, int grid){
    int ng = 2 * grid + 1;

    cv::Mat map, mD;
    map = cv::Mat::zeros(ROW, COL, CV_32FC1);
    mD = cv::Mat::zeros(ROW, COL, CV_32FC1);
    omp_set_num_threads(8);
    #pragma omp parallel
    {
      #pragma omp for nowait
      for (auto i = 0; i < data.cols(); i++){
        map.at<float>(round(data(1, i)), round(data(0, i))) =
          sqrt(pow(data(0, i) - round(data(0, i)), 2) +
               pow(data(1, i) - round(data(1, i)), 2));
        mD.at<float>(round(data(1, i)), round(data(0, i))) = data(2, i);
      }
    }

    cv::Mat output;
    output = cv::Mat::zeros(ROW, COL, CV_32FC1);
    omp_set_num_threads(128);
    #pragma omp parallel
    {
      #pragma omp for nowait
      for (auto i = 0; i < ROW; i++){
        for (auto j = 0; j < COL; j++){
          if (i - grid < 0 || i + grid >= ROW)
            continue;
          if (j - grid < 0 || j + grid >= COL)
            continue;
          float s = 0;
          for (auto r = -grid; r < grid + 1; r++){
            for (auto c = -grid; c < grid + 1; c++){
              float map_val = map.at<float>(i+r, j+c);
              if (map_val != 0){
                output.at<float>(i, j) +=
                  mD.at<float>(i+r, j+c) / map.at<float>(i+r, j+c);
                s += 1 / map.at<float>(i+r, j+c);
              }
            }
          }
          if (s == 0){
            s = 1;
          }
          output.at<float>(i, j) /= s;
        }
      }
    }
    return output;
  }

private:
  const int ROW = 375;
  const int COL = 1242;

};

int main(){
  projection p;
}


//-------------------------------------------------------------------
// The code was written by Vikas C. Raykar
// and is copyrighted under the Lessr GPL:
//
// Copyright (C) 2006 Vikas C. Raykar
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as
// published by the Free Software Foundation; version 2.1 or later.
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
// See the GNU Lesser General Public License for more details.
// You should have received a copy of the GNU Lesser General Public
// License along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston,
// MA 02111-1307, USA.
//
// The author may be contacted via email at: vikas(at)cs(.)umd(.)edu
//-------------------------------------------------------------------

//-------------------------------------------------------------
// File    : FastUnivariateDensityDerivative.cpp
// Purpose : Implementation for for FastUnivariateDensityDerivative
// Author  : Vikas C. Raykar (vikas@cs.umd.edu)
// Date    : September 17, 2005
//-------------------------------------------------------------

#include <math.h>
#include <boost/python.hpp>
#define min(a, b) (((a) < (b)) ? (a) : (b))
#define max(a, b) (((a) > (b)) ? (a) : (b))
#define P_UL 500
#define R 1.0
namespace python = boost::python;
//-------------------------------------------------------------------
// Constructor.
//
// PURPOSE
// -------
// Initialize the class.
// Read the parameters.
// Choose the parameter for the algorithm.
// Space subdivision.
// Compute the constant a.
// Compute B or all the clusters.
//
// PARAMETERS
// ----------
// NSources		      --> number of sources, N.
// MTargets		      --> number of targets, M.
// pSources		      --> pointer to sources, px(N).
// pTargets           --> pointer to the targets, py(M).
// Bandwidth		  --> the source bandwidth, h.
// Order              --> order of the derivative, r.
// epsilon            --> desired error, eps.
// pDensityDerivative --> pointer the the evaluated Density
//-------------------------------------------------------------------

class FastUnivariateDensityDerivative {

public:
    //constructor
    FastUnivariateDensityDerivative(
        int     NSources,
        int     MTargets,
        python::list pSources,
        python::list pTargets,
        double  Bandwidth,
        int     Order,
        double  epsilon);

    python::list pDensityDerivative;

    //destructor
    ~FastUnivariateDensityDerivative();

    //function to evaluate the Density Derivative
    void evaluate();

    //function to evaluate the Hermite polynomial.
    double hermite(double x, int r);

private:
    int N; //number of sources.
    int M; //number of targets.
    double* px; //pointer to sources, (N).
    double* py; //pointer to the targets, (M).
    double h; //the source bandwidth.
    int r; //the rth density derivative.
    double eps; //the desired error
    double* pD; //pointer to the evaluated Density Derivative, (M).

    double rx;
    double rr;
    double ry;
    int K;
    int p;
    double h_square;
    double two_h_square;

    double* pClusterCenter;
    int* pClusterIndex;

    int num_of_a_terms;
    double* a_terms;

    int num_of_b_terms;
    double* b_terms;

    double pi;
    double q;

    void* operator new[](size_t s) { return malloc(s); }
    void operator delete[](void* mem) { free(mem); }

    int factorial(int n);
    double* converter_double(python::list lis);
    python::list converter_list(double* lis);
    void choose_parameters();
    void space_sub_division();
    void compute_a();
    void compute_b();
};

FastUnivariateDensityDerivative::FastUnivariateDensityDerivative(
    int NSources,
    int MTargets,
    python::list pSources,
    python::list pTargets,
    double Bandwidth,
    int Order,
    double epsilon)
{
    // Read the arguments.
    N = NSources;
    M = MTargets;
    px = converter_double(pSources);
    h = Bandwidth;
    r = Order;
    py = converter_double(pTargets);
    pD = new double[N];
    eps = epsilon;

    h_square = h * h;
    two_h_square = 2 * h_square;

    pi = 3.14159265358979;
    q = (pow(-1, r)) / (sqrt(2 * pi) * N * (pow(h, (r + 1))));

    // Choose the parameters for the algorithm.
    choose_parameters();

    // Space sub-division
    space_sub_division();

    // Compute the constant a
    compute_a();

    // Compute the constant B
    compute_b();
}

double* FastUnivariateDensityDerivative::converter_double(python::list lis)
{
    int length = len(lis);
    double* temp = new double[length];
    for (int i = 0; i < length; i++)
    {
        temp[i] = python::extract<double>(lis[i]);
    }
    return temp;
}

python::list FastUnivariateDensityDerivative::converter_list(double* lis)
{
    python::list temp;
    for (int i = 0; i < N; i++)
    {
        temp.append(lis[i]);
    }
    return temp;
}
//-------------------------------------------------------------------
// Destructor.
//-------------------------------------------------------------------
FastUnivariateDensityDerivative::~FastUnivariateDensityDerivative()
{
    delete[] a_terms;
    delete[] b_terms;
}

//-------------------------------------------------------------------
// Compute the factorial.
//-------------------------------------------------------------------
int FastUnivariateDensityDerivative::factorial(int n)
{
    int fact = 1;
    for (int i = 1; i <= n; i++) {
        fact = fact * i;
    }
    return fact;
}

//-------------------------------------------------------------------
// Choose the parameters
// 1. rx --> interval length.
// 2. K  --> number of intervals.
// 3. rr --> cutoff radius.
// 4. ry --> cluster cutoff radius.
// 5. p  --> truncation number.
//-------------------------------------------------------------------

void FastUnivariateDensityDerivative::choose_parameters()
{
    // 1. rx --> interval length.
    rx = h / 2;

    // 2. K  --> number of intervals.
    K = (int)ceil(1.0 / rx);
    rx = 1.0 / K;
    double rx_square = rx * rx;

    // 3. rr --> cutoff radius.
    double r_term = sqrt((double)factorial(r));
    rr = min(R, 2 * h * sqrt(log(r_term / eps)));

    // 4. ry --> cluster cutoff radius.
    ry = rx + rr;

    // 5. p  --> truncation number.
    p = 0;
    double error = 1;
    double temp = 1;
    double comp_eps = eps / r_term;
    while ((error > comp_eps) & (p <= P_UL)) {
        p++;
        double b = min(((rx + sqrt((rx_square) + (8 * p * h_square))) / 2), ry);
        double c = rx - b;
        temp = temp * (((rx * b) / h_square) / p);
        error = temp * (exp(-(c * c) / 2 * two_h_square));
    }
    p = p + 1;
}

//-------------------------------------------------------------------
// Space subdivision
//-------------------------------------------------------------------
void FastUnivariateDensityDerivative::space_sub_division()
{

    // 1. Cluster Centers
    pClusterCenter = new double[K];
    for (int i = 0; i < K; i++) {
        pClusterCenter[i] = (i * rx) + (rx / 2);
    }

    //2. Allocate each source to the corresponding interval
    pClusterIndex = new int[N];
    for (int i = 0; i < N; i++) {
        pClusterIndex[i] = min((int)floor(px[i] / rx), K - 1);
    }
}

//-------------------------------------------------------------------
// Compute the contant term a_{lm}.
// l=0...floor(r/2)
// m=0...r-2l
//-------------------------------------------------------------------
void FastUnivariateDensityDerivative::compute_a()
{
    double r_factorial = (double)factorial(r);
    double* l_constant;
    l_constant = new double[((int)floor((double)r / 2)) + 1];
    l_constant[0] = 1;
    for (int l = 1; l <= (int)floor((double)r / 2); l++) {
        l_constant[l] = l_constant[l - 1] * (-1.0 / (2 * l));
    }
    double* m_constant;
    m_constant = new double[r + 1];
    m_constant[0] = 1;
    for (int m = 1; m <= r; m++) {
        m_constant[m] = m_constant[m - 1] * (-1.0 / m);
    }
    num_of_a_terms = 0;
    for (int l = 0; l <= (int)floor((double)r / 2); l++) {
        for (int m = 0; m <= r - (2 * l); m++) {
            num_of_a_terms++;
        }
    }
    a_terms = new double[num_of_a_terms];
    int k = 0;
    for (int l = 0; l <= (int)floor((double)r / 2); l++) {
        for (int m = 0; m <= r - (2 * l); m++) {
            a_terms[k] = (l_constant[l] * m_constant[m] * r_factorial) / ((double)factorial(r - (2 * l) - m));
            k++;
        }
    }

    delete[] l_constant;
    delete[] m_constant;
}

//-------------------------------------------------------------------
// Compute the contant term B^{n}_{km} for all the clusters.
// n=0...K-1
// k=0...p-1
// m=0...r
//-------------------------------------------------------------------
void FastUnivariateDensityDerivative::compute_b()
{
    num_of_b_terms = K * p * (r + 1);
    b_terms = new double[num_of_b_terms];
    double* k_factorial;
    k_factorial = new double[p];
    k_factorial[0] = 1;
    for (int i = 1; i < p; i++) {
        k_factorial[i] = k_factorial[i - 1] / i;
    }
    double* temp3;
    temp3 = new double[p + r];
    for (int n = 0; n < K; n++) {
        for (int k = 0; k < p; k++) {
            for (int m = 0; m < r + 1; m++) {
                b_terms[(n * p * (r + 1)) + ((r + 1) * k) + m] = 0.0;
            }
        }
    }
    for (int i = 0; i < N; i++) {
        int cluster_number = pClusterIndex[i];
        double temp1 = (px[i] - pClusterCenter[cluster_number]) / h;
        double temp2 = exp(-temp1 * temp1 / 2);
        temp3[0] = 1;
        for (int k = 1; k < p + r; k++) {
            temp3[k] = temp3[k - 1] * temp1;
        }
        for (int k = 0; k < p; k++) {
            for (int m = 0; m < r + 1; m++) {
                b_terms[(cluster_number * p * (r + 1)) + ((r + 1) * k) + m] += (temp2 * temp3[k + m]);
            }
        }
    }
    for (int n = 0; n < K; n++) {
        for (int k = 0; k < p; k++) {
            for (int m = 0; m < r + 1; m++) {
                b_terms[(n * p * (r + 1)) + ((r + 1) * k) + m] *= (k_factorial[k] * q);
            }
        }
    }

    delete[] k_factorial;
    delete[] temp3;
}

//-------------------------------------------------------------------
// Actual function to evaluate the Univariate Density Derivative.
//-------------------------------------------------------------------
void FastUnivariateDensityDerivative::evaluate()
{
    double* temp3 = new double[p + r];
    for (int j = 0; j < M; j++) {
        pD[j] = 0.0;
        int target_cluster_number = min((int)floor(py[j] / rx), K - 1);
        double temp1 = py[j] - pClusterCenter[target_cluster_number];
        double dist = abs(temp1);
        while (dist <= ry && target_cluster_number < K && target_cluster_number >= 0) {
            double temp2 = exp(-temp1 * temp1 / two_h_square);
            double temp1h = temp1 / h;
            temp3[0] = 1;
            for (int i = 1; i < p + r; i++) {
                temp3[i] = temp3[i - 1] * temp1h;
            }
            for (int k = 0; k <= p - 1; k++) {
                int dummy = 0;
                for (int l = 0; l <= (int)floor((double)r / 2); l++) {
                    for (int m = 0; m <= r - (2 * l); m++) {
                        pD[j] = pD[j] + (a_terms[dummy] * b_terms[(target_cluster_number * p * (r + 1)) + ((r + 1) * k) + m] * temp2 * temp3[k + r - (2 * l) - m]);
                        dummy++;
                    }
                }
            }
            target_cluster_number++;
            temp1 = py[j] - pClusterCenter[target_cluster_number];
            dist = abs(temp1);
        }
        target_cluster_number = min((int)floor(py[j] / rx), K - 1) - 1;
        if (target_cluster_number >= 0) {
            double temp1 = py[j] - pClusterCenter[target_cluster_number];
            double dist = abs(temp1);
            while (dist <= ry && target_cluster_number < K && target_cluster_number >= 0) {
                double temp2 = exp(-temp1 * temp1 / two_h_square);
                double temp1h = temp1 / h;
                temp3[0] = 1;
                for (int i = 1; i < p + r; i++) {
                    temp3[i] = temp3[i - 1] * temp1h;
                }
                for (int k = 0; k <= p - 1; k++) {
                    int dummy = 0;
                    for (int l = 0; l <= (int)floor((double)r / 2); l++) {
                        for (int m = 0; m <= r - (2 * l); m++) {
                            pD[j] = pD[j] + (a_terms[dummy] * b_terms[(target_cluster_number * p * (r + 1)) + ((r + 1) * k) + m] * temp2 * temp3[k + r - (2 * l) - m]);
                            dummy++;
                        }
                    }
                }
                target_cluster_number--;
                temp1 = py[j] - pClusterCenter[target_cluster_number];
                dist = abs(temp1);
            }
        }
    }
    pDensityDerivative = converter_list(pD);
    delete[] temp3;
}

BOOST_PYTHON_MODULE(fast_deriv)
{
    namespace python = boost::python;
    python::class_<FastUnivariateDensityDerivative>("FastUnivariateDensityDerivative",
        python::init<int,
                     int,
                     python::list,
                     python::list,
                     double,
                     int,
                     double
                    >())
        .def("evaluate", &FastUnivariateDensityDerivative::evaluate)
        .def_readonly("pD", &FastUnivariateDensityDerivative::pDensityDerivative)
    ;

}


// from http://perso.ens-lyon.fr/philippe.theveny/cise.pdf

#include <iostream>
#include <utility>
#include <cmath> // _Decimal64?

#include <boost/tr1/cmath.hpp>
#include <boost/math/special_functions.hpp>
#include <boost/multiprecision/cpp_dec_float.hpp>
#include <boost/multiprecision/cpp_bin_float.hpp>

typedef boost::multiprecision::number<
            boost::multiprecision::backends::cpp_bin_float<
                106,
                boost::multiprecision::backends::digit_base_2,
                void,
                boost::int16_t, -1022, 1023>,
            boost::multiprecision::et_off>
        cpp_bin_float_double_double;

typedef boost::multiprecision::number<
            boost::multiprecision::backends::cpp_bin_float<
                1000,
                boost::multiprecision::backends::digit_base_2,
                void,
                boost::int16_t, -1022, 1023>,
            boost::multiprecision::et_off>
        cpp_bin_float_1000b;

using boost::multiprecision::cpp_dec_float;
using boost::multiprecision::cpp_bin_float_single;
using boost::multiprecision::cpp_bin_float_double;
using boost::multiprecision::cpp_bin_float_double_extended;
using boost::multiprecision::cpp_bin_float_quad;

typedef boost::multiprecision::number<cpp_dec_float<64> > mp_type;

// This is busted: https://svn.boost.org/trac/boost/ticket/11764
//typedef boost::multiprecision::cpp_bin_float_single mp_type;

//typedef boost::multiprecision::cpp_bin_float_double mp_type;
//typedef boost::multiprecision::cpp_bin_float_double_extended mp_type;
//typedef boost::multiprecision::cpp_bin_float_quad mp_type;
//typedef cpp_bin_float_double_double mp_type;
//typedef cpp_bin_float_1000b mp_type;

//typedef _Decimal64 mp_type;

int main (void)
{
    mp_type a = sin((mp_type)1e22);
    mp_type b = log((mp_type)17.1);
    mp_type c = exp((mp_type)0.42);
    mp_type d = 173746*a + 94228*b - 78487*c;
    std::cout << std::setprecision(std::numeric_limits<mp_type>::max_digits10) << d << std::endl;
    return 0;
}


/**
 * Could really represent any coordinate that has X/Y/Z axises.
 *
 * Copyright 2013 Bruce Ide
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include <Eigen/Core>
#include <Eigen/Geometry>

#ifndef _HPP_XYZ_COORDINATE
#define _HPP_XYZ_COORDINATE

namespace fr {

  namespace coordinates {

    class xyz_coordinate {
    protected:
      double x,y,z;

      inline double limit(double val, double lower, double upper)
      {
	assert(upper > lower);
	double retval = val;
	if (val > upper) {
	  retval = upper;
	} else if (val < lower) {
	  retval = lower;
	}
	return retval;
      }

      inline Eigen::Vector3d unit_vec(const Eigen::Vector3d &vec)
      {
	Eigen::Vector3d retval;
	double nrm = vec.norm();
	retval = vec / nrm;
	return retval;
      }

      // Generic interpolate between two points

      Eigen::Vector3d interpolate_3x3(const Eigen::Vector3d &this_vec, const double &time_now, const Eigen::Vector3d &next_vec, const double &time_other, const double &time_between)
      {
	Eigen::Quaternion<double> quat(1.0, 0.0, 0.0, 0.0);
	double this_norm = this_vec.norm();
	double next_norm = next_vec.norm();
	Eigen::Vector3d this_uv = unit_vec(this_vec);
	Eigen::Vector3d next_uv = unit_vec(next_vec);
	double x = this_uv.dot(next_uv);
	double y = limit(x,-1.0, 1.0);
	double theta = acos(y);
	double adjust = (time_between - time_now) / (time_other - time_now);
	double factor = ((next_norm - this_norm) * adjust + this_norm) / this_norm;
	theta = theta * adjust;
	double sto2 = sin(theta / 2.0);
	Eigen::Vector3d ax = next_vec.cross(this_vec);
	Eigen::Vector3d ax_uv = unit_vec(ax);
	double qx, qy, qz, qw;
	qx = ax_uv(0) * sto2;
	qy = ax_uv(1) * sto2;
	qz = ax_uv(2) * sto2;
	qw = cos(theta/2.0);
	quat = Eigen::Quaternion<double>(qw,qx,qy,qz);
	Eigen::Vector3d z = this_vec * factor;
	Eigen::Quaternion<double> z_q(0.0, z(0), z(1), z(2));
	Eigen::Quaternion<double> q1 = z_q * quat;
	Eigen::Quaternion<double> retaq = quat.inverse();
	Eigen::Quaternion<double> q2 = retaq * q1;
	return q2.vec();
      }

    public:
      xyz_coordinate(const double &x, const double &y, const double &z) : x(x), y(y), z(z)
      {
      }

      xyz_coordinate(const xyz_coordinate &copy) : x(copy.get_x()), y(copy.get_y()), z(copy.get_z())
      {
      }

      // Ah ah, ah ha! Yep, I'm planning children classes!
      virtual ~xyz_coordinate()
      {
      }

      virtual double get_x() const { return x; }
      virtual double get_y() const { return y; }
      virtual double get_z() const { return z; }

      virtual Eigen::Vector3d get_xyz() const
      {
	Eigen::Vector3d retval;
	retval << x,y,z;
	return retval;
      }

    };

  }
}

#endif


/* pcmsolver_copyright_start */
/*
 *     PCMSolver, an API for the Polarizable Continuum Model
 *     Copyright (C) 2013-2016 Roberto Di Remigio, Luca Frediani and contributors
 *
 *     This file is part of PCMSolver.
 *
 *     PCMSolver is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU Lesser General Public License as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 *
 *     PCMSolver 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 Lesser General Public License for more details.
 *
 *     You should have received a copy of the GNU Lesser General Public License
 *     along with PCMSolver.  If not, see <http://www.gnu.org/licenses/>.
 *
 *     For information on the complete list of contributors to the
 *     PCMSolver API, see: <http://pcmsolver.readthedocs.io/>
 */
/* pcmsolver_copyright_end */

#include <cmath>

#include "Config.hpp"

#include <Eigen/Core>
#include <Eigen/LU>

#include "cavity/Cavity.hpp"
#include "cavity/Element.hpp"
#include "green/IGreensFunction.hpp"
#include "utils/MathUtils.hpp"

/*! \file SolverImpl.cpp
 *  \brief Functions common to all solvers
 *  \author Roberto Di Remigio
 *  \date 2015
 */

/*! \brief Builds the **anisotropic** IEFPCM matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \param[in] gf_o Green's function outside the cavity
 *  \return the \f$ \mathbf{K} = \mathbf{T}^{-1}\mathbf{R}\mathbf{A} \f$ matrix
 *
 *  This function calculates the PCM matrix. We use the following definitions:
 *  \f[
 *     \begin{align}
 *       \mathbf{T} &=
 *      \left(2\pi\mathbf{I} - \mathbf{D}_\mathrm{e}\mathbf{A}\right)\mathbf{S}_\mathrm{i}
 *      +\mathbf{S}_\mathrm{e}\left(2\pi\mathbf{I} +
 *      \mathbf{A}\mathbf{D}_\mathrm{i}^\dagger\right) \\
 *      \mathbf{R} &=
 *      \left(2\pi\mathbf{A}^{-1} - \mathbf{D}_\mathrm{e}\right) -
 *      \mathbf{S}_\mathrm{e}\mathbf{S}^{-1}_\mathrm{i}\left(2\pi\mathbf{A}^{-1}-\mathbf{D}_\mathrm{i}\right)
 *     \end{align}
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd anisotropicIEFMatrix(const Cavity & cav, const IGreensFunction & gf_i, const IGreensFunction & gf_o)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI, DI and SE, DE on the whole cavity, regardless of symmetry
  TIMER_ON("Computing SI");
  Eigen::MatrixXd SI = gf_i.singleLayer(cav.elements());
  TIMER_OFF("Computing SI");
  TIMER_ON("Computing DI");
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());
  TIMER_OFF("Computing DI");
  TIMER_ON("Computing SE");
  Eigen::MatrixXd SE = gf_o.singleLayer(cav.elements());
  TIMER_OFF("Computing SE");
  TIMER_ON("Computing DE");
  Eigen::MatrixXd DE = gf_o.doubleLayer(cav.elements());
  TIMER_OFF("Computing DE");

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the fullPCMMatrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    TIMER_ON("Symmetry blocking");
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(DE, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SE, cavitySize, dimBlock, nrBlocks);
    TIMER_OFF("Symmetry blocking");
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  TIMER_ON("Assemble T matrix");
  Eigen::MatrixXd T = ((2 * M_PI * Id - DE * a) * SI + SE * (2 * M_PI * Id + a * DI.adjoint().eval()));
  TIMER_OFF("Assemble T matrix");

  TIMER_ON("Assemble R matrix");
  Eigen::MatrixXd R = ((2 * M_PI * Id - DE * a) - SE * SI.ldlt().solve((2 * M_PI * Id - DI * a)));
  TIMER_OFF("Assemble R matrix");

  TIMER_ON("Assemble T^-1R matrix");
  Eigen::MatrixXd fullPCMMatrix = T.partialPivLu().solve(R);
  TIMER_OFF("Assemble T^-1R matrix");

  return fullPCMMatrix;
}

/*! \brief Builds the **isotropic** IEFPCM matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \param[in] epsilon permittivity outside the cavity
 *  \return the \f$ \mathbf{K} = \mathbf{T}^{-1}\mathbf{R}\mathbf{A} \f$ matrix
 *
 *  This function calculates the PCM matrix. We use the following definitions:
 *  \f[
 *     \begin{align}
 *       \mathbf{T} &=
 *      \left(2\pi\frac{\varepsilon+1}{\varepsilon-1}\mathbf{I} - \mathbf{D}_\mathrm{i}\mathbf{A}\right)\mathbf{S}_\mathrm{i} \\
 *      \mathbf{R} &=
 *      \left(2\pi\mathbf{A}^{-1} - \mathbf{D}_\mathrm{i}\right)
 *     \end{align}
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd isotropicIEFMatrix(const Cavity & cav, const IGreensFunction & gf_i, double epsilon)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI and DI on the whole cavity, regardless of symmetry
  TIMER_ON("Computing SI");
  Eigen::MatrixXd SI = gf_i.singleLayer(cav.elements());
  TIMER_OFF("Computing SI");
  TIMER_ON("Computing DI");
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());
  TIMER_OFF("Computing DI");

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the fullPCMMatrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    TIMER_ON("Symmetry blocking");
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SI, cavitySize, dimBlock, nrBlocks);
    TIMER_OFF("Symmetry blocking");
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  // Tq = -Rv -> q = -(T^-1 * R)v = -Kv
  // T = (2 * M_PI * fact * aInv - DI) * a * SI; R = (2 * M_PI * aInv - DI)
  // fullPCMMatrix_ = K = T^-1 * R * a
  // 1. Form T
  double fact = (epsilon + 1.0)/(epsilon - 1.0);
  TIMER_ON("Assemble T matrix");
  Eigen::MatrixXd T = (2 * M_PI * fact * Id - DI * a) * SI;
  TIMER_OFF("Assemble T matrix");

  TIMER_ON("Assemble R matrix");
  Eigen::MatrixXd R = (2 * M_PI * Id - DI * a);
  TIMER_OFF("Assemble R matrix");

  TIMER_ON("Assemble T^-1R matrix");
  Eigen::MatrixXd fullPCMMatrix = T.partialPivLu().solve(R);
  TIMER_OFF("Assemble T^-1R matrix");

  return fullPCMMatrix;
}

/*! \brief Builds the **anisotropic** \f$ \mathbf{T}_\varepsilon \f$ matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \param[in] gf_o Green's function outside the cavity
 *  \return the \f$ \mathbf{T}_\varepsilon \f$ matrix
 *
 *  We use the following definition:
 *  \f[
 *      \mathbf{T}_\varepsilon =
 *      \left(2\pi\mathbf{I} - \mathbf{D}_\mathrm{e}\mathbf{A}\right)\mathbf{S}_\mathrm{i}
 *      +\mathbf{S}_\mathrm{e}\left(2\pi\mathbf{I} +
 *      \mathbf{A}\mathbf{D}_\mathrm{i}^\dagger\right)
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd anisotropicTEpsilon(const Cavity & cav, const IGreensFunction & gf_i, const IGreensFunction & gf_o)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI, DI and SE, DE on the whole cavity, regardless of symmetry
  Eigen::MatrixXd SI = gf_i.singleLayer(cav.elements());
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());
  Eigen::MatrixXd SE = gf_o.singleLayer(cav.elements());
  Eigen::MatrixXd DE = gf_o.doubleLayer(cav.elements());

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the matrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(DE, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SE, cavitySize, dimBlock, nrBlocks);
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  // Form T
  return ((2 * M_PI * Id - DE * a) * SI + SE * (2 * M_PI * Id + a * DI.adjoint().eval()));
}

/*! \brief Builds the **isotropic** \f$ \mathbf{T}_\varepsilon \f$ matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \param[in] epsilon permittivity outside the cavity
 *  \return the \f$ \mathbf{T}_\varepsilon \f$ matrix
 *
 *  We use the following definition:
 *  \f[
 *      \mathbf{T}_\varepsilon =
 *      \left(2\pi\frac{\varepsilon+1}{\varepsilon-1}\mathbf{I} - \mathbf{D}_\mathrm{i}\mathbf{A}\right)\mathbf{S}_\mathrm{i}
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd isotropicTEpsilon(const Cavity & cav, const IGreensFunction & gf_i, double epsilon)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI, DI and SE, DE on the whole cavity, regardless of symmetry
  Eigen::MatrixXd SI = gf_i.singleLayer(cav.elements());
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the matrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SI, cavitySize, dimBlock, nrBlocks);
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  double fact = (epsilon + 1.0)/(epsilon - 1.0);
  return (2 * M_PI * fact * Id - DI * a) * SI;
}

/*! \brief Builds the **anisotropic** \f$ \mathbf{R}_\infty \f$ matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \param[in] gf_o Green's function outside the cavity
 *  \return the \f$ \mathbf{R}_\infty\mathbf{A} \f$ matrix
 *
 *  We use the following definition:
 *  \f[
 *      \mathbf{R}_\infty =
 *      \left(2\pi\mathbf{A}^{-1} - \mathbf{D}_\mathrm{e}\right) -
 *      \mathbf{S}_\mathrm{e}\mathbf{S}^{-1}_\mathrm{i}\left(2\pi\mathbf{A}^{-1}-\mathbf{D}_\mathrm{i}\right)
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd anisotropicRinfinity(const Cavity & cav, const IGreensFunction & gf_i, const IGreensFunction & gf_o)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI, DI and SE, DE on the whole cavity, regardless of symmetry
  Eigen::MatrixXd SI = gf_i.singleLayer(cav.elements());
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());
  Eigen::MatrixXd SE = gf_o.singleLayer(cav.elements());
  Eigen::MatrixXd DE = gf_o.doubleLayer(cav.elements());

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the matrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SI, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(DE, cavitySize, dimBlock, nrBlocks);
    symmetryBlocking(SE, cavitySize, dimBlock, nrBlocks);
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  // Form R
  return ((2 * M_PI * Id - DE * a) - SE * SI.ldlt().solve((2 * M_PI * Id - DI * a)));
}

/*! \brief Builds the **isotropic** \f$ \mathbf{R}_\infty \f$ matrix
 *  \param[in] cav the discretized cavity
 *  \param[in] gf_i Green's function inside the cavity
 *  \return the \f$ \mathbf{R}_\infty\mathbf{A} \f$ matrix
 *
 *  We use the following definition:
 *  \f[
 *      \mathbf{R}_\infty =
 *      \left(2\pi\mathbf{A}^{-1} - \mathbf{D}_\mathrm{i}\right)
 *  \f]
 *  The matrix is not symmetrized and is not symmetry packed.
 */
inline Eigen::MatrixXd isotropicRinfinity(const Cavity & cav, const IGreensFunction & gf_i)
{
  // The total size of the cavity
  PCMSolverIndex cavitySize = cav.size();
  // The number of irreps in the group
  int nrBlocks = cav.pointGroup().nrIrrep();
  // The size of the irreducible portion of the cavity
  int dimBlock = cav.irreducible_size();

  // Compute SI, DI and SE, DE on the whole cavity, regardless of symmetry
  Eigen::MatrixXd DI = gf_i.doubleLayer(cav.elements());

  // Perform symmetry blocking
  // If the group is C1 avoid symmetry blocking, we will just pack the matrix
  // into "block diagonal" when all other manipulations are done.
  if (cav.pointGroup().nrGenerators() != 0) {
    symmetryBlocking(DI, cavitySize, dimBlock, nrBlocks);
  }

  Eigen::MatrixXd a = cav.elementArea().asDiagonal();
  Eigen::MatrixXd Id = Eigen::MatrixXd::Identity(cavitySize, cavitySize);

  return (2 * M_PI * Id - DI * a);
}


/* Copyright (C) 2012-2017 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */

/**
 * @file PIR.cpp
 * @brief A simple PIR implementation
 */
#include <cstdint>
#include <NTL/GF2X.h>
#include <NTL/BasicThreadPool.h>
#include "intraSlot.h"
#include "matmul.h"

#if (defined(__unix__) || defined(__unix) || defined(unix))
#include <sys/time.h>
#include <sys/resource.h>
#endif

NTL_CLIENT

// Simplistic un-optimized encoding/decoding routines
static void encodeBits(GF2X& poly,
                       const vector<uint8_t>& data, long idx, long d)
{
  FHE_TIMER_START;
  for (long i=0; i<d; i++, idx++) {
    // compute byte index from bit index
    long byteIdx = idx / 8;
    if (byteIdx < lsize(data)) {
      long coef = (data[byteIdx] >> (idx % 8))&1;
      NTL::SetCoeff(poly, i, coef);
    }
  }
}
static void decodeBits(vector<uint8_t>& data, long idx,
                       const GF2X& poly, long d)
{
  FHE_TIMER_START;
  for (long i=0; i<d; i++, idx++) {
    // compute byte index from bit index
    long byteIdx = idx / 8;
    if (byteIdx < lsize(data)) {
      uint8_t coef = NTL::conv<uint>(NTL::coeff(poly,i)) & 1;
      data[byteIdx] |= (coef << (idx % 8));
    }
  }
}

static void decodeDBentry(vector<uint8_t>& entry,
                          const vector<GF2X>& encoded, long d)
{
  long dbSize = divc(lsize(encoded)*d, 8); // size in bytes

}

// Encoding a "database" in a matrix, the number of entries in the
// database is upto the size of the 1st dimension, i.e. sizeOfDim(0).
// The number of bits per entry is upto phi(m)/sizeOfDim(0).
class DBMatrix : public MatMul1D_derived<PA_GF2> {
public:
  PA_INJECT(PA_GF2)

private:
  vector< vector< vector< RX > > > data;
  const EncryptedArray& ea;

public:
  virtual ~DBMatrix() {}
  DBMatrix(const EncryptedArray& _ea, const vector< vector<uint8_t> >&db):
    ea(_ea)
  {
    RBak bak; bak.save(); ea.getAlMod().restoreContext();
    long n = ea.size();
    long d = ea.getDegree();
    long D = ea.sizeOfDimension(0);

    // vector< vector<uint8_t> > db2(lsize(db)); // sanity check
    // for (long i=0; i<lsize(db); i++)
    //   db2[i].resize(db[i].size(), 0);

    // Allocate space for the matrices
    data.resize(n/D);
    for (long k=0; k < n/D; k++) {
      data[k].resize(D);
      for (long i=0; i<D; i++)
	data[k][i].resize(D, RX()); // initialize to D empty polynomials
    }

    /* Each database entry is encoded in upto n/D partial rows, each partial
     * row has D cells of d bits each. In the picture here D=3, n/D=2.
     *
     *  (X1 X2 X3 | X4 X5 X6)  // database entry X
     *  (Y1 Y2 Y3 | Y4 Y5 Y6)  // database entry Y
     *  (Z1 Z2 Z3 | Z4 Z5 Z6)  // database entry Z
     *
     *  For example, the database entry Z is encoded in the two row vectors
     *  (Z1 Z2 Z3), (Z4 Z5 Z6). Each Zi is a degree-d binary polynomial.
     */

    for (long k=0; k<n/D; k++) { // go over row vectors

      // Encode the next d*D bits of each entry db[i]
      for (long i=0; i<D; i++) {
	for (long j=0; j<D; j++) { // encode next d bit of db[i]
          long dataIdx= (k*D+j)*d; // how many bits were already encoded
          if (dataIdx < 8*lsize(db[i])) { // more bits to encode
            encodeBits(data[k][i][j], db[i], dataIdx, d);
            // decodeBits(db2[i], dataIdx, data[k][i][j], d);
          }
	}
      }
    }
    // for (long i=0; i<lsize(db); i++)
    //   for (long j=0; j<lsize(db[i]); j++)
    //     assert(db[i][j] == db2[i][j]);
  }

  const EncryptedArray& getEA() const override { return ea; }
  bool multipleTransforms() const override { return true; }
  long getDim() const override { return 0; }

  bool get(RX& out, long i, long j, long k) const override {
    long n = ea.size();
    long D = ea.sizeOfDimension(0);

    assert(i >= 0 && i < D);
    assert(j >= 0 && j < D);
    assert(k >= 0 && k < n/D);
    if (IsZero(data[k][i][j])) return true;
    out = data[k][i][j];
    return false;
  }
};

static MatMul1D*
buildDBMatrix(const EncryptedArray& ea, const vector< vector<uint8_t> >&db)
{
  FHE_TIMER_START;
  assert (ea.getTag()==PA_GF2_tag);
  return new DBMatrix(ea, db);
}

static vector< vector<uint8_t> >*
buildRandomDB(const EncryptedArray& ea)
{
  FHE_TIMER_START;
  assert (ea.getTag()==PA_GF2_tag);
  vector< vector<uint8_t> >* v = new(vector< vector<uint8_t> >);
  v->resize(ea.sizeOfDimension(0));
  long entrySize = ea.size()/(ea.sizeOfDimension(0)*8);
  for (long i=0; i<ea.sizeOfDimension(0); i++) {
    (*v)[i].resize(entrySize);
    for (long j=0; j<entrySize; j++)
      (*v)[i][j] = NTL::RandomBits_long(8);
  }
  return v;
}

bool DoTest(const EncryptedArray& ea,
            const SecKey& secretKey, bool minimal, bool verbose)
{
  // choose a random database that fits in one ciphertext
  std::unique_ptr<vector< vector<uint8_t> > > db(buildRandomDB(ea));

  // Encode the database as a matrix
  FHE_NTIMER_START(PIR_EncodeDBasMatrix);
  std::unique_ptr< MatMul1D > mat(buildDBMatrix(ea, *db));
  MatMul1D::ExecType mat_exec(*mat, minimal);
  mat_exec.upgrade();
  FHE_NTIMER_STOP(PIR_EncodeDBasMatrix);

  // Choose a random index into the database
  long idx = NTL::RandomBnd(lsize(*db));

  // Encrypt index as a unit vector
  FHE_NTIMER_START(PIR_EncryptSelectionAsVector);
  Ctxt ctxt(secretKey);
  {vector<long> slots(ea.size(), 0);
  for (long i=0; i<ea.size(); i++) {
    if (ea.coordinate(0,i) == idx)
      slots[i]=1;
  }
  ea.encrypt(ctxt, secretKey, slots);
  }
  FHE_NTIMER_STOP(PIR_EncryptSelectionAsVector);

  // Do the matrix-vector multiply
  mat_exec.mul(ctxt);

  // Decrypt and check the result
  FHE_NTIMER_START(PIR_Decrypt_Result);

  vector<ZZX> slots;
  ea.decrypt(ctxt, secretKey, slots); // decrypt the ciphertext vector

  // maximum size of an entry in bits
  long d = ea.getDegree();
  long nOverD = ea.size()/ea.sizeOfDimension(0);
  long entrySize = (ea.size()*d) / ea.sizeOfDimension(0);
  vector<uint8_t> dbEntry(divc(entrySize,8), 0);

  // Go over the slots and put each one in its right place
  for (long i=0; i<ea.size(); i++) {
    // represent i as (ii,jj) with ii index along dim 0 and jj the rest
    std::pair<long,long> pp = ea.getAlMod().getZMStar().breakIndexByDim(i,0);
    long idxx = pp.first*ea.sizeOfDimension(0) + pp.second;
    decodeBits(dbEntry, idxx*d, conv<GF2X>(slots[i]), d);
  }
  FHE_NTIMER_STOP(PIR_Decrypt_Result);

  // check that we've got the right answer
  for (long i=0; i<std::min(lsize(dbEntry), lsize((*db)[idx])); i++)
    if (dbEntry[i] != (*db)[idx][i]) {
      cout << "Grrr@*, entry["<<i<<"]="<<((int)dbEntry[i])<<"!=db["<<idx<<"]["
           <<i<<"]="<<((int)(*db)[idx][i])<<endl;
      return false;
    }
  return true;
}

int ks_strategy = 0;
// 0 == default
// 1 == full
// 2 == BSGS
// 3 == minimal

void  TestIt(Context& context, bool verbose)
{
  if (verbose) {
    context.zMStar.printout();
    std::cout <<"  #threads="<<NTL::AvailableThreads()
              <<", security=" << context.securityLevel()<<endl<< endl;
  }

  SecKey secretKey(context);
  const PubKey& publicKey = secretKey;
  secretKey.GenSecKey(/*w=*/64); // A Hamming-weight-w secret key

  bool minimal = (ks_strategy == 3);

  // we call addSomeFrbMatrices for all strategies except minimal

  switch (ks_strategy) {
  case 0:
    addSome1DMatrices(secretKey);
    addSomeFrbMatrices(secretKey);
    break;
  case 1:
    add1DMatrices(secretKey);
    addSomeFrbMatrices(secretKey);
    break;
  case 2:
    addBSGS1DMatrices(secretKey);
    addSomeFrbMatrices(secretKey);
    break;
  case 3:
    addMinimal1DMatrices(secretKey);
    addMinimalFrbMatrices(secretKey);
    break;

   default:
     Error("bad ks_strategy");
   }
  EncryptedArray ea(context, context.alMod);
  bool okSoFar=true;

  cout << " * "<<ea.sizeOfDimension(0)<<"-entry DB, |entry|="
       << (ea.size()*ea.getDegree()) << " bits. ";
  for (long i=0; i<5; i++)
    if (!DoTest(ea, secretKey, minimal, verbose)) {
      okSoFar = false;
      break;
    }
  if (okSoFar)
    cout << "Nice!!\n\n";

  if (verbose) {
    printAllTimers(cout);
#if (defined(__unix__) || defined(__unix) || defined(unix))
      struct rusage rusage;
      getrusage( RUSAGE_SELF, &rusage );
      cout << "  rusage.ru_maxrss="<<rusage.ru_maxrss << endl;
#endif
  }

}

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

  long m=2047;
  amap.arg("m", m, "defines the cyclotomic polynomial Phi_m(X)");
  long p=2;
  long L=3;
  amap.arg("L", L, "# of levels in the modulus chain");
  long verbose=0;
  amap.arg("verbose", verbose, "print timing and other info");
  long nt=1;
  amap.arg("nt", nt, "# threads");

  amap.arg("force_bsgs", fhe_test_force_bsgs,
           "1 to force on, -1 to force off");
  amap.arg("force_hoist", fhe_test_force_hoist,
           "-1 to force off");
  amap.arg("ks_strategy", ks_strategy,
           "0: default, 1:full, 2:bsgs, 3:minimal");

  NTL::Vec<long> gens;
  amap.arg("gens", gens, "use specified vector of generators", NULL);
  amap.note("e.g., gens='[562 1871 751]'");
  NTL::Vec<long> ords;
  amap.arg("ords", ords, "use specified vector of orders", NULL);
  amap.note("e.g., ords='[4 2 -4]', negative means 'bad'");

  amap.parse(argc, argv);

  if (verbose) {
    cout << "*** matmul1D: m=" << m
	 << ", p^r=2"
	 << ", L=" << L
	 << ", nt=" << nt
	 << ", force_bsgs=" << fhe_test_force_bsgs
	 << ", force_hoist=" << fhe_test_force_hoist
	 << ", ks_strategy=" << ks_strategy
	 << endl;
   }

  vector<long> gens1, ords1;
  convert(gens1, gens);
  convert(ords1, ords);

  if (nt > 1) SetNumThreads(nt);

  setTimersOn();

  Context context(m, 2, 1, gens1, ords1);
  buildModChain(context, L, /*c=*/3);

  TestIt(context, verbose);
}


#pragma once

#include <cassert>
#include <vector>

#include <Eigen/Core>

#include <FMath/detail/Expression_Function.hpp>
#include <FMath/detail/Expression_Lambda.hpp>
#include <FMath/detail/Expression_Slice.hpp>
#include <FMath/detail/Expression_SubSet.hpp>
#include <FMath/detail/Using.hpp>

namespace FMath::detail
{
    template<typename T, typename Container = std::vector<T>>
    class Field
    {
        Container _container;

        // Assignment function which explicitly evaluates an expression
        template<typename Container2>
        void assign(const Container2 & container_from);
        void assign(const T & value);

      public:
        using value_type = T;

        ///////////// Constructors //////////////////////////////////////////////////////

        // Field with initial size
        Field(const std::size_t n) : _container(n) {}

        // Field with initial size and value
        Field(const std::size_t n, const T & value) : _container(n, value) {}

        // Field via initializer list
        Field(const std::initializer_list<T> & list) : _container(list) {}

        // Constructor for underlying container
        Field(const Container & other) : _container(other) {}

        // A Field can be constructed such as to force its evaluation.
        template<typename T2, typename R2>
        Field(const Field<T2, R2> & other) : _container(other.size())
        {
            static_assert(
                std::is_same_v<T, T2>, "FMATH USAGE ERROR: Field<> template parameter "
                                       "must be identical for assignment to work");
            if constexpr (std::is_same_v<T, T2>)
            {
                this->assign(other);
            }
        }

        // Assignment operator for Field of same type (copy assignment)
        Field & operator=(const Field & other)
        {
            assert(size() == other.size());
            this->assign(other);
            return *this;
        }

        ///////////// Assignment ////////////////////////////////////////////////////////

        // Assignment operator for Field of different type
        template<typename T2, typename R2>
        Field & operator=(const Field<T2, R2> & other)
        {
            static_assert(
                std::is_same_v<T, T2>, "FMATH USAGE ERROR: Field<> template parameter "
                                       "must be identical for assignment to work");
            if constexpr (std::is_same_v<T, T2>)
            {
                assert(size() == other.size());
                this->assign(other);
                return *this;
            }
        }

        // Assignment operator for entity
        Field & operator=(const T & value)
        {
            this->assign(value);
            return *this;
        }

        ///////////// Basics ////////////////////////////////////////////////////////////

        void resize(const std::size_t size)
        {
            _container.resize(size);
        }
        void resize(const std::size_t size, const T & value)
        {
            _container.resize(size, value);
        }

        // If the container is a std::vector, data can be retrieved as a pointer
        T * data()
        {
            static_assert(
                std::is_same_v<Container, std::vector<T>>,
                "FMATH USAGE ERROR: data() is only available for evaluated "
                "Fields, not Expressions");
            if constexpr (std::is_same_v<Container, std::vector<T>>)
            {
                return _container.data();
            }
        }

        // Size of underlying container
        auto size() const
        {
            return _container.size();
        }

        // Index operators
        T operator[](const std::size_t i) const
        {
            return _container[i];
        }
        T & operator[](const std::size_t i)
        {
            return _container[i];
        }

        // Returns the underlying data
        const Container & contents() const
        {
            return _container;
        }
        Container & contents()
        {
            return _container;
        }

        ///////////// Transformation ////////////////////////////////////////////////////

        // Re-interpretation as a reference to an Eigen::VectorX
        template<typename RefT>
        Eigen::Ref<RefT> asRef()
        {
            // Field<Vector3> to VectorX of size 3*N
            if constexpr (std::is_same_v<T, Vector3>)
                return Eigen::Ref<RefT>(
                    Eigen::Map<RefT>(this->_container[0].data(), 3 * this->size()));
            // Field<scalar> etc. to VectorX of size N
            else
                return Eigen::Ref<RefT>(Eigen::Map<RefT>(this->data(), this->size()));
        }

        ///////////// Lambda ////////////////////////////////////////////////////////////

        // Applies a given lambda to every entry of the Field.
        // The lambda is passed the index and the entry corresponding to the index.
        // This function is applied immediately, not on assignment.
        template<typename Lambda>
        void apply_lambda(const Lambda & lambda)
        {
            static_assert(
                std::is_convertible<
                    Lambda, std::function<void(std::size_t, const T &)>>::value,
                "FMATH USAGE ERROR: you cannot use apply_lambda with a type that "
                "is not convertible to std::function<void(std::size_t, const T &)>.");
            if constexpr (
                std::is_convertible<
                    Lambda, std::function<void(std::size_t, const T &)>>::value)
            {
                this->assign(Field<T, FieldLambda<T, Container, Lambda>>(
                    FieldLambda<T, Container, Lambda>(this->contents(), lambda)));
            }
        }

        // Applies a given lambda to every entry of the Field.
        // The lambda is passed the index and the entry corresponding to the index.
        template<typename Lambda>
        auto applied_lambda(const Lambda & lambda)
        {
            static_assert(
                std::is_convertible<
                    Lambda, std::function<void(std::size_t, const T &)>>::value,
                "FMATH USAGE ERROR: you cannot use applied_lambda with a type that "
                "is not convertible to std::function<void(std::size_t, const T &)>.");
            if constexpr (
                std::is_convertible<
                    Lambda, std::function<void(std::size_t, const T &)>>::value)
            {
                return Field<T, FieldLambda<T, Container, Lambda>>(
                    FieldLambda<T, Container, Lambda>(this->contents(), lambda));
            }
        }

        ///////////// Reductions ////////////////////////////////////////////////////////

        // Returns the sum over all entries of the Field
        T sum() const;

        // Returns the average over all entries of the Field
        T mean() const;

        // This is only valid for scalar contents
        // This will return the minimum value.
        scalar min() const;

        // This is only valid for scalar contents
        // This will return the maximum value.
        scalar max() const;

        // This is only valid for scalar contents
        // This will return the minimum and maximum value.
        std::pair<scalar, scalar> minmax() const;

        // This is only valid for Vector3 contents
        // Returns the minium and maximum value of the components of all vectorfield
        // entries
        scalar min_component() const;

        // This is only valid for Vector3 contents
        // Returns the minium and maximum value of the components of all vectorfield
        // entries
        scalar max_component() const;

        // This is only valid for Vector3 contents
        // Returns the minium and maximum value of the components of all vectorfield
        // entries
        std::pair<scalar, scalar> minmax_component() const;

        ///////////// VectorField Operations on self ////////////////////////////////////

        // For a VectorField, this returns a Field of the Vector3 norms
        auto norm() const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: norm() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<scalar, NormEx<T, Container>>(
                    NormEx<T, Container>(this->contents()));
            }
        }

        // For a VectorField, this returns a Field of the squared Vector3 norms
        auto squaredNorm() const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: squaredNorm() is only available on "
                "Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<scalar, SquaredNormEx<T, Container>>(
                    SquaredNormEx<T, Container>(this->contents()));
            }
        }

        // Normalizes the Vector3 entries of a VectorField to norm 1.
        // If a norm is zero, nothing is done.
        // This function is applied immediately, not on assignment.
        void normalize()
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: normalize() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                this->assign(Field<T, NormalizedEx<T, Container>>(
                    NormalizedEx<T, Container>(this->contents())));
            }
        }

        // Normalizes the Vector3 entries of a VectorField to norm 1.
        // If a norm is zero, nothing is done.
        auto normalized() const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: normalized() is only available on "
                "Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<T, NormalizedEx<T, Container>>(
                    NormalizedEx<T, Container>(this->contents()));
            }
        }

        ///////////// VectorField Operations with others ////////////////////////////////

        // Element-wise dot-product between vector-fields, yielding a scalar-field
        template<typename Container2>
        auto dot(const Field<Vector3, Container2> & field) const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: dot() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<scalar, FieldDotFieldEx<T, Container, Container2>>(
                    FieldDotFieldEx<T, Container, Container2>(
                        this->contents(), field.contents()));
            }
        }

        // Element-wise dot-product between a vector-field and a vector, yielding a scalar-field
        auto dot(const Vector3 & vec) const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: dot() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<scalar, VectorDotFieldEx<T, Container>>(
                    VectorDotFieldEx<T, Container>(this->contents(), vec));
            }
        }

        // Element-wise cross-product between vector-fields, yielding a vector-field
        template<typename Container2>
        auto cross(const Field<Vector3, Container2> & field) const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: cross() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<Vector3, FieldCrossFieldEx<T, Container, Container2>>(
                    FieldCrossFieldEx<T, Container, Container2>(
                        this->contents(), field.contents()));
            }
        }

        // Element-wise cross-product between a vector-field and a vector, yielding a vector-field
        auto cross(const Vector3 & vec) const
        {
            static_assert(
                std::is_same_v<T, Vector3>,
                "FMATH USAGE ERROR: cross() is only available on Field<Vector3>");
            if constexpr (std::is_same_v<T, Vector3>)
            {
                return Field<Vector3, VectorCrossFieldEx<T, Container>>(
                    VectorCrossFieldEx<T, Container>(this->contents(), vec));
            }
        }

        ///////////// SubSet Extraction /////////////////////////////////////////////////

        // Extract a 1D slice of a Field
        auto slice(
            const std::size_t begin = 0, const std::optional<std::size_t> end = {},
            const std::size_t stride = 1)
        {
            return Field<T, SliceEx<T, Container>>(
                SliceEx<T, Container>(this->contents(), begin, end, stride));
        }

        // Extract a subset of a Field's values via a list of indices
        auto operator[](const std::vector<std::size_t> & indices)
        {
            return Field<T, SubSetEx<T, Container>>(
                SubSetEx<T, Container>(this->contents(), indices));
        }
    };
}


#include <algorithm>
using namespace std;

#include "regression_tree.h"
#include <boost/bind.hpp>
#include <boost/thread/thread.hpp>
using namespace boost;

extern vector<vector<int> > fInds;

int feature_to_split_on;

/*bool mysortpred(const tuple* d1, const tuple* d2) {
  return d1->features[feature_to_split_on] < d2->features[feature_to_split_on];
}
*/
/*void sort_data_by_feature(vector<tuple*>& data, int f) {
  feature_to_split_on = f;
  sort(data.begin(), data.end(), mysortpred);
}
*/

void sort_data_by_feature(vector<int>& location,  vector<int> dataCount, vector<int> invertIdx, int f){
	int cur=0;
        for (int i = 0; i < fInds[f].size(); i ++){
                int z = fInds[f][i];
		int loc=invertIdx[z];
                for (int j = 0; j < dataCount[z]; j ++)
			location[cur++]=loc;
        }
}

/*bool mysortpred2(const pair<tuple*, int> tk1, const pair<tuple*, int> tk2) {
  return tk1.first->features[tk1.second] < tk2.first->features[tk2.second];
}*/
/////////

//double best_fc_in_feature(vector<tuple*> data, int f, pair<int, double>& fc) {
double best_fc_in_feature(vector<tuple*> data, vector<int> dataCount, vector<int> invertIdx, int f, int& fs, double& vs, const args_t& args) {  //if (!XX)
/*  vector< pair<tuple*, int> > tk;
  int z;

  for (z = 0; z < data.size(); z++)
    tk.push_back( pair<tuple*,int>(data[z], f) );
  sort(tk.begin(), tk.end(), mysortpred2);
  for (z = 0; z < data.size(); z++)
    data[z] = tk[z].first;
*/
  int n = data.size(), i;
  double min = MY_DBL_MAX;
        vector<int> location(n, -1);
        sort_data_by_feature(location, dataCount, invertIdx, f);

  // get impurity (squared loss) for missing data
  double M = 0.0;
  double W = 0.0;
  int missing = 0;
	int loc = location[missing];
  while (data[loc]->features[f] == UNKNOWN && missing < n-1) {
    //M += data[missing]->target * 1.0;
	M += data[loc]->target * data[loc]->weight;
	W += data[loc]->weight;
    	missing++;
	loc = location[missing];
  }
  if (missing == n-1) // all data is missing
    return MY_DBL_MAX;
  if (missing) {
    //double mbar = M * 1.0 / missing;
	double mbar = M/W;
    M = 0.0;
    for (i = 0; i < missing; i++){
	loc = location[i];
      M += data[loc]->weight * (data[loc]->target - mbar) * (data[loc]->target - mbar);
    }
  }
  int nn = n - missing; // number of data points that arent missing

  // we have impurity = E_{i=0..n} (yi - ybar)^2  (E = summation)
  // this factors to impurty = E(yi^2) + (n+1)ybar^2 - 2*(n+1)*ybar^2
  // we want the impurity for left and right splits for all splits k
  // let ybl = ybar left, ybr = ybar right
  //     s = E_{i=0..k} yi^2
  //     r = E_{i=k+1..n} yi^2

  //impurity =\sum_i  wi (yi - ybar)^2 = \sum_i wi*yi^2 - 2*(\sum_i wi*yi)*ybar + (\sum_i wi) ybar^2
  // ybar = (\sum_i wi*yi) / (sum_i wi) = ywr/WR
  //  r = \sum_i wi*yi^2, ywr = \sum_i wi*yi,  WR = \sum_i wi
  // impurity = r - 2*ywr*ywr/WR + ywr*ywr/WR = r - ywr*ywr/WR
  double ybl, ybr, s, r, L, R, I, ywl, ywr, WL, WR;
  ybl = ybr = s = r = ywl = ywr = WL = WR = 0.0;

  // put all data into right side R
  // get ybar right and r
  int start = missing;
  for (i = start; i < n; i++) {
	loc = location[i];
    	r += data[loc]->target * data[loc]->target * data[loc]->weight;
    	//ybr += data[i]->target;
    	ywr += data[loc]->target * data[loc]->weight;
    	WR += data[loc]->weight;
  }
  //ybr /= 1.0 * nn;
  //r += 0.0000000001 for precision errors

  // for every i
  // put yi into left side, remove it from the right side, and calculate squared lost
  // impurity of putting all data points into the right tree equals putting all the data into the left tree (such cases are not considered)
  for (i = start; i < n-1; i++) {
    int j = i - missing;
    double yn = data[location[i]]->target;
    double w = data[location[i]]->weight;

    s += w *  yn * yn;
    r -= w * yn * yn;
    ywr -= w * yn;
    ywl += w * yn;
    WL += w;
    WR -= w;

    if (r < 0 && r > -0.000001) r = 0;
    if (r < 0) r = 0;

    //ybl = (j*ybl + yn) / (j+1.0);
    //ybr = ((nn-j)*ybr - yn) / (nn-j-1.0);

    //L = s + WL*ybl*ybl - 2*ybl*ywl;
    //R = r + WR*ybr*ybr - 2*ybr*ywr;
    L = s - ywl*ywl/WL;
    R = r - ywr*ywr/WR;

    //L = s + (j+1)*ybl*ybl - 2*(j+1)*ybl*ybl;
    //R = r + (nn-j-1.0)*ybr*ybr - 2*(nn-j-1.0)*ybr*ybr;

    // precision errors?
    if (L < 0 && L > -0.0000001) L = 0;
    if (R < 0 && R > -0.0000001) R = 0;
    if (R < 0) R = 0;
    if (L < 0) L = 0;

    if(0)
    if (L < 0 || R < 0 || r < 0)
      printf("Problem %lf %lf %lf\n", L, R, r);

    // do not consider splitting here if data is the same as next
    if (data[location[i]]->features[f] == data[location[i+1]]->features[f])
      continue;

    I = L + R;// + M;
    I = L + R + M;

    if (I < min) {
      min = I;
      fs = f;
      vs = (data[location[i]]->features[f] + data[location[i+1]]->features[f]) / 2;
    }
  }

  return min;
}

// find best feature to split on in range [start, end)
// store results int I, fs, vs : impurity, feature, value
//void find_split_in_range(vector<tuple*> data, vector<int> dataCount, vector<int> invertIdx, int start, int  end, int& fs, double& vs, double& I, vector<bool> skip, const args_t& args) {
void find_split_in_range(vector<tuple*> data, vector<int> dataCount, vector<int> invertIdx, pair<int,int> range, int& fs, double& vs, double& I, vector<bool>skip, const args_t& args) {

  double min = MY_DBL_MAX;
	int start = range.first, end = range.second;

  for (int i = start; i < end; i++) {
    int f = i + 1;
    if (skip[f]) continue;

    int fi;
    double vi, Ii;
    Ii = best_fc_in_feature(data, dataCount, invertIdx, f, fi, vi, args);
    if (Ii < min) {
      min = Ii;
      fs = fi;
      vs = vi;
    }
  }

  I = min;
}

bool find_split_p(vector<tuple*> data, vector<int> dataCount, vector<int> invertIdx, int NF, int& f_split, double& v_split, vector<bool>& skip, const args_t& args) {

  f_split = -1;
  double min = MY_DBL_MAX;
  int n = data.size(), i;

  pair<int, double>* fc = new pair<int,double>[NF];
  double* I = new double[NF];

  int numthreads = args.processors;
  thread** threads = new thread*[numthreads];

  int* F = new int[numthreads];
  double* V = new double[numthreads];
  double* Imp = new double[numthreads];

  for (i = 0; i < numthreads; i++)
//	find_split_in_range(data, dataCount, invertIdx, i*(NF-1)/numthreads, (i+1)*(NF-1)/numthreads, F[i], V[i], Imp[i], skip, args);
//    threads[i] = new thread(find_split_in_range, data, dataCount, invertIdx, i*(NF-1)/numthreads, (i+1)*(NF-1)/numthreads, ref(F[i]), ref(V[i]), ref(Imp[i]), ref(skip), cref(args));
	threads[i] = new thread(bind(find_split_in_range, data, dataCount, invertIdx, make_pair(i*(NF-1)/numthreads, (i+1)*(NF-1)/numthreads), ref(F[i]), ref(V[i]), ref(Imp[i]), ref(skip), cref(args) ));
  for (i = 0; i < numthreads; i++) {
    threads[i]->join();
    delete threads[i];
  }
  delete[] threads;

  for (i = 0; i < numthreads;i++)
    if (Imp[i] < min) {
      min = Imp[i];
      f_split = F[i];
      v_split = V[i];
    }

  delete[] fc;  delete[] I;  delete[] V;  delete[] Imp; delete[] F;
  return min != MY_DBL_MAX;
}

///////////////////////

bool dt_node::entropy_split(data_t data, vector<int> dataCount, vector<int> invertIdx, int NF, int& f_split, double& v_split, int K, bool par) {
  f_split = -1;
  double min = MY_DBL_MAX;
  int n = data.size(), i;

  vector<bool> skip;

  //min E(i=1..k) pi * log(1/pi)

  for (i = 0; i <= NF; i++)
    skip.push_back( (K > 0) ? true : false);

  for (i = 0; i < K; i++) {
    int f;
    do
      f = rand() % (NF-2) + 1;
    while (!skip[f]);
    skip[f] = false;
  }

  //if (K <= 0)
  //return find_split_p(data, NF, f_split, v_split, skip);

  vector<int> location(n, -1);
  for (int f = 1; f < NF; f++) {
	if (skip[f]) continue;
	sort_data_by_feature(location, dataCount, invertIdx, f);
    //sort_data_by_feature(data,f);

    //if (skip[f]) continue;

    /*
 vector< pair<tuple*, int> > tk;
  int z;

  for (z = 0; z < data.size(); z++)
  tk.push_back( pair<tuple*,int>(data[z], f) );
  sort(tk.begin(), tk.end(), mysortpred2);

  for (z = 0; z < data.size(); z++)
  data[z] = tk[z].first;
    */

  int num_c = 7;
  vector<double> c_miss, c_left, c_right;
  //vector<double> p_miss, p_left, p_right;
  //vector<int> freq;
  //int n_left = 0, n_right = 0;
  for (i=0;i<num_c;i++) {
	c_miss.push_back(0.0);
//	p_miss.push_back(0.0);
    c_left.push_back(0.0);
    c_right.push_back(0.0);
//    p_left.push_back(0.0);
//    p_right.push_back(0.0);
//    freq.push_back(0);
  }

    // get impurity (entropy) for missing data
    double M = 0.0, L, R, W_miss = 0.0, W_left = 0.0, W_right = 0.0;
    int missing = 0;
	int loc = location[missing];
    while (data[loc]->features[f] == UNKNOWN && missing < n-1) {
	c_miss[(int)data[loc]->target] += data[loc]->weight;
	W_miss += data[loc]->weight;
        missing++;
	loc = location[missing];
    }

    if (missing == n-1) // all data is missing
      continue;

    int nn = n - missing; // number of data points that arent missing
    // entropy
    // put all data into right side R
    int start = missing;
    for (i = start; i < n; i++) {
	loc = location[i];
      	c_right[(int)data[loc]->target]+=data[loc]->weight;
      	W_right+=data[loc]->weight;
    }
	if (W_right){
    		for (i = 0; i < num_c; i++) {
      			if (c_right[i])
				R += c_right[i]/W_right * log(W_right / c_right[i]);
    		}
	}

    	if (missing && W_miss) {
                for (i = 0; i < num_c; i ++){
                        if (c_miss[i])
                                M += c_miss[i]/W_miss * log(W_miss/c_miss[i]);
                }
    	}
    	L = 0.0;

    // for every i
    // put yi into left side, remove it from the right side, and calculate squared lost
    for (i = start; i < n-1; i++) {
      int j = i - missing;
      int yn = (int)data[loc]->target;

	W_right -= data[loc]->weight;
	W_left += data[loc]->weight;
      //n_right--;
      //n_left++;
	c_left[yn] += data[loc]->weight;
	c_right[yn] -= data[loc]->weight;
      //c_left[yn]++;
      //c_right[yn]--;

      //p_left[yn] += 1.0 / freq[yn];
      //p_right[yn] -= 1.0 / freq[yn];

      // do not consider splitting here if data is the same as next
      if (data[location[i]]->features[f] == data[location[i+1]]->features[f])
	continue;

      L = 0.0;
      int k;
      for (k = 0; k < num_c; k++) {
      	if (c_left[k])
	  L +=  c_left[k]/W_left * log(W_left / c_left[k]);
      }
      R = 0.0;
      for (k = 0; k < num_c; k++) {
	if (c_right[k])
	  R += c_right[k]/W_right * log(W_right / c_right[k]);
      }
      double ssum = W_left+W_right+W_miss;
      double I = W_left/ssum * L +  W_right/ssum * R + W_miss/ssum * M;
      //I = 1.0*n_left/n *

      /*
      L = 0.0, R = 0.0;
      for (i = 0; i < num_c; i++)
	L += p_left;
      */

      if (I < min) {
	min = I;
	f_split = f;
	v_split = (data[location[i]]->features[f] + data[location[i+1]]->features[f])/2;
      }
    }
  }

  return min != MY_DBL_MAX;
}

//////////////////////

bool dt_node::find_split(vector<tuple*> data, vector<int> dataCount, vector<int> invertIdx, int NF, int& f_split, double& v_split, int K, bool par, const args_t& args) {
  if (args.loss == ALG_ENTROPY)
    return entropy_split(data, dataCount, invertIdx, NF, f_split, v_split, K, par);

  f_split = -1;
  double min = MY_DBL_MAX;
  int n = data.size(), i;

  vector<bool> skip;

  //K = NF/2;

  // pick K random features to split on, if specified
  for (i = 0; i <= NF; i++)
    skip.push_back( (K > 0) ? true : false);
  for (i = 0; i < K; i++) {
    int f;
    do
      f = rand() % (NF-2) + 1;
    while (!skip[f]);
    skip[f] = false;
  }

  //if (K >= 10)
  if (args.alg != ALG_FOREST && args.processors!=1)
    return find_split_p(data, dataCount, invertIdx, NF, f_split, v_split, skip, args);

	vector<int> location(n, -1);
  for (int f = 1; f < NF; f++) {
        if (skip[f]) continue;
    // sort data
/*    vector< pair<tuple*, int> > tk;
    int z;
    for (z = 0; z < data.size(); z++)
      tk.push_back( pair<tuple*,int>(data[z], f) );
    sort(tk.begin(), tk.end(), mysortpred2);
    for (z = 0; z < data.size(); z++)
      data[z] = tk[z].first;
*/
	sort_data_by_feature(location, dataCount, invertIdx, f);

    // get impurity (squared loss) for missing data
/*    double M = 0.0;
    int missing = 0;
    while (data[missing]->features[f] == UNKNOWN && missing < n-1) {
      M += data[missing]->target * 1.0;
      missing++;
    }
    if (missing == n-1) // all data is missing
      continue;
    if (missing) {
      double mbar = M * 1.0 / missing;
      M = 0.0;
      for (i = 0; i < missing; i++)
	M += (data[i]->target - mbar) * (data[i]->target - mbar);
    }
    int nn = n - missing; // number of data points that arent missing
*/
  double M = 0.0;
  double W = 0.0;
  int missing = 0;
	int loc = location[missing];
  while (data[loc]->features[f] == UNKNOWN && missing < n-1) {
    //M += data[missing]->target * 1.0;
        M += data[loc]->target * data[loc]->weight;
        W += data[loc]->weight;
    	missing++;
	loc = location[missing];
  }
  if (missing == n-1) // all data is missing
    return MY_DBL_MAX;
  if (missing) {
    //double mbar = M * 1.0 / missing;
        double mbar = M/W;
    M = 0.0;
    for (i = 0; i < missing; i++){
	loc = location[i];
      M += data[loc]->weight * (data[loc]->target - mbar) * (data[loc]->target - mbar);
   }
  }

  double ybl, ybr, s, r, L, R, I, ywl, ywr, WL, WR;
  ybl = ybr = s = r = ywl = ywr = WL = WR = 0.0;

  // put all data into right side R
  // get ybar right and r
  int start = missing;
  for (i = start; i < n; i++) {
	loc = location[i];
    	r += data[loc]->target * data[loc]->target * data[loc]->weight;
    	//ybr += data[i]->target;
    	ywr += data[loc]->target * data[loc]->weight;
    	WR += data[loc]->weight;
  }
  //ybr /= 1.0 * nn;
  //r += 0.0000000001 for precision errors

  // for every i
  // put yi into left side, remove it from the right side, and calculate squared lost
  // impurity of putting all data points into the right tree equals putting all the data into the left tree (such cases are not considered)
  for (i = start; i < n-1; i++) {
    int j = i - missing;
    double yn = data[location[i]]->target;
    double w = data[location[i]]->weight;

    s += w *  yn * yn;
    r -= w * yn * yn;
    ywr -= w * yn;
    ywl += w * yn;
    WL += w;
    WR -= w;

    if (r < 0 && r > -0.000001) r = 0;
    if (r < 0) r = 0;
    //ybl = (j*ybl + yn) / (j+1.0);
    //ybr = ((nn-j)*ybr - yn) / (nn-j-1.0);
    //L = s + WL*ybl*ybl - 2*ybl*ywl;
    //R = r + WR*ybr*ybr - 2*ybr*ywr;
    L = s - ywl*ywl/WL;
    R = r - ywr*ywr/WR;
    // precision errors?
    if (L < 0 && L > -0.0000001) L = 0;
    if (R < 0 && R > -0.0000001) R = 0;
    if (R < 0) R = 0;
    if (L < 0) L = 0;

    if(0)
    if (L < 0 || R < 0 || r < 0)
      printf("Problem %lf %lf %lf\n", L, R, r);

    // do not consider splitting here if data is the same as next
    if (data[location[i]]->features[f] == data[location[i+1]]->features[f])
      continue;

    I = L + R + M ;

      if (I < min) {
        min = I;
        f_split = f;
        v_split = (data[location[i]]->features[f] + data[location[i+1]]->features[f])/2;
      }

  }
}
  return min != MY_DBL_MAX;
}


#include <engine/Neighbours.hpp>
#include <engine/Vectormath.hpp>
#include <engine/Manifoldmath.hpp>
#include <utility/Logging.hpp>
#include <utility/Exception.hpp>

#include <Eigen/Dense>

#include <numeric>
#include <iostream>
#include <cstdio>
#include <cmath>
#include <algorithm>

using namespace Utility;

namespace Engine
{
	namespace Neighbours
	{
		std::vector<scalar> Get_Shell_Radius(const Data::Geometry & geometry, const int n_shells)
		{
			auto shell_radius = std::vector<scalar>(n_shells);

			Vector3 a = geometry.bravais_vectors[0];
			Vector3 b = geometry.bravais_vectors[1];
			Vector3 c = geometry.bravais_vectors[2];

			scalar current_radius=0, dx, min_distance=0;
			int i=0, j=0, k=0;
			int ii, jj, kk;

			// The 15 is a value that is big enough by experience to
			// produce enough needed shells, but is small enough to run sufficiently fast
			int imax = 15, jmax = 15, kmax = 15;
			Vector3 x0={0,0,0}, x1={0,0,0};

			// Abort condidions for all 3 vectors
			if (a.norm() == 0.0) imax = 0;
			if (b.norm() == 0.0) jmax = 0;
			if (c.norm() == 0.0) kmax = 0;

			for (int n = 0; n < n_shells; ++n)
			{
				current_radius = min_distance;
				min_distance = 1e10;
				for (int iatom = 0; iatom < geometry.n_cell_atoms; ++iatom)
				{
				    x0 = geometry.cell_atoms[iatom][0] * a + geometry.cell_atoms[iatom][1] * b + geometry.cell_atoms[iatom][2] * c;
					for (ii = imax; ii >= -imax; --ii)
					{
						for (jj = jmax; jj >= -jmax; --jj)
						{
							for (kk = kmax; kk >= -kmax; --kk)
							{
								for (int jatom = 0; jatom < geometry.n_cell_atoms; ++jatom)
								{
									if ( !( iatom==jatom && ii==0 && jj==0 && kk==0 ) )
									{
									    x1 = geometry.cell_atoms[jatom][0] * a + geometry.cell_atoms[jatom][1] * b + geometry.cell_atoms[jatom][2] * c + ii*a + jj*b + kk*c;
										dx = (x0-x1).norm();

										if (dx - current_radius > 1e-6 && dx < min_distance)
										{
											min_distance = dx;
											shell_radius[n] = dx;
										}
									}
								}//endfor jatom
							}//endfor kk
						}//endfor jj
					}//endfor ii
				}//endfor iatom
			}

			return shell_radius;
		}

		pairfield Get_Pairs_in_Shells(const Data::Geometry & geometry, int nShells)
		{
			auto pairs = pairfield(0);

			auto shell_radius = Get_Shell_Radius(geometry, nShells);

			Vector3 a = geometry.bravais_vectors[0];
			Vector3 b = geometry.bravais_vectors[1];
			Vector3 c = geometry.bravais_vectors[2];

			// The nShells + 10 is a value that is big enough by experience to
			// produce enough needed shells, but is small enough to run sufficiently fast
			int tMax = nShells + 10;
			int imax = tMax, jmax = tMax, kmax = tMax;
			int i,j,k;
			scalar dx, delta, radius;
			Vector3 x0={0,0,0}, x1={0,0,0};

			// Abort condidions for all 3 vectors
			if (a.norm() == 0.0) imax = 0;
			if (b.norm() == 0.0) jmax = 0;
			if (c.norm() == 0.0) kmax = 0;

			for (int iatom = 0; iatom < geometry.n_cell_atoms; ++iatom)
			{
				x0 = geometry.cell_atoms[iatom];
				for (int ishell = 0; ishell < nShells; ++ishell)
				{
					radius = shell_radius[ishell];
					for (i = imax; i >= -imax; --i)
					{
						for (j = jmax; j >= -jmax; --j)
						{
							for (k = kmax; k >= -kmax; --k)
							{
								for (int jatom = 0; jatom < geometry.n_cell_atoms; ++jatom)
								{
									x1 = geometry.cell_atoms[jatom] + i*a + j*b + k*c;
									dx = (x0-x1).norm();
									delta = std::abs(dx - radius);
									if (delta < 1e-6)
									{
										pairs.push_back( {iatom, jatom, {i, j, k} } );
									}
								}//endfor jatom
							}//endfor k
						}//endfor j
					}//endfor i
				}//endfor ishell
			}//endfor iatom

			return pairs;
		}

		neighbourfield Get_Neighbours_in_Shells(const Data::Geometry & geometry, int nShells)
		{
			auto neighbours = neighbourfield(0);

			auto shell_radius = Get_Shell_Radius(geometry, nShells);

			Vector3 a = geometry.bravais_vectors[0];
			Vector3 b = geometry.bravais_vectors[1];
			Vector3 c = geometry.bravais_vectors[2];

			// The nShells + 10 is a value that is big enough by experience to
			// produce enough needed shells, but is small enough to run sufficiently fast
			int tMax = nShells + 10;
			int imax = std::min(tMax, geometry.n_cells[0]-1), jmax = std::min(tMax, geometry.n_cells[1]-1), kmax = std::min(tMax, geometry.n_cells[2]-1);
			int i,j,k;
			scalar dx, delta, radius;
			Vector3 x0={0,0,0}, x1={0,0,0};

			// Abort condidions for all 3 vectors
			if (a.norm() == 0.0) imax = 0;
			if (b.norm() == 0.0) jmax = 0;
			if (c.norm() == 0.0) kmax = 0;

			for (int iatom = 0; iatom < geometry.n_cell_atoms; ++iatom)
			{
				x0 = geometry.cell_atoms[iatom][0] * a + geometry.cell_atoms[iatom][1] * b + geometry.cell_atoms[iatom][2] * c;
				for (int ishell = 0; ishell < nShells; ++ishell)
				{
					radius = shell_radius[ishell];
					for (i = imax; i >= -imax; --i)
					{
						for (j = jmax; j >= -jmax; --j)
						{
							for (k = kmax; k >= -kmax; --k)
							{
								for (int jatom = 0; jatom < geometry.n_cell_atoms; ++jatom)
								{
									x1 = geometry.cell_atoms[jatom][0] * a + geometry.cell_atoms[jatom][1] * b + geometry.cell_atoms[jatom][2] * c + i*a + j*b + k*c;
									dx = (x0-x1).norm();
									delta = std::abs(dx - radius);
									if (delta < 1e-6)
									{
										Neighbour neigh;
										neigh.i = iatom;
										neigh.j = jatom;
										neigh.translations[0] = i;
										neigh.translations[1] = j;
										neigh.translations[2] = k;
										neigh.idx_shell = ishell;
										neighbours.push_back( neigh );
									}
								}//endfor jatom
							}//endfor k
						}//endfor j
					}//endfor i
				}//endfor ishell
			}//endfor iatom

			return neighbours;
		}

		pairfield Get_Pairs_in_Radius(const Data::Geometry & geometry, scalar radius)
		{
			auto pairs = pairfield(0);

			if (radius > 1e-6)
			{
				Vector3 a = geometry.bravais_vectors[0];
				Vector3 b = geometry.bravais_vectors[1];
				Vector3 c = geometry.bravais_vectors[2];

				Vector3 bounds_diff = geometry.bounds_max - geometry.bounds_min;
				Vector3 ratio = {
					bounds_diff[0]/std::max(1, geometry.n_cells[0]),
					bounds_diff[1]/std::max(1, geometry.n_cells[1]),
					bounds_diff[2]/std::max(1, geometry.n_cells[2]) };

				// This should give enough translations to contain all DDI pairs
				int imax = 0, jmax = 0, kmax = 0;
				if ( bounds_diff[0] > 0 )
					imax = std::min(geometry.n_cells[0], (int)(1.1 * radius * geometry.n_cells[0] / bounds_diff[0]));
				if ( bounds_diff[1] > 0 )
					jmax = std::min(geometry.n_cells[1], (int)(1.1 * radius * geometry.n_cells[1] / bounds_diff[1]));
				if ( bounds_diff[2] > 0 )
					kmax = std::min(geometry.n_cells[2], (int)(1.1 * radius * geometry.n_cells[2] / bounds_diff[2]));

				int i,j,k;
				scalar dx;
				Vector3 x0={0,0,0}, x1={0,0,0};

				// Abort condidions for all 3 vectors
				if (a.norm() == 0.0) imax = 0;
				if (b.norm() == 0.0) jmax = 0;
				if (c.norm() == 0.0) kmax = 0;

				for (int iatom = 0; iatom < geometry.n_cell_atoms; ++iatom)
				{
					x0 = geometry.cell_atoms[iatom];
					for (i = imax; i >= -imax; --i)
					{
						for (j = jmax; j >= -jmax; --j)
						{
							for (k = kmax; k >= -kmax; --k)
							{
								for (int jatom = 0; jatom < geometry.n_cell_atoms; ++jatom)
								{
									x1 = geometry.cell_atoms[jatom] + i*a + j*b + k*c;
									dx = (x0-x1).norm();
									if (dx < radius)
									{
										pairs.push_back( {iatom, jatom, {i, j, k} } );
									}
								}//endfor jatom
							}//endfor k
						}//endfor j
					}//endfor i
				}//endfor iatom
			}

			return pairs;
		}

		neighbourfield Get_Neighbours_in_Radius(const Data::Geometry & geometry, scalar radius)
		{
			auto neighbours = neighbourfield(0);

			if (radius > 1e-6)
			{
				Vector3 a = geometry.bravais_vectors[0];
				Vector3 b = geometry.bravais_vectors[1];
				Vector3 c = geometry.bravais_vectors[2];

				Vector3 bounds_diff = geometry.bounds_max - geometry.bounds_min;
				Vector3 ratio = {
					bounds_diff[0]/std::max(1, geometry.n_cells[0]),
					bounds_diff[1]/std::max(1, geometry.n_cells[1]),
					bounds_diff[2]/std::max(1, geometry.n_cells[2]) };

				// This should give enough translations to contain all DDI pairs
				int imax = 0, jmax = 0, kmax = 0;
				if ( bounds_diff[0] > 0 )
					imax = std::min(geometry.n_cells[0], (int)(1.1 * radius * geometry.n_cells[0] / bounds_diff[0]));
				if ( bounds_diff[1] > 0 )
					jmax = std::min(geometry.n_cells[1], (int)(1.1 * radius * geometry.n_cells[1] / bounds_diff[1]));
				if ( bounds_diff[2] > 0 )
					kmax = std::min(geometry.n_cells[2], (int)(1.1 * radius * geometry.n_cells[2] / bounds_diff[2]));

				int i, j, k;
				scalar dx;
				Vector3 x0 = { 0,0,0 }, x1 = { 0,0,0 };

				// Abort condidions for all 3 vectors
				if (a.norm() == 0.0) imax = 0;
				if (b.norm() == 0.0) jmax = 0;
				if (c.norm() == 0.0) kmax = 0;

				for (int iatom = 0; iatom < geometry.n_cell_atoms; ++iatom)
				{
					x0 = geometry.cell_atoms[iatom];
					for (i = imax; i >= -imax; --i)
					{
						for (j = jmax; j >= -jmax; --j)
						{
							for (k = kmax; k >= -kmax; --k)
							{
								for (int jatom = 0; jatom < geometry.n_cell_atoms; ++jatom)
								{
									x1 = geometry.cell_atoms[jatom] + i*a + j*b + k*c;
									dx = (x0 - x1).norm();
									if (dx < radius)
									{
										Neighbour neigh;
										neigh.i = iatom;
										neigh.j = jatom;
										neigh.translations[0] = i;
										neigh.translations[1] = j;
										neigh.translations[2] = k;
										neigh.idx_shell = 0;
										neighbours.push_back( neigh );
									}
								}//endfor jatom
							}//endfor k
						}//endfor j
					}//endfor i
				}//endfor iatom
			}

			return neighbours;
		}

		Vector3 DMI_Normal_from_Pair(const Data::Geometry & geometry, const Pair & pair, int chirality)
		{
			Vector3 ta = geometry.bravais_vectors[0];
			Vector3 tb = geometry.bravais_vectors[1];
			Vector3 tc = geometry.bravais_vectors[2];

			int da = pair.translations[0];
			int db = pair.translations[1];
			int dc = pair.translations[2];

			Vector3 ipos = geometry.cell_atoms[pair.i];
			Vector3 jpos = geometry.cell_atoms[pair.j] + da*ta + db*tb + dc*tc;

			if (chirality == 1)
			{
				// Bloch chirality
				return (jpos - ipos).normalized();
			}
			else if (chirality == -1)
			{
				// Inverse Bloch chirality
				return (ipos - jpos).normalized();
			}
			else if (chirality == 2)
			{
				// Neel chirality (surface)
				return (jpos - ipos).normalized().cross(Vector3{0,0,1});
			}
			else if (chirality == -2)
			{
				// Inverse Neel chirality (surface)
				return Vector3{0,0,1}.cross((jpos - ipos).normalized());
			}
			else
			{
				return Vector3{ 0,0,0 };
			}
		}

		void DDI_from_Pair(const Data::Geometry & geometry, const Pair & pair, scalar & magnitude, Vector3 & normal)
		{
			Vector3 ta = geometry.bravais_vectors[0];
			Vector3 tb = geometry.bravais_vectors[1];
			Vector3 tc = geometry.bravais_vectors[2];

			int da = pair.translations[0];
			int db = pair.translations[1];
			int dc = pair.translations[2];

			Vector3 ipos = geometry.cell_atoms[pair.i];
			Vector3 jpos = geometry.cell_atoms[pair.j] + da*ta + db*tb + dc*tc;

			// Calculate positions and difference vector
			Vector3 vector_ij = jpos - ipos;

			// Length of difference vector
			magnitude = vector_ij.norm();
			normal = vector_ij.normalized();
		}

	}// end Namespace Neighbours
}// end Namespace Engine


////////////////////////////////////////////////////////////////////////////////
// IncompressibleBalloonEnergy.hh
////////////////////////////////////////////////////////////////////////////////
/*! @file
//  Implementation of the incompressible neo-Hookean-based strain energy
//  density used in Skouras 2014: Designing Inflatable Structures (before
//  homogenizing away the wrinkles using a relaxed energy density).
//
//  This energy is implemented as a function of the right Green-Green
//  deformation tensor.
*/
//  Author:  Julian Panetta (jpanetta), julian.panetta@gmail.com
//  Created:  04/04/2019 18:19:10
////////////////////////////////////////////////////////////////////////////////
#ifndef INCOMPRESSIBLEBALLOONENERGY_HH
#define INCOMPRESSIBLEBALLOONENERGY_HH

#include <Eigen/Dense>
#include <array>
#include <MeshFEM/EnergyDensities/Tensor.hh>

template<typename Real>
struct IncompressibleBalloonEnergy {
    using V2d  = Eigen::Matrix<Real, 2, 1>;
    using M2d  = Eigen::Matrix<Real, 2, 2>;

    IncompressibleBalloonEnergy() { }

    template<typename Derived>
    IncompressibleBalloonEnergy(const Eigen::MatrixBase<Derived> &C) { setMatrix(C); }

    template<typename Derived>
    void setMatrix(const Eigen::MatrixBase<Derived> &C) {
        static_assert((Derived::RowsAtCompileTime == 2) && (Derived::ColsAtCompileTime == 2), "Only 2x2 supported for now");

        Real a = C(0, 0),
             b = C(0, 1),
             c = C(1, 1);
        if (std::abs(b - C(1, 0)) > 1e-15) throw std::runtime_error("Asymmetric matrix");

        m_C = C;
        m_trace_C = C.trace();
        m_det_C = a * c - b * b;
        m_grad_det_C <<  c, -b,
                        -b,  a;
    }

    Real energy() const {
        return stiffness * (m_trace_C + 1.0 / m_det_C - 3.0);
    }

    Real denergy(const M2d &dC) const {
        return stiffness * (dC.trace() - (1.0 / (m_det_C * m_det_C)) * doubleContract(m_grad_det_C, dC));
    }

    M2d denergy() const {
        return stiffness * (M2d::Identity() - (1.0 / (m_det_C * m_det_C)) * m_grad_det_C);
    }

    Real d2energy(const M2d &dC_a, const M2d &dC_b) const {
        return stiffness * ((2.0 / (m_det_C * m_det_C * m_det_C)) * doubleContract(m_grad_det_C, dC_a) * doubleContract(m_grad_det_C, dC_b)
                          - (1.0 / (m_det_C * m_det_C)) * (dC_a(1, 1) * dC_b(0, 0) + dC_a(0, 0) * dC_b(1, 1) - 2 * dC_a(0, 1) * dC_b(0, 1)));
    }

    M2d delta_denergy(const M2d &dC) const {
        M2d adj_dC;
        adj_dC << dC(1, 1), -dC(0, 1),
                 -dC(1, 0),  dC(0, 0);
        return (((2.0 * stiffness / (m_det_C * m_det_C * m_det_C)) * doubleContract(m_grad_det_C, dC)) * m_grad_det_C
                     - (stiffness / (m_det_C * m_det_C))                                               * adj_dC);
    }

    // Second derivatives evaluated at the reference configuration
    M2d delta_denergy_undeformed(const M2d &dC) const {
        return stiffness * (dC.trace() * M2d::Identity() + dC);
    }

    Real stiffness = 1.0;

    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
private:
    Real m_trace_C, m_det_C;
    M2d m_grad_det_C;
    M2d m_C;
};

#endif /* end of include guard: INCOMPRESSIBLEBALLOONENERGY_HH */


/* =========================================================================
   Copyright (c) 2010-2014, 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
============================================================================= */

/* Computation of eigenvalues of a symmetric, tridiagonal matrix using
 * bisection.
 */

#ifndef NDEBUG
  #define NDEBUG
#endif

// includes, system
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

// includes, project

#include "viennacl/scalar.hpp"
#include "viennacl/vector.hpp"
#include "viennacl/matrix.hpp"

#include "viennacl/linalg/bisect_gpu.hpp"
#include "viennacl/linalg/bisect.hpp"
#include "viennacl/linalg/tql2.hpp"

#include <examples/benchmarks/benchmark-utils.hpp>
//#include <Eigen/Eigenvalues>
//using namespace Eigen;

#define EPS 10.0e-4

typedef float NumericT;

////////////////////////////////////////////////////////////////////////////////
/// \brief initInputData   Initialize the diagonal and superdiagonal elements of
///                        the matrix
/// \param diagonal        diagonal elements of the matrix
/// \param superdiagonal   superdiagonal elements of the matrix
/// \param mat_size        Dimension of the matrix
///
void
initInputData(viennacl::vector<NumericT> &diagonal, viennacl::vector<NumericT> &superdiagonal, const unsigned int mat_size)
{

  srand(time(NULL));

#define RANDOM_VALUES 0
  if (RANDOM_VALUES == true)
  {
    // Initialize diagonal and superdiagonal elements with random values
    for (unsigned int i = 0; i < mat_size; ++i)
    {
        diagonal[i] =      static_cast<NumericT>(2.0 * (((double)rand()
                                     / (double) RAND_MAX) - 0.5));
        superdiagonal[i] = static_cast<NumericT>(2.0 * (((double)rand()
                                     / (double) RAND_MAX) - 0.5));
    }
  }
  else
  {
    // Initialize diagonal and superdiagonal elements with modulo values
    // This will cause in many multiple eigenvalues.
    for (unsigned int i = 0; i < mat_size; ++i)
    {
       diagonal[i] = ((NumericT)(i % 37)) - 4.5f;
       superdiagonal[i] = ((NumericT)(i % 5)) - 4.5f;
    }
  }
  // the first element of s is used as padding on the device (thus the
  // whole vector is copied to the device but the kernels are launched
  // with (s+1) as start address
  superdiagonal[0] = 0.0f;
}

////////////////////////////////////////////////////////////////////////////////
//! Run a simple test
////////////////////////////////////////////////////////////////////////////////
bool
runTest(const int mat_size, std::vector<double> &av_time_all, unsigned int time_index)
    {
    bool bResult = false;
    viennacl::vector<NumericT> diagonal(mat_size);
    viennacl::vector<NumericT> superdiagonal(mat_size);
    viennacl::vector<NumericT> eigenvalues_bisect(mat_size);
    std::vector<NumericT> eigenvalues_bisect_cpu(mat_size);

    // -------Start the bisection algorithm------------
    std::cout << "Matrix size: " << mat_size << std::endl;

    unsigned int iterations = 10;
    double time_all     = 0.0;

    for(unsigned int i = 0; i < iterations; i++)
    {
      initInputData(diagonal, superdiagonal, mat_size);

      Timer timer;
      timer.start();

      // bisection - gpu
      bResult = viennacl::linalg::bisect(diagonal, superdiagonal, eigenvalues_bisect);
      viennacl::backend::finish();     // sync
      //---Run the tql algorithm-----------------------------------
     // viennacl::linalg::tql1<NumericT>(mat_size, diagonal, superdiagonal);
     // bResult = true;

      // Run the bisect algorithm for CPU only
      //eigenvalues_bisect_cpu = viennacl::linalg::bisect(diagonal, superdiagonal);
      // bResult = true;

      time_all     += timer.get() * 1000;
      if (bResult == false)
       return false;

    }

    std::cout << "Time: \t" << time_all / (double)iterations << "ms" << std::endl << std::endl;

    av_time_all[time_index] = time_all / (double)iterations;

  return bResult;

}

////////////////////////////////////////////////////////////////////////////////
// Program main
////////////////////////////////////////////////////////////////////////////////
int
main(int argc, char **argv)
{
    bool test_result = true;
    unsigned int time_index = 0;
    std::vector<double> av_time_all(500);
    std::vector<unsigned int> mat_sizes(500);

    for( unsigned int mat_size = 16;
         mat_size < 600;
         mat_size = mat_size * 1.15, time_index++)
      {
      test_result = runTest(mat_size, av_time_all, time_index);
      mat_sizes[time_index] = mat_size;

      if(test_result == true)
      {
        //std::cout << "Success!" << std::endl << std::endl;
      }
      else
      {
        std::cout << "---FAIL---" << std::endl;
        exit(EXIT_FAILURE);
      }
    }

    std::cout << "Times" << std::endl;
    for(unsigned int i = 0; i < time_index; i++)
    {
      std::cout <<  mat_sizes[i] << "\t" << av_time_all[i] << std::endl;
    }

}


#include <iostream>

#include <Eigen/Sparse>
#include <Eigen/Dense>

#include <unsupported/Eigen/SparseExtra> // For reading MatrixMarket files

#include <amgcl/backend/eigen.hpp>
#include <amgcl/make_solver.hpp>
#include <amgcl/solver/bicgstab.hpp>
#include <amgcl/amg.hpp>
#include <amgcl/coarsening/smoothed_aggregation.hpp>
#include <amgcl/relaxation/spai0.hpp>

int main(int argc, char *argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <matrix.mm>" << std::endl;
        return 1;
    }

    // Read sparse matrix from MatrixMarket format.
    // In general this should come pre-assembled.
    Eigen::SparseMatrix<double, Eigen::RowMajor> A;
    Eigen::loadMarket(A, argv[1]);

    // Use vector of ones as RHS for simplicity:
    Eigen::VectorXd f = Eigen::VectorXd::Constant(A.rows(), 1.0);

    // Zero initial approximation:
    Eigen::VectorXd x = Eigen::VectorXd::Zero(A.rows());

    // Setup the solver:
    typedef amgcl::make_solver<
        amgcl::amg<
            amgcl::backend::eigen<double>,
            amgcl::coarsening::smoothed_aggregation,
            amgcl::relaxation::spai0
            >,
        amgcl::solver::bicgstab<amgcl::backend::eigen<double> >
        > Solver;

    Solver solve(A);
    std::cout << solve << std::endl;

    // Solve the system for the given RHS:
    int    iters;
    double error;
    std::tie(iters, error) = solve(f, x);

    std::cout << iters << " " << error << std::endl;
}