//
// Libor Novak
// 04/12/2017
//
// Various kinds of statistics on the accumulators - to find out if it will gives us some interesting extra
// information.
//

#include <caffe/caffe.hpp>
#include "caffe/util/bbtxt.hpp"

// This code only works with OpenCV!
#ifdef USE_OPENCV

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <algorithm>
#include <iosfwd>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include <boost/program_options.hpp>
namespace po = boost::program_options;

namespace {

    class Hist2D
    {
    public:

        Hist2D (double xmin, double xmax, double ymin, double ymax, int num_bins)
            : xmin(xmin), xmax(xmax), ymin(ymin), ymax(ymax),
              xspread(xmax-xmin), yspread(ymax-ymin),
              hist(num_bins, num_bins, CV_64FC1, cv::Scalar(0)),
              counts(num_bins, num_bins, CV_64FC1, cv::Scalar(0))
        {
        }

        void addEntry (double x, double y, double weight=1.0)
        {
            int col = std::round((x - this->xmin) / this->xspread * hist.cols);
            int row = std::round((y - this->ymin) / this->yspread * hist.rows);

            if (col >= 0 && row >= 0 && col < hist.cols && row < hist.rows)
            {
                hist.at<double>(row, col) += weight;
                counts.at<double>(row, col) += 1.0;
            }
            else
            {
                std::cout << "Out of bounds: " << x << ", " << y << " => " << col << ", " << row << std::endl;
            }
        }

        cv::Mat normalized ()
        {
            cv::Mat hist_norm; this->hist.copyTo(hist_norm);
            double m; cv::minMaxLoc(hist_norm, 0, &m);

            hist_norm *= 1.0 / m;

            return hist_norm;
        }

        cv::Mat countNormalized ()
        {
            cv::Mat hist_norm;
            cv::divide(this->hist, this->counts, hist_norm);
            return hist_norm;
        }

        cv::Mat countNormalizedNormalized ()
        {
            cv::Mat hist_norm;
            cv::divide(this->hist, this->counts, hist_norm);

            double m; cv::minMaxLoc(hist_norm, 0, &m);
            hist_norm *= 1.0 / m;

            return hist_norm;
        }

        friend std::ostream& operator<< (std::ostream& os, const Hist2D &h);

        // --------------------------------------- PUBLIC MEMBERS ---------------------------------------- //
        double xmin, xmax, ymin, ymax;
        double xspread, yspread;
        cv::Mat hist;
        cv::Mat counts;
        int total;
    };

    std::ostream& operator<< (std::ostream& os, const Hist2D &h)
    {
        os << h.hist;
        return os;
    }

    std::vector<Hist2D> hist_wh_neg;
    std::vector<Hist2D> hist_tl_neg;
    std::vector<Hist2D> hist_br_neg;
    std::vector<Hist2D> hist_wh_pos;
    std::vector<Hist2D> hist_tl_pos;
    std::vector<Hist2D> hist_br_pos;

    std::vector<Hist2D> hist_wh_car_g_bb;
    std::vector<Hist2D> hist_wh_notcar_g_bb;

}

/**
 * @brief Wraps the input layer into a vector of cv::Mat so we could assign data to it more easily
 * @param input_layer Pointer to the net input layer blob
 * @param input_channels Vector of cv::Mat, which will be assigned
 */
void wrapInputLayer (caffe::Blob<float>* input_layer, std::vector<cv::Mat> &out_input_channels)
{
    out_input_channels.clear();

    int height = input_layer->shape(2);
    int width  = input_layer->shape(3);

    float* input_data = input_layer->mutable_cpu_data();

    for (int i = 0; i < input_layer->shape(1); ++i)
    {
        cv::Mat channel(height, width, CV_32FC1, input_data);
        out_input_channels.push_back(channel);
        input_data += width * height;
    }
}

void histogramOfCoords (caffe::Blob<float> *output, int a, const std::string &name, const std::vector<BB2D> &gt_bbs)
{
    // Build probabilistic target (ground truth) accumulator - we need it to determine positive and negative
    // pixels
    static std::map<std::string, std::pair<double, double>> size_bounds;
    static std::map<std::string, double> scales;
    if (size_bounds.size() == 0)
    {
        // WARNING! These are size bounds for "macc_0.3_r2_x2_to_x16"!!
        size_bounds.insert(std::make_pair("acc_x2", std::make_pair(22.25, 55.5)));
        size_bounds.insert(std::make_pair("acc_x4", std::make_pair(44.5, 111.0)));
        size_bounds.insert(std::make_pair("acc_x8", std::make_pair(89.0, 222.0)));
        size_bounds.insert(std::make_pair("acc_x16", std::make_pair(178.0, 444.0)));
        scales.insert(std::make_pair("acc_x2", 2.0));
        scales.insert(std::make_pair("acc_x4", 4.0));
        scales.insert(std::make_pair("acc_x8", 8.0));
        scales.insert(std::make_pair("acc_x16", 16.0));
    }
    cv::Mat acc_gt_prob(output->shape(2), output->shape(3), CV_32FC1, cv::Scalar(0.0f));
    for (const BB2D &gt_bb: gt_bbs)
    {
        double size = std::max(gt_bb.width(), gt_bb.height());
        if (size > size_bounds[name].first && size < size_bounds[name].second)
        {
            // This ground truth should be detected by this accumulator
            cv::Point2d co = gt_bb.center();
            cv::circle(acc_gt_prob, cv::Point(co.x/scales[name], co.y/scales[name]), 3, cv::Scalar(1.0f), -1);
        }
    }

//    cv::imshow("gt acc " + std::to_string(a), acc_gt_prob);

    float *data_output = output->mutable_cpu_data();

    // 2D bounding box
    cv::Mat acc_prob(output->shape(2), output->shape(3), CV_32FC1, data_output+output->offset(0, 0));
    cv::Mat acc_xmin(output->shape(2), output->shape(3), CV_32FC1, data_output+output->offset(0, 1));
    cv::Mat acc_ymin(output->shape(2), output->shape(3), CV_32FC1, data_output+output->offset(0, 2));
    cv::Mat acc_xmax(output->shape(2), output->shape(3), CV_32FC1, data_output+output->offset(0, 3));
    cv::Mat acc_ymax(output->shape(2), output->shape(3), CV_32FC1, data_output+output->offset(0, 4));

//    cv::imshow("detected acc " + std::to_string(a), acc_prob);

    // Extract detected boxes - only extract local maxima from 3x3 neighborhood
    for (int i = 0; i < acc_prob.rows; ++i)
    {
        for (int j = 0; j < acc_prob.cols; ++j)
        {
            float label = acc_gt_prob.at<float>(i, j);

            double w = acc_xmax.at<float>(i, j) - acc_xmin.at<float>(i, j);
            double h = acc_ymax.at<float>(i, j) - acc_ymin.at<float>(i, j);

            if (label > 0.0f)
            {
                // This is a positive pixel
                hist_wh_pos[a].addEntry(w, h);
                hist_tl_pos[a].addEntry(acc_xmin.at<float>(i, j), acc_ymin.at<float>(i, j));
                hist_br_pos[a].addEntry(acc_xmax.at<float>(i, j), acc_ymax.at<float>(i, j));
                hist_wh_car_g_bb[a].addEntry(w, h, 1.0);
                hist_wh_notcar_g_bb[a].addEntry(w, h, 0.0);
            }
            else
            {
                // This is a background pixel
                hist_wh_neg[a].addEntry(w, h);
                hist_tl_neg[a].addEntry(acc_xmin.at<float>(i, j), acc_ymin.at<float>(i, j));
                hist_br_neg[a].addEntry(acc_xmax.at<float>(i, j), acc_ymax.at<float>(i, j));
                hist_wh_car_g_bb[a].addEntry(w, h, 0.0);
                hist_wh_notcar_g_bb[a].addEntry(w, h, 1.0);
            }
        }
    }
}

void computeStatistics (const std::string &path_image, const std::shared_ptr<caffe::Net<float>> &net,
                        const std::map<std::string, std::vector<BB2D>> &gt_bbs_list)
{
    caffe::Blob<float>* input_layer  = net->input_blobs()[0];

    std::vector<cv::Mat> input_channels;

    // Read the image
    cv::Mat image = cv::imread(path_image, CV_LOAD_IMAGE_COLOR);
    // Convert to zero mean and unit variance
    cv::Mat imagef; image.convertTo(imagef, CV_32FC3);
    imagef -= cv::Scalar(128.0f, 128.0f, 128.0f);
    imagef *= 1.0f/128.0f;

    // Ground truth bounding boxes
    std::vector<BB2D> gt_bbs;
    auto gt_bbsi = gt_bbs_list.find(path_image);
    if (gt_bbsi == gt_bbs_list.end())
    {
        LOG(WARNING) << "No ground truth for image '" << path_image << "'";
    }
    else
    {
        gt_bbs = (*gt_bbsi).second;
    }

    // Reshape the network
    input_layer->Reshape(1, input_layer->shape(1), imagef.rows, imagef.cols);
    net->Reshape();

    // Prepare the cv::Mats for input
    wrapInputLayer(input_layer, input_channels);
    // Copy the image to the input layer of the network
    cv::split(imagef, input_channels);

    net->Forward();

    // For each accumulator
    for (int a = 0; a < net->output_blobs().size(); ++a)
    {
        histogramOfCoords(net->output_blobs()[a], a, net->blob_names()[net->output_blob_indices()[a]], gt_bbs);
    }

//    cv::imshow("image", image);
//    cv::waitKey(0);
}

void runStatisticsComputation (const std::string &path_prototxt, const std::string &path_caffemodel,
                               const std::string &path_image_list, const std::string &path_gt_bbtxt,
                               const std::string &path_out)
{
#ifdef CPU_ONLY
    caffe::Caffe::set_mode(caffe::Caffe::CPU);
#else
    caffe::Caffe::set_mode(caffe::Caffe::GPU);
#endif

    // Create network and load trained weights from caffemodel file
    auto net = std::make_shared<caffe::Net<float>>(path_prototxt, caffe::TEST);
    net->CopyTrainedLayersFrom(path_caffemodel);

    caffe::Blob<float>* input_layer  = net->input_blobs()[0];
    caffe::Blob<float>* output_layer = net->output_blobs()[0];

    CHECK_EQ(net->num_inputs(), 1) << "Network should have exactly one input.";
    CHECK_EQ(input_layer->shape(1), 3) << "Input layer must have 3 channels.";
    CHECK_EQ(output_layer->shape(1), 5) << "Unsupported network, only 5 channels!";

    std::ifstream infile(path_image_list.c_str());
    CHECK(infile) << "Unable to open image list TXT file '" << path_image_list << "'!";
    std::string line;

    // Load ground truth
    std::map<std::string, std::vector<BB2D>> gt_bbs_list = readBBTXTFile(path_gt_bbtxt);

    for (int i = 0; i < net->output_blobs().size(); ++i)
    {
        hist_wh_neg.emplace_back(0, 2, 0, 2, 200);
        hist_tl_neg.emplace_back(-1, 1, -1, 1, 200);
        hist_br_neg.emplace_back(0, 2, 0, 2, 200);
        hist_wh_pos.emplace_back(0, 2, 0, 2, 200);
        hist_tl_pos.emplace_back(-1, 1, -1, 1, 200);
        hist_br_pos.emplace_back(0, 2, 0, 2, 200);
        hist_wh_car_g_bb.emplace_back(0, 2, 0, 2, 200);
        hist_wh_notcar_g_bb.emplace_back(0, 2, 0, 2, 200);
    }

    // -- RUN THE DETECTOR ON EACH IMAGE -- //
    while (std::getline(infile, line))
    {
        LOG(INFO) << line;
        CHECK(boost::filesystem::exists(line)) << "Image '" << line << "' not found!";

        // Detect bbs on the image
        computeStatistics(line, net, gt_bbs_list);
    }

    for (int i = 0; i < net->output_blobs().size(); ++i)
    {
//        {
//            std::vector<cv::Mat> chs_tl;
//            chs_tl.push_back(cv::Mat::zeros(hist_tl_pos[i].hist.size(), CV_64FC1));
//            chs_tl.push_back(hist_tl_pos[i].normalized());
//            chs_tl.push_back(hist_tl_neg[i].normalized());
//            cv::line(chs_tl[0], cv::Point(chs_tl[0].cols/2,0), cv::Point(chs_tl[0].cols/2, chs_tl[0].rows), cv::Scalar(0.5), 2);
//            cv::line(chs_tl[0], cv::Point(0,chs_tl[0].rows/2), cv::Point(chs_tl[0].cols, chs_tl[0].rows/2), cv::Scalar(0.5), 2);
//            cv::Mat comb_tl; cv::merge(chs_tl, comb_tl);
//            cv::imshow("Normalized tl histogram " + std::to_string(i), comb_tl);
//            cv::imwrite(path_out + "/hist_tl_" + net->blob_names()[net->output_blob_indices()[i]] + ".png", comb_tl*255);
//        }

//        {
//            std::vector<cv::Mat> chs_br;
//            chs_br.push_back(cv::Mat::zeros(hist_br_pos[i].hist.size(), CV_64FC1));
//            chs_br.push_back(hist_br_pos[i].normalized());
//            chs_br.push_back(hist_br_neg[i].normalized());
//            cv::line(chs_br[0], cv::Point(chs_br[0].cols/2,0), cv::Point(chs_br[0].cols/2, chs_br[0].rows), cv::Scalar(0.5));
//            cv::line(chs_br[0], cv::Point(0,chs_br[0].rows/2), cv::Point(chs_br[0].cols, chs_br[0].rows/2), cv::Scalar(0.5));
//            cv::Mat comb_br; cv::merge(chs_br, comb_br);
//            cv::imshow("Normalized br histogram " + std::to_string(i), comb_br);
//            cv::imwrite(path_out + "/hist_br_" + net->blob_names()[net->output_blob_indices()[i]] + ".png", comb_br*255);
//        }

        {
            // P(BB|CAR_GT) WxH
            std::vector<cv::Mat> chs;
            chs.push_back(cv::Mat::zeros(hist_wh_pos[i].hist.size(), CV_64FC1));
            chs.push_back(hist_wh_pos[i].normalized());
            chs.push_back(hist_wh_neg[i].normalized());
            cv::line(chs[0], cv::Point(chs[0].cols/2,0), cv::Point(chs[0].cols/2, chs[0].rows), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,chs[0].rows/2), cv::Point(chs[0].cols, chs[0].rows/2), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,0), cv::Point(chs[0].cols, chs[0].rows), cv::Scalar(0.5));
            cv::Mat comb_wh; cv::merge(chs, comb_wh);
//            cv::imshow("P(BB|CAR_GT) WxH " + std::to_string(i), comb_wh);
            cv::imwrite(path_out + "/hist_wh_bb_g_car_notcar_" + net->blob_names()[net->output_blob_indices()[i]] + ".png", comb_wh*255);
        }

        {
            // P(CAR_GT|BB) WxH
            std::vector<cv::Mat> chs;
            chs.push_back(cv::Mat::zeros(hist_wh_car_g_bb[i].hist.size(), CV_64FC1));
            chs.push_back(hist_wh_car_g_bb[i].countNormalized());
            chs.push_back(cv::Mat::zeros(hist_wh_car_g_bb[i].hist.size(), CV_64FC1));
            cv::line(chs[0], cv::Point(chs[0].cols/2,0), cv::Point(chs[0].cols/2, chs[0].rows), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,chs[0].rows/2), cv::Point(chs[0].cols, chs[0].rows/2), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,0), cv::Point(chs[0].cols, chs[0].rows), cv::Scalar(0.5));
            cv::Mat comb_cf; cv::merge(chs, comb_cf);
//            cv::imshow("P(CAR_GT|BB) WxH " + std::to_string(i), comb_cf);
            cv::imwrite(path_out + "/hist_wh_car_g_bb_" + net->blob_names()[net->output_blob_indices()[i]] + ".png", comb_cf*255);
        }

        {
            // P(NOT_CAR_GT|BB) WxH
            std::vector<cv::Mat> chs;
            chs.push_back(cv::Mat::zeros(hist_wh_car_g_bb[i].hist.size(), CV_64FC1));
            chs.push_back(cv::Mat::zeros(hist_wh_car_g_bb[i].hist.size(), CV_64FC1));
            chs.push_back(hist_wh_notcar_g_bb[i].countNormalized());
            cv::line(chs[0], cv::Point(chs[0].cols/2,0), cv::Point(chs[0].cols/2, chs[0].rows), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,chs[0].rows/2), cv::Point(chs[0].cols, chs[0].rows/2), cv::Scalar(0.5));
            cv::line(chs[0], cv::Point(0,0), cv::Point(chs[0].cols, chs[0].rows), cv::Scalar(0.5));
            cv::Mat comb_cf; cv::merge(chs, comb_cf);
//            cv::imshow("P(NOT_CAR_GT|BB) WxH " + std::to_string(i), comb_cf);
            cv::imwrite(path_out + "/hist_wh_notcar_g_bb_" + net->blob_names()[net->output_blob_indices()[i]] + ".png", comb_cf*255);
        }

    }

//    cv::waitKey();
}

// -----------------------------------------------  MAIN  ------------------------------------------------ //

struct ProgramArguments
{
    std::string path_prototxt;
    std::string path_caffemodel;
    std::string path_image_list;
    std::string path_gt_bbtxt;
    std::string path_out;
};

/**
 * @brief Parses arguments of the program
 */
void parseArguments (int argc, char** argv, ProgramArguments &pa)
{
    try {
        po::options_description desc("Arguments");
        desc.add_options()
            ("help", "Print help")
            ("prototxt", po::value<std::string>(&pa.path_prototxt)->required(),
             "Model file of the network (*.prototxt)")
            ("caffemodel", po::value<std::string>(&pa.path_caffemodel)->required(),
             "Weight file of the network (*.caffemodel)")
            ("image_list", po::value<std::string>(&pa.path_image_list)->required(),
             "Path to a TXT file with paths to the images to be tested")
            ("gt_bbtxt", po::value<std::string>(&pa.path_gt_bbtxt)->required(),
             "Path to a BBTXT file with ground truth annotation for the images in image list")
            ("path_out", po::value<std::string>(&pa.path_out)->required(),
             "Path to the output folder")
        ;

        po::positional_options_description positional;
        positional.add("prototxt", 1);
        positional.add("caffemodel", 1);
        positional.add("image_list", 1);
        positional.add("gt_bbtxt", 1);
        positional.add("path_out", 1);

        // Parse the input arguments
        po::variables_map vm;
        po::store(po::command_line_parser(argc, argv).options(desc).positional(positional).run(), vm);

        if (vm.count("help")) {
            std::cout << "Usage: ./macc_statistics path/f.prototxt path/f.caffemodel path/image_list.txt path/out/folder\n";
            std::cout << desc;
            exit(EXIT_SUCCESS);
        }

        po::notify(vm);

        if (!boost::filesystem::exists(pa.path_prototxt))
        {
            std::cerr << "ERROR: File '" << pa.path_prototxt << "' does not exist!" << std::endl;
            exit(EXIT_FAILURE);
        }
        if (!boost::filesystem::exists(pa.path_caffemodel))
        {
            std::cerr << "ERROR: File '" << pa.path_caffemodel << "' does not exist!" << std::endl;
            exit(EXIT_FAILURE);
        }
        if (!boost::filesystem::exists(pa.path_image_list))
        {
            std::cerr << "ERROR: File '" << pa.path_image_list << "' does not exist!" << std::endl;
            exit(EXIT_FAILURE);
        }
        if (!boost::filesystem::exists(pa.path_gt_bbtxt))
        {
            std::cerr << "ERROR: File '" << pa.path_gt_bbtxt << "' does not exist!" << std::endl;
            exit(EXIT_FAILURE);
        }
        if (not boost::filesystem::exists(pa.path_out))
        {
            std::cerr << "ERROR: Output folder '" << pa.path_out << "' does not exist!" << std::endl;
            exit(EXIT_FAILURE);
        }
    }
    catch(std::exception& e)
    {
        std::cerr << e.what() << "\n";
        exit(EXIT_FAILURE);
    }
}

int main (int argc, char** argv)
{
    FLAGS_logtostderr = 1;
    FLAGS_minloglevel = ::google::INFO;
    ::google::InitGoogleLogging(argv[0]);

    ProgramArguments pa;
    parseArguments(argc, argv, pa);

    runStatisticsComputation(pa.path_prototxt, pa.path_caffemodel, pa.path_image_list, pa.path_gt_bbtxt, pa.path_out);

    return EXIT_SUCCESS;
}

#else
int main(int argc, char** argv) {
    LOG(FATAL) << "This example requires OpenCV; compile with USE_OPENCV.";
}
#endif  // USE_OPENCV


#include <cmath>
#include <fstream>
#include <iostream>
#include <iterator>
#include <algorithm>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include "cmdline.h"
#include "vcf.h"

using std::size_t;

// TODO
// Chi-squared test for independance
// void calc_chisq(const std::vector<allele_t> &x, const std::vector<allele_t> &y);

namespace {

struct Parameter
{
    std::string vcf;
    std::string out;
    std::string loc;
    double rsq = 0.5;
    int maxdist = 500000;
} par ;

// LD (D')  -  Lewontin, R.C. (1964). Genetics 49(1), 49-67.
// LD (r^2) -  Hill, W.G., and Robertson, A. (1968). Theor Appl Genet 38(6), 226-231.

void calc_dprime_rsq_kernel(double pa, double pb, double pab, double &dprime, double &rsq)
{
    auto D = pab - pa*pb;
    auto Dmax = D > 0 ? std::min(pa*(1-pb),(1-pa)*pb) : std::min(pa*pb,(1-pa)*(1-pb));
    dprime = std::fabs(D) / Dmax;  // fabs is required for multi-allelic D'
    rsq = D*D / (pa * (1-pa) * pb * (1-pb));
}

void calc_dprime_rsq_hap2(const std::vector<allele_t> &x, const std::vector<allele_t> &y, double &dprime, double &rsq)
{
    static const allele_t a = 1, b = 1;

    auto n = x.size();
    int nt = 0, na = 0, nb = 0, nab = 0;

    for (size_t i = 0; i < n; ++i) {
        if (x[i] && y[i]) {
            ++nt;
            if (x[i] == a) {
                ++na;
                if (y[i] == b)
                    ++nab;
            }
            if (y[i] == b)
                ++nb;
        }
    }

    dprime = rsq = std::numeric_limits<double>::quiet_NaN();

    if (nt == 0)
        return;

    double fnt = nt;

    calc_dprime_rsq_kernel(na/fnt, nb/fnt, nab/fnt, dprime, rsq);
}

void calc_dprime_rsq_hap(const std::vector<allele_t> &x, const std::vector<allele_t> &y, double &dprime, double &rsq)
{
    auto n = x.size();

    auto xmax = * std::max_element(x.begin(), x.end());
    auto ymax = * std::max_element(y.begin(), y.end());

    dprime = rsq = std::numeric_limits<double>::quiet_NaN();

    if (xmax < 2 || ymax < 2)
        return;

    if (xmax == 2 && ymax == 2) {
        calc_dprime_rsq_hap2(x, y, dprime, rsq);
        return;
    }

    std::vector<size_t> idx;
    idx.reserve(n);
    for (size_t i = 0; i < n; ++i) {
        if (x[i] && y[i])
            idx.push_back(i);
    }

    if ( idx.empty() )
        return;

    dprime = rsq = 0.0;

    for (allele_t a = 1; a <= xmax; ++a) {
        for (allele_t b = 1; b <= ymax; ++b) {
            int na = 0, nb = 0, nab = 0;

            for (auto i : idx) {
                if (x[i] == a) {
                    ++na;
                    if (y[i] == b)
                        ++nab;
                }
                if (y[i] == b)
                    ++nb;
            }

            if (na == 0 || nb == 0)
                continue;

            double nt = idx.size();
            double pa = na / nt;
            double pb = nb / nt;
            double pab = nab / nt;

            double t1 = 0.0, t2 = 0.0;
            calc_dprime_rsq_kernel(pa, pb, pab, t1, t2);

            pab = pa * pb;
            dprime += pab * t1;
            rsq += pab * t2;
        }
    }
}

// Estimates haplotype frequencies via the EM algorithm
// AB, Ab, aB, ab, AaBb
void calc_hap_prob_EM(int n11, int n12, int n21, int n22, int ndh, double &p11, double &p12, double &p21, double &p22)
{
    static const int maxit = 1000;
    static const double tol = 1e-10;

    double n = n11 + n12 + n21 + n22 + ndh * 2;
    p11 = n11 / n;
    p12 = n12 / n;
    p21 = n21 / n;
    p22 = n22 / n;

    if (ndh == 0)
        return;

    auto cp11 = p11;
    auto cp12 = p12;
    auto cp21 = p21;
    auto cp22 = p22;

    auto h = ndh / n;
    auto x = h / 2;
    auto y = h - x;

    for (int i = 0; i < maxit; ++i) {
        p11 = cp11 + x;
        p12 = cp12 + y;
        p21 = cp21 + y;
        p22 = cp22 + x;
        auto z = h * p11 * p22 / (p11 * p22 + p12 * p21);
        if (std::fabs(x - z) < tol)
            break;
        x = z;
        y = h - x;
    }
}

void calc_dprime_rsq_dip2(const std::vector<allele_t> &x, const std::vector<allele_t> &y, double &dprime, double &rsq)
{
    auto n = x.size() / 2;

    int nt = 0;
    int f[3][3] = { { 0,0,0 }, { 0,0,0 }, { 0,0,0 } };

    for (size_t i = 0; i < n; ++i) {
        auto j = i*2, k = i*2+1;
        if (x[j] && x[k] && y[j] && y[k]) {
            auto a = x[j] + x[k] - 2;
            auto b = y[j] + y[k] - 2;
            ++f[a][b];
            ++nt;
        }
    }

    dprime = rsq = std::numeric_limits<double>::quiet_NaN();

    if (nt == 0)
        return;

    int n11 = f[0][0] * 2 + f[0][1] + f[1][0];
    int n12 = f[0][2] * 2 + f[0][1] + f[1][2];
    int n21 = f[2][0] * 2 + f[1][0] + f[2][1];
    int n22 = f[2][2] * 2 + f[2][1] + f[1][2];
    int ndh = f[1][1];

    double p11 = 0.0, p12 = 0.0, p21 = 0.0, p22 = 0.0;
    calc_hap_prob_EM(n11, n12, n21, n22, ndh, p11, p12, p21, p22);

    calc_dprime_rsq_kernel(p11 + p12, p11 + p21, p11, dprime, rsq);
}

// presuming that gametic phase is known
void calc_dprime_rsq_dip(const std::vector<allele_t> &x, const std::vector<allele_t> &y, double &dprime, double &rsq)
{
    auto n = x.size() / 2;

    auto xmax = * std::max_element(x.begin(), x.end());
    auto ymax = * std::max_element(y.begin(), y.end());

    dprime = rsq = std::numeric_limits<double>::quiet_NaN();

    if (xmax < 2 || ymax < 2)
        return;

    if (xmax == 2 && ymax == 2) {
        calc_dprime_rsq_dip2(x, y, dprime, rsq);
        return;
    }

    std::vector<size_t> idx;
    idx.reserve(n);
    for (size_t i = 0; i < n; ++i) {
        auto j = i*2, k = i*2+1;
        if (x[j] && x[k] && y[j] && y[k])
            idx.push_back(i);
    }

    if ( idx.empty() )
        return;

    dprime = rsq = 0.0;

    for (allele_t a = 1; a <= xmax; ++a) {
        for (allele_t b = 1; b <= ymax; ++b) {
            int na = 0, nb = 0, nab = 0;

            for (auto i : idx) {
                auto j = i*2, k = i*2+1;
                if (x[j] == a) {
                    ++na;
                    if (y[j] == b)
                        ++nab;
                }

                if (x[k] == a) {
                    ++na;
                    if (y[k] == b)
                        ++nab;
                }

                if (y[j] == b)
                    ++nb;

                if (y[k] == b)
                    ++nb;
            }

            if (na == 0 || nb == 0)
                continue;

            double nt = idx.size() * 2;
            auto pa = na / nt;
            auto pb = nb / nt;
            auto pab = nab / nt;

            double t1 = 0.0, t2 = 0.0;
            calc_dprime_rsq_kernel(pa, pb, pab, t1, t2);

            pab = pa * pb;
            dprime += pab * t1;
            rsq += pab * t2;
        }
    }
}

std::vector<std::string> read_string_list(const std::string &filename)
{
    std::vector<std::string> vs;

    std::ifstream ifs(filename);

    if ( ! ifs )
        std::cerr << "ERROR: can't open file for reading: " << filename << "\n";
    else
        std::copy(std::istream_iterator<std::string>(ifs), std::istream_iterator<std::string>(),
                  std::back_inserter(vs));

    return vs;
}

int ldstat_list(const Genotype &gt)
{
    auto loc = read_string_list(par.loc);

    std::sort(loc.begin(), loc.end());

    std::ofstream ofs(par.out + ".sub");
    if ( ! ofs ) {
        std::cerr << "ERROR: can't open file: " << par.out << ".list\n";
        return 1;
    }

    ofs << "Locus1\tChromosome1\tPosition1\tLocus2\tChromosome2\tPosition2\tDPrime\tRSquare\n";

    auto m = gt.loc.size();

    for (size_t i = 0; i < m; ++i) {
        if ( ! std::binary_search(loc.begin(), loc.end(), gt.loc[i]) )
            continue;

        std::cerr << "INFO: locus " << gt.loc[i] << "\n";

        for (size_t j = 0; j < m; ++j) {
            if (gt.loc[j] == gt.loc[i])
                continue;

            auto n1 = gt.allele[i].size();
            auto n2 = gt.allele[j].size();

            if (n1 < 2 || n2 < 2)
                continue;

            auto dprime = std::numeric_limits<double>::quiet_NaN();
            auto rsq = std::numeric_limits<double>::quiet_NaN();

            if (n1 == 2 && n2 == 2) {
                if (gt.ploidy == 1)
                    calc_dprime_rsq_hap2(gt.dat[i], gt.dat[j], dprime, rsq);
                else
                    calc_dprime_rsq_dip2(gt.dat[i], gt.dat[j], dprime, rsq);
            }
            else {
                if (gt.ploidy == 1)
                    calc_dprime_rsq_hap(gt.dat[i], gt.dat[j], dprime, rsq);
                else
                    calc_dprime_rsq_dip(gt.dat[i], gt.dat[j], dprime, rsq);
            }

            if ( ! std::isfinite(dprime) || ! std::isfinite(rsq) )
                continue;

            if (rsq < par.rsq)
                continue;

            ofs << gt.loc[i] << "\t" << gt.chr[i] << "\t" << gt.pos[i] << "\t"
                << gt.loc[j] << "\t" << gt.chr[j] << "\t" << gt.pos[j] << "\t"
                << dprime << "\t" << rsq << "\n";
        }
    }

    return 0;
}

} // namespace

int ldstat(int argc, char *argv[])
{
    std::cerr << "LDSTAT (Built on " __DATE__ " " __TIME__ ")\n";

    CmdLine cmd;

    cmd.add("--vcf", "VCF file", "");
    cmd.add("--loc", "locus list file", "");
    cmd.add("--loc-min-r2", "minimum LD (r2) threshold for locus list", "0.5");
    cmd.add("--out", "output file", "ldstat.out");
    cmd.add("--maxdist", "maximum inter-variant distance", "500000");

    cmd.parse(argc, argv);

    if (argc < 2) {
        cmd.show();
        return 1;
    }

    par.vcf = cmd.get("--vcf");
    par.loc = cmd.get("--loc");
    par.out = cmd.get("--out");
    par.rsq = std::stod(cmd.get("--loc-min-r2"));
    par.maxdist = std::stoi(cmd.get("--maxdist"));

    Genotype gt;

    std::cerr << "INFO: reading genotype file...\n";
    if (read_vcf(par.vcf, gt) != 0)
        return 2;
    std::cerr << "INFO: " << gt.ind.size() << " individuals, " << gt.loc.size() << " loci\n";

    if ( ! par.loc.empty() )
        return ldstat_list(gt);

    std::ofstream ofs1(par.out + ".all");
    if ( ! ofs1 ) {
        std::cerr << "ERROR: can't open file: " << par.out << ".all\n";
        return 1;
    }

    ofs1 << "Chromosome\tLocus1\tPosition1\tLocus2\tPosition2\tDistance\tDPrime\tRSquare\n";

    int n = par.maxdist / 1000;
    using namespace boost::accumulators;
    std::vector< accumulator_set<double, stats<tag::count, tag::mean> > > acc1(n), acc2(n);

    auto m = gt.loc.size();

    for (size_t i = 0; i < m; ++i) {
        for (size_t j = i + 1; j < m; ++j) {
            if (gt.chr[j] != gt.chr[i])
                continue;

            auto dist = std::abs(gt.pos[j] - gt.pos[i]);
            if (dist > par.maxdist)
                continue;

            auto n1 = gt.allele[i].size();
            auto n2 = gt.allele[j].size();

            if (n1 < 2 || n2 < 2)
                continue;

            auto dprime = std::numeric_limits<double>::quiet_NaN();
            auto rsq = std::numeric_limits<double>::quiet_NaN();

            if (n1 == 2 && n2 == 2) {
                if (gt.ploidy == 1)
                    calc_dprime_rsq_hap2(gt.dat[i], gt.dat[j], dprime, rsq);
                else
                    calc_dprime_rsq_dip2(gt.dat[i], gt.dat[j], dprime, rsq);
            }
            else {
                if (gt.ploidy == 1)
                    calc_dprime_rsq_hap(gt.dat[i], gt.dat[j], dprime, rsq);
                else
                    calc_dprime_rsq_dip(gt.dat[i], gt.dat[j], dprime, rsq);
            }

            if ( ! std::isfinite(dprime) || ! std::isfinite(rsq) )
                continue;

            int k = dist / 1000;
            if (dist % 1000 == 0)
                --k;

            acc1[k](dprime);
            acc2[k](rsq);

            ofs1 << gt.chr[i] << "\t"
                 << gt.loc[i] << "\t" << gt.pos[i] << "\t"
                 << gt.loc[j] << "\t" << gt.pos[j] << "\t"
                 << dist << "\t" << dprime << "\t" << rsq << "\n";
        }
    }

    std::ofstream ofs2(par.out + ".sum");
    if ( ! ofs2 ) {
        std::cerr << "ERROR: can't open file: " << par.out << ".sum\n";
        return 1;
    }

    ofs2 << "Group\tCount\tDPrime\tRSquare\n";

    for (int i = 0; i < n; ++i) {
        auto c = count(acc1[i]);
        ofs2 << (i + 1) * 1000 << "\t" << c << "\t";
        if (c != 0)
            ofs2 << mean(acc1[i]) << "\t" << mean(acc2[i]) << "\n";
        else
            ofs2 << "NA\tNA\n";
    }

    return 0;
}


#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <string>
#include <cstdlib>
#include <ctime>
#include <cmath>
#include <algorithm>
#include <sstream>
#include <Eigen/Eigenvalues>
#include <unsupported/Eigen/MatrixFunctions>
#include "BitUtility.h"
#include "StateGenerator.h"
#include "ScriptGenerator.h"
#include "Spectrum.h"
#include "NormCalculator.h"
#include "otoc.h"

using namespace std;

struct Options {
    int L = -1;
    int M = -1;
    f_type N = std::numeric_limits<f_type>::infinity();
    f_type beta = std::numeric_limits<f_type>::infinity();
    f_type bt = std::numeric_limits<f_type>::infinity();
    f_type tMin;
    f_type tMax;
    f_type tSteps;
    std::string W;
    std::string V;
    std::string Op;
    char ComputeNorm = 0;
    char ExcludeZeroStates = 0;
    char Parts = 1;
    f_type alpha = 1.0;
    f_type sigma = 0.0;
};

void TestBlockMatrixExp() {
    IOperator* op = IOperator::Create("XXXZ");
    int L = 11;
    int N = 13;
    BlockMatrix res(L);
    op->ToMatrixN(L, N, res);
    Stopwatch watch;
    BlockMatrix exp1 = res.exp();
    std::cout << "elapsed: " << watch.Elapsed() << std::endl;

    Matrix mat = res.matrix();
    watch.Start();
    Matrix exp2 = mat.exp();
    std::cout << "elapsed: " << watch.Elapsed() << std::endl;

    std::cout << "norm=" << exp2.norm() << std::endl;
    std::cout << "diff=" << (exp1.matrix() - exp2).norm() << std::endl;

    delete op;
}

void TestMatrixExp() {
    DilatationOperator ham;
    int L = 13;
    int M = 6;
    f_type N = 10;
    f_type beta = 0.9;
    int size = StateCollection::Inst()->StateNumber(L, M);
    Matrix res(size, size);
    ham.ToMatrixN(L, M, N, beta, res);

    Stopwatch watch;
/*    watch.Start();
    std::vector<var_t> v1;
    for (f_type t = 0.0; t <= 10; t += 0.5) {
        std::cout << "t=" << t << " ";
        var_t z(0, -t);
        Matrix exp1 = (res * z).exp();
        v1.push_back(exp1.trace());
    }

    std::cout << endl;
    std::cout << "v1=" << v1 << std::endl;

    std::cout << "time used:" << watch.Stop() << " seconds." << std::endl << std::endl;;*/

    watch.Start();
    Eigen::ComplexEigenSolver<Matrix> solver(res, true);
    Eigen::ArrayXcd diag = solver.eigenvalues().array();
//    Vector diag = solver.eigenvalues();
    Matrix V = solver.eigenvectors();
    Matrix Vinv = V.inverse();

//    std::cout << "diag=" << diag << std::endl;
    std::vector<var_t> v2;
    std::cout << "initial time= " << watch.Elapsed() << std::endl;

    for (f_type t = 0.0; t <= 10; t += 0.5) {
        std::cout << "t=" << t << " ";
        var_t z(0, -t);
        Matrix exp2 = V * (diag * z).exp().matrix().asDiagonal() * Vinv;
//        Matrix exp2 = V * (diag.array() * z).exp().matrix().asDiagonal() * Vinv;
//        std::cout << "(diag*z).exp()=" << (diag * z).exp() << std::endl;
        v2.push_back(exp2.trace());
    }
    std::cout << std::endl;
    std::cout << "v2=" << v2 << std::endl;

    std::cout << "time used:" << watch.Stop() << " seconds." << std::endl;
}

void TestOtoc() {
    int L = 4;
    int M = -1;
    f_type N = 17.0;
    f_type beta = 0.9;
    f_type b = 0.5;
    OtocParameters params = {"XZ", "ZX", L, M, N, beta, b, 1};
    std::cout << "test otoc" << std::endl;
    IOtoc* otoc = IOtoc::Create(params);
    otoc->Init();
    std::cout << "finished initialization..." << std::endl;
    for (int i = 0; i < 10; i++) {
        std::cout << i << ' ' << otoc->Compute((f_type)i) << std::endl;
    }
}

void TestForAnything() {
    Matrix m;
    std::cout << m.rows() << " " << m.cols() << std::endl;
}

void TestOp() {
    FourFieldOperator ffp("XZXZ");
    TraceState ts;
    ts.AddTrace(SingleTrace("ZXZX"));
    MixState ms;
    ffp.ApplyOn(ts, ms);
    std::cout << ms << std::endl;
}

void TestNormCalculator() {
    StateGenerator generator;
    generator.GenerateAllStates();
    generator.InitStateCollection(StateCollection::Inst());

    SingleTrace trace("ZZXX");
    TraceState ts;
    ts.AddTrace(trace);
    BruteForceCalculator calc;
    PowerSeries<i64> res = calc.Calculate(ts, ts);
    for (int i = 0; i <= res.HighestOrder(); i++) {
        std::cout << res.Coef(i) << " ";
    }
    std::cout << std::endl;
}

void TestNormCalculator2() {
    StateGenerator generator;
    generator.GenerateAllStates();
    generator.InitStateCollection(StateCollection::Inst());
    std::cout << "init done" << std::endl;
    /*    StateId id(0, 0);
    std::cout << StateCollection::Inst()->GetState(id) << std::endl;*/

    SingleTrace trace("ZZ");
    SingleTrace trace2("ZZXX");
    SingleTrace trace3("ZZXZZX");
    trace.Normalize();
    trace2.Normalize();
    trace3.Normalize();
    TraceState ts;
    ts.AddTrace(trace);
    ts.AddTrace(trace2);
    ts.QuickNormalize();
    TraceState ts2;
    ts2.AddTrace(trace3);
    RecursiveNormCalculator calc;
    PowerSeries<nType> res = calc.Calculate(ts, ts2);
    std::cout << "<" << trace << "^" << trace << ">=";
    std::cout << res << std::endl;

    BruteForceCalculator calc2;
    PowerSeries<nType> res2 = calc2.Calculate(ts, ts2);
    std::cout << "res=" << res2 << std::endl;
}

void ComputeOtoc(OtocParameters params, f_type tMin, f_type tMax, int steps, bool parts = false) {
    string rootFolder = "";
    ScriptGenerator gen(rootFolder);
    gen.SaveOtoc(params, tMin, tMax, steps, parts);
}

void GenerateHamMatrix(int L, int M)
{
    string rootFolder = "";
    ScriptGenerator gen(rootFolder);
    if (M >= 0 && M <= L) {
        gen.SaveHamMatrix(L, M);
        gen.SaveStates(L, M);
    }
    else {
        for (int i = 0; i <= L; i++) {
            gen.SaveHamMatrix(L, i);
            gen.SaveStates(L, i);
        }
    }
}

void GenerateOpMatrix(std::string opStr, int L, int M = -1)
{
    string rootFolder = "";
    ScriptGenerator gen(rootFolder);
    if (M >= 0) {
        gen.SaveOpMatrix(opStr, L, M);
    } else {
        IOperator* op = dynamic_cast<ISimpleOperator*>(IOperator::Create(opStr));
        if (!op->MagnonFixed()) {
            gen.SaveOpMatrix(opStr, L);
        }
        else {
            for (int m = 0; m <= L; m++) {
                gen.SaveOpMatrix(opStr, L, m);
            }
        }
        delete op;
    }
}

void GenerateNormMatrix(int L, int M, bool su)
{
    string rootFolder = "";
    ScriptGenerator gen(rootFolder);
    if (M < 0) {
        if (su) {
            gen.SaveNormToDataFile2(L);
        }
        else {
            gen.SaveNormToDataFile(L);
        }
        return;
    }
    if (su) {
        gen.SaveNormToDataFile2(L, M);
    }
    else {
        gen.SaveNormToDataFile(L, M);
    }
}

void ComputeSpectrum(int L, int M, int N, f_type beta)
{
    string rootFolder = "";
    ScriptGenerator gen(rootFolder);
    if (M >= 0) {
        gen.SaveSpectrum(L, M, N, beta);
    }
    else {
        for (int i = 0; i <= L; i++) {
            gen.SaveSpectrum(L, i, N, beta);
        }
    }
}

std::vector<std::string> optionDesc {
    "-L <integer number>", "Length of spin chain, 2 <= L <= " + ToString(MAX_BIT_TO_GENERATE),
    "-M <integer number>", "Number of magnons.",
    "-N <real number>", "Rank of SU(N) color group. In general it's integer but could be real.",
    "-b <real number>", "The beta parameter of beta-deformed N=4 SYM.",
    "-bt <real number>", "The reciprocal of temperature.",
    "-tmin <real number>", "The initial time of OTOC.",
    "-tmax <real number>", "The final time of OTOC.",
    "-tstep <integer number>", "Number of steps go from initial time to final time.",
    "-o <string>", "The operator to generate matrix for. The values can be NTr, NTrX, ",
    " ", "or any two-letter or four-letter string built out of X and Z.",
    " ", "NTr indicates the operator counts number of traces;",
    " ", "NTrX indicates the operator counts number of traces containing X;",
    " ", "XZ indicates the trace operator Tr(X d/dZ); XZXZ indicates the operator Tr(XZ d/dX d/dZ)",
    "-W <string>", "The W operator of OTOC. Its scope of values are the same as -o option.",
    "-V <string>", "The V operator of OTOC. Its scope of values are the same as -o option.",
    "-n", "The option only used in 'otoc' command. If not spefied, the norm matrix will be computed;",
    " ", "otherwise, the norm matrix will be loaded form local files generated by 'norm' command",
    "-nozero", "If the option is specified, zero energy states is excluded in C(t) computation.",
    "-alpha <real number>", "The alpha parameter of the OTOC regularization",
    "-sigma <real number>", "The sigma parameter of the OTOC regularization"
};

void DisplayOptions() {
    std::cout << "General Options:" << std::endl;
    for (size_t i = 0; i < optionDesc.size(); i += 2) {
        std::cout << "  " << std::left << std::setfill(' ') << std::setw(25) << optionDesc[i];
        std::cout << "\t\t" << optionDesc[i + 1] << std::endl;
    }
}

std::vector<std::string> exampleDesc {
    "Generate Hamiltonian matrix and trace states for L=4, M=2",                         // 1
    "Generate all Hamiltonian matrix and trace states for L=4",                          // 2
    "Generate norm matrix for L=4, M=2",                                //  3
    "Generate all norm matrix for L=4",                                 // 4
    "Compute energy spectrum for L=4, M=2, N=17, beta=0.9",             // 5
    "Generate matrix for operator counting traces containing X for L=4",
    "Compute OTOC for W=Tr(Xd/dZ), V=Tr(Zd/dX), L=4, N=5, beta=0.9, beta of temperature=0.5,\n \
\t\tand time run from 0.0 to 5.0 with step size 0.1. Note that to run the example one needs\n\
\t\tto run \"otoc norm - L 4\" first to generate all norm matrices for L = 4 and then copy the\n \
\t\tgenerated files to otocdata folder",
    "Compute OTOC for W=Tr(XZ d/dX d/dZ), V=Tr(ZX d/dX d/dZ), L=4, M=2, beta=0.9, beta of temperature=0.5,\n \
\t\tand N=Infinity. As -n option is specified, it will compute norm matrix rather than load it from local files.",
    "Compute C(t) for W=Tr(XX d/dX d/dZ), V=Tr(ZZ d/dZ d/dX), L=4, N=5. Zero energy states are excluded.",
    "Compute C(t) for W=Tr(XX d/dX d/dZ), V=Tr(ZZ d/dZ d/dX), L=4, N=5. With regularization parameters alpha=0.5 and sigma=0.25.",
    "Example of computing OTOC for composite operators. Currently, addition (+), substraction(-) , \n\
\t\tmultiplication(*), and exponentiate (exp), and normal ordered (~) are supported. No division (/) \n\
\t\tand space is allowd in the expression. Complex number should write in the form (real,imag).",
};

std::vector<std::string> examples {
    "otoc ham -L 4 -M 2",                       // 1
    "otoc ham -L 4",                            // 2
    "otoc norm -L 4 -M 2",                      // 3
    "otoc norm -L 4",                           // 4
    "otoc spec -L 4 -M 2 -N 17 -b 0.9",         // 5
    "otoc operator -o NTrX -L 4",                          // 6
    "otoc otoc -W XZ -V ZX -L 4 -N 5 -b 0.9 -bt 0.5 -tmin 0.0 -tmax 5.0 -tstep 50",
    "otoc otoc -W XZXZ -V XZXZ -L 4 -M 2 -b 0.9 -bt 0.5 -tmin 0.0 -tmax 5.0 -tstep 50 -n",
    "otoc otoc -W XXXZ -V ZZZX -L 4 -N 5 -b 0.9 -bt 0.5 -tmin 0.0 -tmax 5.0 -tstep 50 -nozero",
    "otoc otoc -W XXXZ -V ZZZX -alpha 0.5 -sigma 0.25 -L 4 -N 5 -b 0.9 -bt 0.5 -tmin 0.0 -tmax 5.0 -tstep 50",
    "otoc otoc -W -0.2*exp(0.2*XXXZ+0.3*XZ)+(0.3,-0.5) -V ~(-XZXZ+0.5*ZZZX-0.2*ZX) -L 4 -N 5 -b 0.9 -bt 0.5 -tmin 0.0 -tmax 5.0 -tstep 50",
};

void DisplayExample() {
    std::cout << "Examples:" << std::endl;
    for (int i = 0; i < examples.size(); i++) {
        std::cout << "  Example " << i + 1 << ":\t" << exampleDesc[i] << std::endl;
        std::cout << "\t\t\t" << examples[i] << std::endl << std::endl;
    }
}

std::vector<std::string> commands{ "ham", "norm", "spec", "operator", "otoc" };

void DisplayHelp() {
    std::cout << "Usage:\n  otoc <command> [options]" << std::endl << std::endl;
    std::cout << "Commands:" << std::endl;

    std::vector<std::string> commandDesc {
        "Generate and save Hamiltonian matrix and trace states.",
        "Compute and save spectrum of Hamiltonian.",
        "Generate and save norm matrix of trace states.",
        "Generate and save matrix of operators.",
        "Compute and save out-of-time-order correlator."
    };

    for (int i = 0; i < commands.size(); i++) {
        std::cout << "  " << std::left << std::setfill(' ') << std::setw(25) << commands[i];
        std::cout << "\t\t" << commandDesc[i] << std::endl;
    }

    std::cout << std::endl;
    DisplayOptions();

    std::cout << std::endl;
    DisplayExample();
}

std::vector<string> Tokenize(std::string str, char seperator) {
    std::vector<std::string> ret;
    stringstream ss(str);
    std::string token;
    while (getline(ss, token, seperator)) {
        ret.push_back(token);
    }

    return ret;
}

void InitStates(int maxL) {
    StateGenerator generator(maxL);
    generator.GenerateAllStates();
    generator.InitStateCollection(StateCollection::Inst());
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        DisplayHelp();
        return -1;
    }

    std::string command = argv[1];
    if (command == "test") {
        InitStates(16);
//        TestForAnything();
//        TestOtoc();
        //TestOp();
//        TestMatrixExp();
        TestBlockMatrixExp();
        return 0;
    }

    if (std::find(commands.begin(), commands.end(), command) == commands.end()) {
        DisplayHelp();
        return -1;
    }

    Options options;
    bool displayHelp = false;

    for (int i = 2; i < argc; i++) {
        string cmd = argv[i];
        if (cmd == "-L") {
            //
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.L = ToExpression(val, -1.0L);
        }
        else if (cmd == "-M") {
            //
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.M = ToExpression(val, -1);
        }
        else if (cmd == "-N") {
            //
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.N = ToExpression(val, -1.0);
        }
        else if (cmd == "-b") {
            //
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.beta = ToExpression(val, 1000000.0);
            beta_nsym = options.beta;
        }
        else if (cmd == "-bt") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.bt = ToExpression(val, options.bt);
            inverseTempeture = options.bt;
        }
        else if (cmd == "-tmax") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.tMax = ToExpression(val, options.tMax);
        }
        else if (cmd == "-tmin") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.tMin = ToExpression(val, options.tMin);
        }
        else if (cmd == "-tstep") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            string val = argv[i];
            options.tSteps = ToExpression(val, options.tSteps);
        } else if (cmd == "-o") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            options.Op = argv[i];
        }
        else if (cmd == "-W") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            options.W = argv[i];
        }
        else if (cmd == "-V") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            options.V = argv[i];
        }
        else if (cmd == "-alpha") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            options.alpha = ToExpression(argv[i], options.alpha);
        }
        else if (cmd == "-sigma") {
            i++;
            if (i >= argc) {
                displayHelp = true;
                break;
            }
            options.sigma = ToExpression(argv[i], options.sigma);;
        }
        else if (cmd == "-n") {
            options.ComputeNorm = 1;
        }
        else if (cmd == "-nozero") {
            options.ExcludeZeroStates = 1;
        }
//        else if (cmd == "-parts") {
//            options.Parts = 1;
 //       }
        else {
            std::cerr << "invalid option " << cmd << std::endl;
            return -1;
        }
    }

    if (displayHelp) {
        DisplayHelp();
        return -1;
    }

    if (options.L > MAX_BIT_TO_GENERATE || options.L <= 1 || options.M > options.L || options.M < -1) {
        displayHelp = true;
    }

    if (options.N < 1) {
        displayHelp = true;
    }

    if (displayHelp) {
        DisplayHelp();
        return -1;
    }

    Stopwatch watch;
    watch.Start();
    LOG("Initializing trace states...", Verbos);
    InitStates(options.L);

    try {
        if (command == commands[0]) {
            GenerateHamMatrix(options.L, options.M);
        }
        else if (command == commands[1]) {
            GenerateNormMatrix(options.L, options.M, true);
        }
        else if (command == commands[2]) {
            ComputeSpectrum(options.L, options.M, options.N, options.beta);
        }
        else if (command == commands[3]) {
            GenerateOpMatrix(options.Op, options.L, options.M);
        }
        else if (command == commands[4]) {
            OtocParameters params = { options.W, options.V, options.L, options.M, options.N,
                options.beta, options.bt, options.ComputeNorm, options.ExcludeZeroStates,
                options.alpha, options.sigma };
            ComputeOtoc(params, options.tMin, options.tMax, options.tSteps, options.Parts);
        }

        LOG("The task is completed in " << watch.Stop() << " seconds.", Info);
    }
    catch (OtocException& ex){
        LOG(ex.what(), Error);
        LOG("Elapsed time: " << watch.Stop() << " seconds.", Info);
    }
}


/*
 * StokesM2L.cpp
 *
 *  Created on: Oct 12, 2016
 *      Author: wyan
 */

#include "SVD_pvfmm.hpp"

#include <Eigen/Dense>

#include <chrono>
#include <iomanip>
#include <iostream>

#define DIRECTLAYER 2
#define PI314 3.1415926535897932384626433

namespace Stokes3D3D {

inline double ERFC(double x) { return std::erfc(x); }
inline double ERF(double x) { return std::erf(x); }

/*
 * def AEW(xi,rvec):
 r=np.sqrt(rvec.dot(rvec))
 A = 2*(xi*np.exp(-(xi**2)*(r**2))/(np.sqrt(np.pi)*r**2)+ss.erfc(xi*r)/(2*r**3))
 *(r*r*np.identity(3)+np.outer(rvec,rvec)) -
 4*xi/np.sqrt(np.pi)*np.exp(-(xi**2)*(r**2))*np.identity(3)
 return A
 *
 * */
inline Eigen::Matrix3d AEW(const double xi, const Eigen::Vector3d &rvec) {
    const double r = rvec.norm();
    Eigen::Matrix3d A = 2 * (xi * exp(-(xi * xi) * (r * r)) / (sqrt(PI314) * r * r) + erfc(xi * r) / (2 * r * r * r)) *
                            (r * r * Eigen::Matrix3d::Identity() + (rvec * rvec.transpose())) -
                        4 * xi / sqrt(PI314) * exp(-(xi * xi) * (r * r)) * Eigen::Matrix3d::Identity();
    return A;
}

/*
 *
 def BEW(xi,kvec):
 k=np.sqrt(kvec.dot(kvec))
 B =
 8*np.pi*(1+k*k/(4*(xi**2)))*((k**2)*np.identity(3)-np.outer(kvec,kvec))/(k**4)
 return B*np.exp(-k**2/(4*xi**2))
 *
 * */
inline Eigen::Matrix3d BEW(const double xi, const Eigen::Vector3d &kvec) {
    const double k = kvec.norm();
    Eigen::Matrix3d B = 8 * PI314 * (1 + k * k / (4 * (xi * xi))) *
                        ((k * k) * Eigen::Matrix3d::Identity() - (kvec * kvec.transpose())) / (k * k * k * k);
    B *= exp(-k * k / (4 * xi * xi));
    return B;
}

/*
 * def stokes3DEwald(rvec,force):
 xi = 2
 r=np.sqrt(rvec.dot(rvec))
 real = 0
 N=4
 for i in range(-N,N+1):
 for j in range(-N,N+1):
 for k in range(-N,N+1):
 real = real + AEW(xi,rvec+1.0*np.array([i,j,k])).dot(force)
 wave = 0
 N=4
 for i in range(-N,N+1):
 for j in range(-N,N+1):
 for k in range(-N,N+1):
 kvec=2*np.pi*np.array([i,j,k]) # L = 1
 if(i==0 and j==0 and k==0):
 continue
 else:
 wave = wave + BEW(xi,kvec).dot(force)*np.exp(-complex(0,1)*kvec.dot(rvec))

 return (np.real(wave)+real)

 * */
inline void GkernelEwald(const Eigen::Vector3d &rvec, Eigen::Matrix3d &Gsum) {
    const double xi = 2;
    const double r = rvec.norm();
    Eigen::Matrix3d real = Eigen::Matrix3d::Zero();
    const int N = 5;
    if (r < 1e-14) {
        auto Gself = -4 * xi / sqrt(PI314) * Eigen::Matrix3d::Identity(); // the self term
        for (int i = -N; i < N + 1; i++) {
            for (int j = -N; j < N + 1; j++) {
                for (int k = -N; k < N + 1; k++) {
                    if (i == 0 && j == 0 && k == 0) {
                        continue;
                    }
                    real = real + AEW(xi, rvec + Eigen::Vector3d(i, j, k));
                }
            }
        }
        real += Gself;
    } else {
        for (int i = -N; i < N + 1; i++) {
            for (int j = -N; j < N + 1; j++) {
                for (int k = -N; k < N + 1; k++) {
                    real = real + AEW(xi, rvec + Eigen::Vector3d(i, j, k));
                }
            }
        }
    }
    Eigen::Matrix3d wave = Eigen::Matrix3d::Zero();

    for (int i = -N; i < N + 1; i++) {
        for (int j = -N; j < N + 1; j++) {
            for (int k = -N; k < N + 1; k++) {
                Eigen::Vector3d kvec(2 * PI314 * i, 2 * PI314 * j, 2 * PI314 * k);
                if (i == 0 and j == 0 and k == 0) {
                    continue;
                } else {
                    wave = wave + BEW(xi, kvec) * cos(kvec.dot(rvec));
                }
            }
        }
    }
    Gsum = real + wave;
}

inline void Gkernel(const Eigen::Vector3d &target, const Eigen::Vector3d &source, Eigen::Matrix3d &answer) {
    auto rst = target - source;
    double rnorm = rst.norm();
    if (rnorm < 1e-14) {
        answer = Eigen::Matrix3d::Zero();
        return;
    }
    auto part2 = rst * rst.transpose() / (rnorm * rnorm * rnorm);
    auto part1 = Eigen::Matrix3d::Identity() / rnorm;
    answer = part1 + part2;
}

/*
 *
 def stokes3DM2L(rvec,force):
 uEwald=stokes3DEwald(rvec,force)
 uNB=0
 N=3
 for i in range(-N,N+1):
 for j in range(-N,N+1):
 for k in range(-N,N+1):
 uNB=uNB+Gkernel(rvec-np.array([i,j,k])).dot(force)
 return uEwald-uNB
 * */
// Out of Layer 1
inline void GkernelEwaldO1(const Eigen::Vector3d &rvec, Eigen::Matrix3d &GsumO1) {
    Eigen::Matrix3d Gfree = Eigen::Matrix3d::Zero();
    GkernelEwald(rvec, GsumO1);
    const int N = DIRECTLAYER;
    for (int i = -N; i < N + 1; i++) {
        for (int j = -N; j < N + 1; j++) {
            for (int k = -N; k < N + 1; k++) {
                Gkernel(rvec, Eigen::Vector3d(i, j, k), Gfree);
                GsumO1 -= Gfree;
            }
        }
    }
}

/**
 * \brief Returns the coordinates of points on the surface of a cube.
 * \param[in] p Number of points on an edge of the cube is (n+1)
 * \param[in] c Coordinates to the centre of the cube (3D array).
 * \param[in] alpha Scaling factor for the size of the cube.
 * \param[in] depth Depth of the cube in the octree.
 * \return Vector with coordinates of points on the surface of the cube in the
 * format [x0 y0 z0 x1 y1 z1 .... ].
 */

template <class Real_t>
std::vector<Real_t> surface(int p, Real_t *c, Real_t alpha, int depth) {
    size_t n_ = (6 * (p - 1) * (p - 1) + 2); // Total number of points.

    std::vector<Real_t> coord(n_ * 3);
    coord[0] = coord[1] = coord[2] = -1.0;
    size_t cnt = 1;
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = -1.0;
            coord[cnt * 3 + 1] = (2.0 * (i + 1) - p + 1) / (p - 1);
            coord[cnt * 3 + 2] = (2.0 * j - p + 1) / (p - 1);
            cnt++;
        }
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = (2.0 * i - p + 1) / (p - 1);
            coord[cnt * 3 + 1] = -1.0;
            coord[cnt * 3 + 2] = (2.0 * (j + 1) - p + 1) / (p - 1);
            cnt++;
        }
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = (2.0 * (i + 1) - p + 1) / (p - 1);
            coord[cnt * 3 + 1] = (2.0 * j - p + 1) / (p - 1);
            coord[cnt * 3 + 2] = -1.0;
            cnt++;
        }
    for (size_t i = 0; i < (n_ / 2) * 3; i++)
        coord[cnt * 3 + i] = -coord[i];

    Real_t r = 0.5 * pow(0.5, depth);
    Real_t b = alpha * r;
    for (size_t i = 0; i < n_; i++) {
        coord[i * 3 + 0] = (coord[i * 3 + 0] + 1.0) * b + c[0];
        coord[i * 3 + 1] = (coord[i * 3 + 1] + 1.0) * b + c[1];
        coord[i * 3 + 2] = (coord[i * 3 + 2] + 1.0) * b + c[2];
    }
    return coord;
}

int main(int argc, char **argv) {
    Eigen::initParallel();
    Eigen::setNbThreads(1);
    std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();

    const int pEquiv = atoi(argv[1]); // (8-1)^2*6 + 2 points
    const int pCheck = atoi(argv[1]);
    const double scaleEquiv = 1.05;
    const double scaleCheck = 2.95;
    const double pCenterEquiv[3] = {-(scaleEquiv - 1) / 2, -(scaleEquiv - 1) / 2, -(scaleEquiv - 1) / 2};
    const double pCenterCheck[3] = {-(scaleCheck - 1) / 2, -(scaleCheck - 1) / 2, -(scaleCheck - 1) / 2};

    const double scaleLEquiv = 1.05;
    const double scaleLCheck = 2.95;
    const double pCenterLEquiv[3] = {-(scaleLEquiv - 1) / 2, -(scaleLEquiv - 1) / 2, -(scaleLEquiv - 1) / 2};
    const double pCenterLCheck[3] = {-(scaleLCheck - 1) / 2, -(scaleLCheck - 1) / 2, -(scaleLCheck - 1) / 2};

    auto pointMEquiv = surface(pEquiv, (double *)&(pCenterEquiv[0]), scaleEquiv,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0
    auto pointMCheck = surface(pCheck, (double *)&(pCenterCheck[0]), scaleCheck,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0

    auto pointLEquiv = surface(pEquiv, (double *)&(pCenterLCheck[0]), scaleLCheck,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0
    auto pointLCheck = surface(pCheck, (double *)&(pCenterLEquiv[0]), scaleLEquiv,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0

    //	for (int i = 0; i < pointLEquiv.size() / 3; i++) {
    //		std::cout << pointLEquiv[3 * i] << " " << pointLEquiv[3 * i + 1]
    //<< " " << pointLEquiv[3 * i + 2] << " "
    //				<< std::endl;
    //	}
    //
    //	for (int i = 0; i < pointLCheck.size() / 3; i++) {
    //		std::cout << pointLCheck[3 * i] << " " << pointLCheck[3 * i + 1]
    //<< " " << pointLCheck[3 * i + 2] << " "
    //				<< std::endl;
    //	}

    // const int imageN = 100; // images to sum
    // calculate the operator M2L with least square
    const int equivN = pointMEquiv.size() / 3;
    const int checkN = pointLCheck.size() / 3;
    Eigen::MatrixXd M2L(3 * equivN, 3 * equivN);
    Eigen::MatrixXd A(3 * checkN, 3 * equivN);
#pragma omp parallel for
    for (int k = 0; k < checkN; k++) {
        Eigen::Matrix3d G = Eigen::Matrix3d::Zero();
        Eigen::Vector3d Cpoint(pointLCheck[3 * k], pointLCheck[3 * k + 1], pointLCheck[3 * k + 2]);
        for (int l = 0; l < equivN; l++) {
            const Eigen::Vector3d Lpoint(pointLEquiv[3 * l], pointLEquiv[3 * l + 1], pointLEquiv[3 * l + 2]);
            Gkernel(Cpoint, Lpoint, G);
            A.block<3, 3>(3 * k, 3 * l) = G;
        }
    }
    Eigen::MatrixXd ApinvU(A.cols(), A.rows());
    Eigen::MatrixXd ApinvVT(A.cols(), A.rows());
    pinv(A, ApinvU, ApinvVT);

#pragma omp parallel for
    for (int i = 0; i < equivN; i++) {
        const Eigen::Vector3d Mpoint(pointMEquiv[3 * i], pointMEquiv[3 * i + 1], pointMEquiv[3 * i + 2]);
        //		std::cout<<"debug:"<<Mpoint<<std::endl;
        // assemble linear system
        Eigen::MatrixXd f(3 * checkN, 3);
        for (int k = 0; k < checkN; k++) {
            Eigen::Matrix3d temp = Eigen::Matrix3d::Zero();
            Eigen::Vector3d Cpoint(pointLCheck[3 * k], pointLCheck[3 * k + 1], pointLCheck[3 * k + 2]);
            //			std::cout<<"debug:"<<k<<std::endl;
            // sum the images
            // use 3D Ewald subtract the first layer
            GkernelEwaldO1(Cpoint - Mpoint, temp);
            f.block<3, 3>(3 * k, 0) = temp;
        }

        M2L.block(0, 3 * i, 3 * equivN, 3) = (ApinvU.transpose() * (ApinvVT.transpose() * f));
    }
    std::chrono::high_resolution_clock::time_point t2 = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::microseconds>(t2 - t1).count();

    // dump M2L
    for (int i = 0; i < 3 * equivN; i++) {
        for (int j = 0; j < 3 * equivN; j++) {
            std::cout << i << " " << j << " " << std::scientific << std::setprecision(18) << M2L(i, j) << std::endl;
        }
    }

    std::cout << "Precomputing time:" << duration / 1e6 << std::endl;

    /*
     * pointForce=[(np.array([1.0,0,0]),np.array([0.1,0.55,0.2]))
     ,(np.array([-1.0,1.0,1.0]),np.array([0.5,0.1,0.3]))
     ,(np.array([0.0,0.0,-1.0]),np.array([0.8,0.5,0.7]))]
     * */
    std::vector<Eigen::Vector3d, Eigen::aligned_allocator<Eigen::Vector3d>> forcePoint(3);
    std::vector<Eigen::Vector3d, Eigen::aligned_allocator<Eigen::Vector3d>> forceValue(3);
    forcePoint[0] = Eigen::Vector3d(0.1, 0.55, 0.2);
    forceValue[0] = Eigen::Vector3d(1, 0, 0);
    forcePoint[1] = Eigen::Vector3d(0.5, 0.1, 0.3);
    forceValue[1] = Eigen::Vector3d(-1, 1, 1);
    forcePoint[2] = Eigen::Vector3d(0.8, 0.5, 0.7);
    forceValue[2] = Eigen::Vector3d(0, 0, -1);

    // solve M
    A.resize(3 * checkN, 3 * equivN);
    ApinvU.resize(A.cols(), A.rows());
    ApinvVT.resize(A.cols(), A.rows());
    Eigen::VectorXd f(3 * checkN);
    for (int k = 0; k < checkN; k++) {
        Eigen::Vector3d temp = Eigen::Vector3d::Zero();
        Eigen::Matrix3d G = Eigen::Matrix3d::Zero();
        Eigen::Vector3d Cpoint(pointMCheck[3 * k], pointMCheck[3 * k + 1], pointMCheck[3 * k + 2]);
        for (size_t p = 0; p < forcePoint.size(); p++) {
            Gkernel(Cpoint, forcePoint[p], G);
            temp = temp + G * (forceValue[p]);
        }
        f.block<3, 1>(3 * k, 0) = temp;
        for (int l = 0; l < equivN; l++) {
            Eigen::Vector3d Mpoint(pointMEquiv[3 * l], pointMEquiv[3 * l + 1], pointMEquiv[3 * l + 2]);
            Gkernel(Cpoint, Mpoint, G);
            A.block<3, 3>(3 * k, 3 * l) = G;
        }
    }
    pinv(A, ApinvU, ApinvVT);
    Eigen::VectorXd Msource = (ApinvU.transpose() * (ApinvVT.transpose() * f));
    // impose net charge equal
    double fx = 0, fy = 0, fz = 0;
    for (int i = 0; i < equivN; i++) {
        fx += Msource[3 * i];
        fy += Msource[3 * i + 1];
        fz += Msource[3 * i + 2];
    }
    std::cout << "fx svd before correction: " << fx << std::endl;
    std::cout << "fy svd before correction: " << fy << std::endl;
    std::cout << "fz svd before correction: " << fz << std::endl;
    double fnetx = 0;
    double fnety = 0;
    double fnetz = 0;
    for (size_t p = 0; p < forcePoint.size(); p++) {
        fnetx += (forceValue[p][0]);
        fnety += (forceValue[p][1]);
        fnetz += (forceValue[p][2]);
    }
    /*
     * fx=(fx-fnet[0])/len(MPoints)
     fy=(fy-fnet[1])/len(MPoints)
     fz=(fz-fnet[2])/len(MPoints)
     * */
    fx = (fx - fnetx) / equivN;
    fy = (fy - fnety) / equivN;
    fz = (fz - fnetz) / equivN;
    for (int i = 0; i < equivN; i++) {
        Msource[3 * i] -= fx;
        Msource[3 * i + 1] -= fy;
        Msource[3 * i + 2] -= fz;
    }
    std::cout << "Msource: " << Msource << std::endl;

    std::vector<Eigen::Vector3d, Eigen::aligned_allocator<Eigen::Vector3d>> forcePointExt(0);
    std::vector<Eigen::Vector3d, Eigen::aligned_allocator<Eigen::Vector3d>> forceValueExt(0);
    for (size_t p = 0; p < forcePoint.size(); p++) {
        for (int i = -DIRECTLAYER; i < DIRECTLAYER + 1; i++) {
            for (int j = -DIRECTLAYER; j < DIRECTLAYER + 1; j++) {
                for (int k = -DIRECTLAYER; k < DIRECTLAYER + 1; k++) {
                    forcePointExt.push_back(Eigen::Vector3d(i, j, k) + forcePoint[p]);
                    forceValueExt.push_back(forceValue[p]);
                }
            }
        }
    }

    Eigen::VectorXd M2Lsource = M2L * (Msource);

    Eigen::Vector3d samplePoint(0.5, 0.5, 0.5);
    Eigen::Vector3d Usample(0, 0, 0);
    Eigen::Vector3d UsampleSP(0, 0, 0);
    Eigen::Matrix3d G;
    for (size_t p = 0; p < forcePointExt.size(); p++) {
        Gkernel(samplePoint, forcePointExt[p], G);
        Usample = Usample + G * (forceValueExt[p]);
    }
    std::cout << "Usample Direct:" << Usample << std::endl;
    for (int p = 0; p < equivN; p++) {
        Eigen::Vector3d Lpoint(pointLEquiv[3 * p], pointLEquiv[3 * p + 1], pointLEquiv[3 * p + 2]);
        Eigen::Vector3d Fpoint(M2Lsource[3 * p], M2Lsource[3 * p + 1], M2Lsource[3 * p + 2]);
        Gkernel(samplePoint, Lpoint, G);
        UsampleSP = UsampleSP + G * (Fpoint);
    }

    std::cout << "Usample M2L:" << UsampleSP << std::endl;
    std::cout << "Usample M2L total:" << UsampleSP + Usample << std::endl;

    Eigen::Vector3d UsampleDirect = 0 * Usample;
    for (size_t p = 0; p < forcePoint.size(); p++) {
        GkernelEwald(samplePoint - forcePoint[p], G);
        UsampleDirect += G * (forceValue[p]);
    }
    std::cout << "Usample Ewald:" << UsampleDirect << std::endl;

    std::cout << "error" << UsampleSP + Usample - UsampleDirect << std::endl;

    return 0;
}

} // namespace Stokes3D3D

#undef PI314
#undef DIRECTLAYER


//
// OpenTissue, A toolbox for physical based simulation and animation.
// Copyright (C) 2007 Department of Computer Science, University of Copenhagen
//
#include <OpenTissue/configuration.h>

#include <OpenTissue/core/math/math_basic_types.h>
#include <OpenTissue/core/geometry/geometry_compute_volume_length_quality_measure.h>

#define BOOST_AUTO_TEST_MAIN
#include <OpenTissue/utility/utility_push_boost_filter.h>
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>
#include <OpenTissue/utility/utility_pop_boost_filter.h>

BOOST_AUTO_TEST_SUITE(opentissue_geometry_util_compute_volume_length_quality_measure);

  BOOST_AUTO_TEST_CASE(case_by_case_testing)
  {
    using std::sqrt;

    typedef OpenTissue::math::BasicMathTypes<double, size_t> math_types;
    typedef math_types::vector3_type                         vector3_type;
    typedef math_types::real_type                            real_type;

    real_type tol = 0.0001;

    {
      // semi regular
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 0, 2);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;

      BOOST_CHECK_CLOSE( quality, 0.975423, tol);
    }
    {
      // close to flat
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 0, 0.01);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0465344, tol);
    }
    {
      // flat but still a simplex
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 0, 0);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0, tol);
    }
    {
      // flat but no longer a simplex - quad-like
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 1, 1, 0);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0, tol);
    }
    {
      // inverted
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 0, -1);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.865855, tol);
    }
    {
      // degenerate 2 nodes
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 1,-1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 1, 0);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0, tol);
    }
    {
      // degenerate 3 nodes
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 0, 1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 1, 0);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0, tol);
    }
    {
      // degenerate 4 nodes
      vector3_type p0(-1,-1, 0);
      vector3_type p1( 0, 1, 0);
      vector3_type p2( 0, 1, 0);
      vector3_type p3( 0, 1, 0);
      real_type quality = OpenTissue::geometry::compute_volume_length_quality_measure(p0,p1,p2,p3);
      std::cout << quality << std::endl;
      BOOST_CHECK_CLOSE( quality, 0.0, tol);
    }

  }

BOOST_AUTO_TEST_SUITE_END();


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

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

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

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

#include <geometry_test_common.hpp>

#include <boost/concept_check.hpp>

#include <boost/geometry/extensions/gis/geographic/strategies/vincenty.hpp>

#include <boost/geometry/strategies/strategies.hpp>
#include <boost/geometry/algorithms/assign.hpp>
#include <boost/geometry/geometries/point.hpp>
#include <test_common/test_point.hpp>

#ifdef HAVE_TTMATH
#  include <boost/geometry/extensions/contrib/ttmath_stub.hpp>
#endif

template <typename P1, typename P2>
void test_vincenty(double lon1, double lat1, double lon2, double lat2, double expected_km)
{
    // Set radius type, but for integer coordinates we want to have floating point radius type
    typedef typename bg::promote_floating_point
        <
            typename bg::coordinate_type<P1>::type
        >::type rtype;

    typedef bg::strategy::distance::vincenty<rtype> vincenty_type;

    BOOST_CONCEPT_ASSERT(
        (
            bg::concept::PointDistanceStrategy<vincenty_type, P1, P2>)
        );

    vincenty_type vincenty;
    typedef typename bg::strategy::distance::services::return_type<vincenty_type, P1, P2>::type return_type;

    P1 p1, p2;

    bg::assign_values(p1, lon1, lat1);
    bg::assign_values(p2, lon2, lat2);

    BOOST_CHECK_CLOSE(vincenty.apply(p1, p2), return_type(1000.0) * return_type(expected_km), 0.001);
}

template <typename P1, typename P2>
void test_all()
{
    test_vincenty<P1, P2>(0, 89, 1, 80, 1005.1535769); // sub-polar
    test_vincenty<P1, P2>(4, 52, 4, 52, 0.0); // no point difference
    test_vincenty<P1, P2>(4, 52, 3, 40, 1336.039890); // normal case
}

template <typename P>
void test_all()
{
    test_all<P, P>();
}

int test_main(int, char* [])
{

    //test_all<float[2]>();
    //test_all<double[2]>();
    test_all<bg::model::point<int, 2, bg::cs::geographic<bg::degree> > >();
    test_all<bg::model::point<float, 2, bg::cs::geographic<bg::degree> > >();
    test_all<bg::model::point<double, 2, bg::cs::geographic<bg::degree> > >();

#if defined(HAVE_TTMATH)
    test_all<bg::model::point<ttmath::Big<1,4>, 2, bg::cs::geographic<bg::degree> > >();
    test_all<bg::model::point<ttmath_big, 2, bg::cs::geographic<bg::degree> > >();
#endif

    return 0;
}


//  Copyright (c) 2007 John Maddock
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
// Computes test data for the various bessel functions using
// archived - deliberately naive - version of the code.
// We'll rely on the high precision of boost::math::ntl::RR to get us out of
// trouble and not worry about how long the calculations take.
// This provides a reasonably independent set of test data to
// compare against newly added asymptotic expansions etc.
//
#include <fstream>

#include <boost/math/tools/test_data.hpp>
#include "ntl_rr_lanczos.hpp"

#include <boost/math/special_functions/bessel.hpp>

using namespace boost::math::tools;
using namespace boost::math;
using namespace boost::math::detail;
using namespace std;

// Compute J(v, x) and Y(v, x) simultaneously by Steed's method, see
// Barnett et al, Computer Physics Communications, vol 8, 377 (1974)
template <typename T>
int bessel_jy_bare(T v, T x, T* J, T* Y, int kind = need_j|need_y)
{
    // Jv1 = J_(v+1), Yv1 = Y_(v+1), fv = J_(v+1) / J_v
    // Ju1 = J_(u+1), Yu1 = Y_(u+1), fu = J_(u+1) / J_u
    T u, Jv, Ju, Yv, Yv1, Yu, Yu1, fv, fu;
    T W, p, q, gamma, current, prev, next;
    bool reflect = false;
    int n, k, s;

    using namespace std;
    using namespace boost::math::tools;
    using namespace boost::math::constants;

    if (v < 0)
    {
        reflect = true;
        v = -v;                             // v is non-negative from here
        kind = need_j|need_y;               // need both for reflection formula
    }
    n = real_cast<int>(v + 0.5L);
    u = v - n;                              // -1/2 <= u < 1/2

    if (x < 0)
    {
       *J = *Y = policies::raise_domain_error<T>("",
          "Real argument x=%1% must be non-negative, complex number result not supported", x, policies::policy<>());
        return 1;
    }
    if (x == 0)
    {
       *J = *Y = policies::raise_overflow_error<T>(
          "", 0, policies::policy<>());
       return 1;
    }

    // x is positive until reflection
    W = T(2) / (x * pi<T>());               // Wronskian
    if (x <= 2)                           // x in (0, 2]
    {
       if(temme_jy(u, x, &Yu, &Yu1, policies::policy<>()))             // Temme series
        {
           // domain error:
           *J = *Y = Yu;
           return 1;
        }
        prev = Yu;
        current = Yu1;
        for (k = 1; k <= n; k++)            // forward recurrence for Y
        {
            next = 2 * (u + k) * current / x - prev;
            prev = current;
            current = next;
        }
        Yv = prev;
        Yv1 = current;
        CF1_jy(v, x, &fv, &s, policies::policy<>());                 // continued fraction CF1
        Jv = W / (Yv * fv - Yv1);           // Wronskian relation
    }
    else                                    // x in (2, \infty)
    {
        // Get Y(u, x):
        CF1_jy(v, x, &fv, &s, policies::policy<>());
        // tiny initial value to prevent overflow
        T init = sqrt(tools::min_value<T>());
        prev = fv * s * init;
        current = s * init;
        for (k = n; k > 0; k--)             // backward recurrence for J
        {
            next = 2 * (u + k) * current / x - prev;
            prev = current;
            current = next;
        }
        T ratio = (s * init) / current;     // scaling ratio
        // can also call CF1() to get fu, not much difference in precision
        fu = prev / current;
        CF2_jy(u, x, &p, &q, policies::policy<>());                  // continued fraction CF2
        T t = u / x - fu;                   // t = J'/J
        gamma = (p - t) / q;
        Ju = sign(current) * sqrt(W / (q + gamma * (p - t)));

        Jv = Ju * ratio;                    // normalization

        Yu = gamma * Ju;
        Yu1 = Yu * (u/x - p - q/gamma);

        // compute Y:
        prev = Yu;
        current = Yu1;
        for (k = 1; k <= n; k++)            // forward recurrence for Y
        {
            next = 2 * (u + k) * current / x - prev;
            prev = current;
            current = next;
        }
        Yv = prev;
    }

    if (reflect)
    {
        T z = (u + n % 2) * pi<T>();
        *J = cos(z) * Jv - sin(z) * Yv;     // reflection formula
        *Y = sin(z) * Jv + cos(z) * Yv;
    }
    else
    {
        *J = Jv;
        *Y = Yv;
    }

    return 0;
}

int progress = 0;

template <class T>
T cyl_bessel_j_bare(T v, T x)
{
   T j, y;
   bessel_jy_bare(v, x, &j, &y);

   std::cout << progress++ << ":   J(" << v << ", " << x << ") = " << j << std::endl;

   if(fabs(j) > 1e30)
      throw std::domain_error("");

   return j;
}

template <class T>
T cyl_bessel_i_bare(T v, T x)
{
   using namespace std;
   if(x < 0)
   {
      // better have integer v:
      if(floor(v) == v)
      {
         T r = cyl_bessel_i_bare(v, -x);
         if(tools::real_cast<int>(v) & 1)
            r = -r;
         return r;
      }
      else
         return policies::raise_domain_error<T>(
            "",
            "Got x = %1%, but we need x >= 0", x, policies::policy<>());
   }
   if(x == 0)
   {
      return (v == 0) ? 1 : 0;
   }
   T I, K;
   boost::math::detail::bessel_ik(v, x, &I, &K, 0xffff, policies::policy<>());

   std::cout << progress++ << ":   I(" << v << ", " << x << ") = " << I << std::endl;

   if(fabs(I) > 1e30)
      throw std::domain_error("");

   return I;
}

template <class T>
T cyl_bessel_k_bare(T v, T x)
{
   using namespace std;
   if(x < 0)
   {
      return policies::raise_domain_error<T>(
         "",
         "Got x = %1%, but we need x > 0", x, policies::policy<>());
   }
   if(x == 0)
   {
      return (v == 0) ? policies::raise_overflow_error<T>("", 0, policies::policy<>())
         : policies::raise_domain_error<T>(
         "",
         "Got x = %1%, but we need x > 0", x, policies::policy<>());
   }
   T I, K;
   bessel_ik(v, x, &I, &K, 0xFFFF, policies::policy<>());

   std::cout << progress++ << ":   K(" << v << ", " << x << ") = " << K << std::endl;

   if(fabs(K) > 1e30)
      throw std::domain_error("");

   return K;
}

template <class T>
T cyl_neumann_bare(T v, T x)
{
   T j, y;
   bessel_jy(v, x, &j, &y, 0xFFFF, policies::policy<>());

   std::cout << progress++ << ":   Y(" << v << ", " << x << ") = " << y << std::endl;

   if(fabs(y) > 1e30)
      throw std::domain_error("");

   return y;
}

template <class T>
T sph_bessel_j_bare(T v, T x)
{
   std::cout << progress++ << ":   j(" << v << ", " << x << ") = ";
   if((v < 0) || (floor(v) != v))
      throw std::domain_error("");
   T r = sqrt(constants::pi<T>() / (2 * x)) * cyl_bessel_j_bare(v+0.5, x);
   std::cout << r << std::endl;
   return r;
}

template <class T>
T sph_bessel_y_bare(T v, T x)
{
   std::cout << progress++ << ":   y(" << v << ", " << x << ") = ";
   if((v < 0) || (floor(v) != v))
      throw std::domain_error("");
   T r = sqrt(constants::pi<T>() / (2 * x)) * cyl_neumann_bare(v+0.5, x);
   std::cout << r << std::endl;
   return r;
}

enum
{
   func_J = 0,
   func_Y,
   func_I,
   func_K,
   func_j,
   func_y
};

int main(int argc, char* argv[])
{
   std::cout << std::setprecision(17) << std::scientific;
   std::cout << sph_bessel_j_bare(0., 0.1185395751953125e4) << std::endl;
   std::cout << sph_bessel_j_bare(22., 0.6540834903717041015625) << std::endl;

   parameter_info<boost::math::ntl::RR> arg1, arg2;
   test_data<boost::math::ntl::RR> data;

   boost::math::ntl::RR::SetPrecision(1000);
   boost::math::ntl::RR::SetOutputPrecision(40);

   int functype = 0;
   std::string letter = "J";

   if(argc == 2)
   {
      if(std::strcmp(argv[1], "--Y") == 0)
      {
         functype = func_Y;
         letter = "Y";
      }
      else if(std::strcmp(argv[1], "--I") == 0)
      {
         functype = func_I;
         letter = "I";
      }
      else if(std::strcmp(argv[1], "--K") == 0)
      {
         functype = func_K;
         letter = "K";
      }
      else if(std::strcmp(argv[1], "--j") == 0)
      {
         functype = func_j;
         letter = "j";
      }
      else if(std::strcmp(argv[1], "--y") == 0)
      {
         functype = func_y;
         letter = "y";
      }
      else
         assert(0);
   }

   bool cont;
   std::string line;

   std::cout << "Welcome.\n"
      "This program will generate spot tests for the Bessel " << letter << " function\n\n";
   do{
      get_user_parameter_info(arg1, "v");
      get_user_parameter_info(arg2, "x");
      boost::math::ntl::RR (*fp)(boost::math::ntl::RR, boost::math::ntl::RR);
      if(functype == func_J)
         fp = cyl_bessel_j_bare;
      else if(functype == func_I)
         fp = cyl_bessel_i_bare;
      else if(functype == func_K)
         fp = cyl_bessel_k_bare;
      else if(functype == func_Y)
         fp = cyl_neumann_bare;
      else if(functype == func_j)
         fp = sph_bessel_j_bare;
      else if(functype == func_y)
         fp = sph_bessel_y_bare;
      else
         assert(0);

      data.insert(fp, arg1, arg2);

      std::cout << "Any more data [y/n]?";
      std::getline(std::cin, line);
      boost::algorithm::trim(line);
      cont = (line == "y");
   }while(cont);

   std::cout << "Enter name of test data file [default=bessel_j_data.ipp]";
   std::getline(std::cin, line);
   boost::algorithm::trim(line);
   if(line == "")
      line = "bessel_j_data.ipp";
   std::ofstream ofs(line.c_str());
   line.erase(line.find('.'));
   ofs << std::scientific;
   write_code(ofs, data, line.c_str());

   return 0;
}


#include <iostream>
#include <vector>
#include <string>
#include <complex>

#include <boost/program_options.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>

#include <amgcl/backend/builtin.hpp>
#include <amgcl/value_type/complex.hpp>
#include <amgcl/adapter/crs_tuple.hpp>

#include <amgcl/mpi/util.hpp>
#include <amgcl/mpi/make_solver.hpp>
#include <amgcl/mpi/preconditioner.hpp>
#include <amgcl/mpi/solver/runtime.hpp>

#include <amgcl/io/mm.hpp>
#include <amgcl/io/binary.hpp>
#include <amgcl/profiler.hpp>

namespace amgcl {
    profiler<> prof;
}

namespace math = amgcl::math;

//---------------------------------------------------------------------------
ptrdiff_t assemble_poisson3d(amgcl::mpi::communicator comm,
        ptrdiff_t n, int block_size,
        std::vector<ptrdiff_t> &ptr,
        std::vector<ptrdiff_t> &col,
        std::vector<std::complex<double>> &val,
        std::vector<std::complex<double>> &rhs)
{
    ptrdiff_t n3 = n * n * n;

    ptrdiff_t chunk = (n3 + comm.size - 1) / comm.size;
    if (chunk % block_size != 0) {
        chunk += block_size - chunk % block_size;
    }
    ptrdiff_t row_beg = std::min(n3, chunk * comm.rank);
    ptrdiff_t row_end = std::min(n3, row_beg + chunk);
    chunk = row_end - row_beg;

    ptr.clear(); ptr.reserve(chunk + 1);
    col.clear(); col.reserve(chunk * 7);
    val.clear(); val.reserve(chunk * 7);

    rhs.resize(chunk);
    std::fill(rhs.begin(), rhs.end(), 1.0);

    const double h2i = (n - 1) * (n - 1);
    ptr.push_back(0);

    for (ptrdiff_t idx = row_beg; idx < row_end; ++idx) {
        ptrdiff_t k = idx / (n * n);
        ptrdiff_t j = (idx / n) % n;
        ptrdiff_t i = idx % n;

        if (k > 0)  {
            col.push_back(idx - n * n);
            val.push_back(-h2i);
        }

        if (j > 0)  {
            col.push_back(idx - n);
            val.push_back(-h2i);
        }

        if (i > 0) {
            col.push_back(idx - 1);
            val.push_back(-h2i);
        }

        col.push_back(idx);
        val.push_back(6 * h2i);

        if (i + 1 < n) {
            col.push_back(idx + 1);
            val.push_back(-h2i);
        }

        if (j + 1 < n) {
            col.push_back(idx + n);
            val.push_back(-h2i);
        }

        if (k + 1 < n) {
            col.push_back(idx + n * n);
            val.push_back(-h2i);
        }

        ptr.push_back( col.size() );
    }

    return chunk;
}

//---------------------------------------------------------------------------
void solve_scalar(
        amgcl::mpi::communicator comm,
        ptrdiff_t chunk,
        const std::vector<ptrdiff_t> &ptr,
        const std::vector<ptrdiff_t> &col,
        const std::vector<std::complex<double>> &val,
        const boost::property_tree::ptree &prm,
        const std::vector<std::complex<double>> &rhs
        )
{
    typedef amgcl::backend::builtin<std::complex<double>> Backend;

    typedef
        amgcl::mpi::make_solver<
            amgcl::runtime::mpi::preconditioner<Backend>,
            amgcl::runtime::mpi::solver::wrapper<Backend>
            >
        Solver;

    using amgcl::prof;

    prof.tic("setup");
    Solver solve(comm, std::tie(chunk, ptr, col, val), prm);
    prof.toc("setup");

    if (comm.rank == 0) {
        std::cout << solve << std::endl;
    }

    std::vector<std::complex<double>> x(chunk);

    int    iters;
    double error;

    prof.tic("solve");
    std::tie(iters, error) = solve(rhs, x);
    prof.toc("solve");

    if (comm.rank == 0) {
        std::cout
            << "Iterations: " << iters << std::endl
            << "Error:      " << error << std::endl
            << prof << std::endl;
    }
}

//---------------------------------------------------------------------------
int main(int argc, char *argv[]) {
    amgcl::mpi::init_thread mpi(&argc, &argv);
    amgcl::mpi::communicator comm(MPI_COMM_WORLD);

    if (comm.rank == 0)
        std::cout << "World size: " << comm.size << std::endl;

    using amgcl::prof;

    // Read configuration from command line
    namespace po = boost::program_options;
    po::options_description desc("Options");

    desc.add_options()
        ("help,h", "show help")
        (
         "size,n",
         po::value<ptrdiff_t>()->default_value(128),
         "domain size"
        )
        ("prm-file,P",
         po::value<std::string>(),
         "Parameter file in json format. "
        )
        (
         "prm,p",
         po::value< std::vector<std::string> >()->multitoken(),
         "Parameters specified as name=value pairs. "
         "May be provided multiple times. Examples:\n"
         "  -p solver.tol=1e-3\n"
         "  -p precond.coarse_enough=300"
        )
        ;

    po::positional_options_description p;
    p.add("prm", -1);

    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).options(desc).positional(p).run(), vm);
    po::notify(vm);

    if (vm.count("help")) {
        if (comm.rank == 0) std::cout << desc << std::endl;
        return 0;
    }

    boost::property_tree::ptree prm;
    if (vm.count("prm-file")) {
        read_json(vm["prm-file"].as<std::string>(), prm);
    }

    if (vm.count("prm")) {
        for(const std::string &v : vm["prm"].as<std::vector<std::string> >()) {
            amgcl::put(prm, v);
        }
    }

    ptrdiff_t n;
    std::vector<ptrdiff_t> ptr;
    std::vector<ptrdiff_t> col;
    std::vector<std::complex<double>> val;
    std::vector<std::complex<double>> rhs;

    prof.tic("assemble");
    n = assemble_poisson3d(comm, vm["size"].as<ptrdiff_t>(), 1, ptr, col, val, rhs);
    prof.toc("assemble");

    solve_scalar(comm, n, ptr, col, val, prm, rhs);
}