#include <stdio.h>
#include <cfloat>
#include <list>
#include <vector>
#include <math.h>
#include <iostream>
#include <algorithm>
#include <numeric>
#include <boost/algorithm/clamp.hpp>
#include <random>
#include<opencv2/opencv.hpp>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/util/work_sharder.h"
#include <boost/algorithm/clamp.hpp>
#include <boost/mpl/map.hpp>
#include <boost/mpl/int.hpp>
#include <boost/mpl/at.hpp>
#include "bboxes.h"
#include <future>
#include "wtoolkit.h"

using namespace tensorflow;
using namespace std;
namespace ba=boost::algorithm;
namespace bm=boost::mpl;

typedef Eigen::ThreadPoolDevice CPUDevice;

/*
 *
 * masks:[batch_size,Nr,h,w]
 * labels: [batch_size,Nr]
 * lens:[batch_size]
 * output_bboxes:[batch_size,nr,4] (ymin,xmin,ymax,xmax)
 * output_labels:[batch_size,nr]
 * output_lens:[batch_size]
 * output_ids:[batch_size,nr] 用于表示实例的编号，如第一个batch中的第二个实例所生成的所有的box的ids为3(id的编号从1开始)
 */
REGISTER_OP("MaskLineBboxes")
    .Attr("T: {int64,int32}")
	.Attr("max_output_nr:int")
    .Input("mask: uint8")
    .Input("labels: T")
    .Input("lens: int32")
	.Output("output_bboxes:float")
	.Output("output_labels:T")
	.Output("output_lens:int32")
	.Output("output_ids:int32")
	.SetShapeFn([](shape_inference::InferenceContext* c) {
            int        nr = -1;

            c->GetAttr("max_output_nr",&nr);

            const auto batch_size = c->Value(c->Dim(c->input(0),0));
            const auto shape0     = c->MakeShape({batch_size,nr,4});
            const auto shape1     = c->Matrix(batch_size,nr);
            const auto shape2     = c->Vector(batch_size);

			c->set_output(0, shape0);
			c->set_output(1, shape1);
			c->set_output(2, shape2);
			c->set_output(3, shape1);
			return Status::OK();
			});

template <typename Device,typename T>
class MaskLineBboxesOp: public OpKernel {
    private:
        using bbox_t = tuple<float,float,float,float>;
	public:
		explicit MaskLineBboxesOp(OpKernelConstruction* context) : OpKernel(context) {
			OP_REQUIRES_OK(context, context->GetAttr("max_output_nr", &max_output_nr_));
		}

		void Compute(OpKernelContext* context) override
		{
            TIME_THISV1("MaskLineBboxes");
			const Tensor &_mask= context->input(0);
			const Tensor &_labels= context->input(1);
			const Tensor &_lens = context->input(2);
			auto mask= _mask.template tensor<uint8_t,4>();
            auto labels = _labels.template tensor<T,2>();
            auto lens = _lens.template tensor<int32_t,1>();

			OP_REQUIRES(context, _mask.dims() == 4, errors::InvalidArgument("mask data must be 4-dimensional"));
			OP_REQUIRES(context, _labels.dims() == 2, errors::InvalidArgument("labels data must be 2-dimensional"));
			OP_REQUIRES(context, _lens.dims() == 1, errors::InvalidArgument("lens data must be 1-dimensional"));

			const auto     batch_size = _mask.dim_size(0);
			const auto     data_nr   = _mask.dim_size(1);
            list<vector<bbox_t>> out_bboxes;
            list<vector<int>> out_labels;
            list<vector<int>> out_ids;

            for(auto i=0; i<batch_size; ++i) {
                vector<bbox_t> res;
                vector<int> res_labels;
                vector<int> res_ids;
                res.reserve(1024);
                for(auto j=0; j<lens(i); ++j) {
                    const auto label = labels(i,j);
                    auto res0 = get_bboxes(mask.chip(i,0).chip(j,0));
                    if(!res0.empty()) {
                        res.insert(res.end(),res0.begin(),res0.end());
                        res_labels.insert(res_labels.end(),res0.size(),label);
                        res_ids.insert(res_ids.end(),res0.size(),j+1);
                    }
                }
			    OP_REQUIRES(context, res.size() == res_labels.size(), errors::InvalidArgument("size of bboxes should equal size of labels."));
                out_bboxes.push_back(std::move(res));
                out_labels.push_back(std::move(res_labels));
                out_ids.push_back(std::move(res_ids));
            }

            auto output_nr = max_output_nr_;

            if(output_nr<=0) {
                auto it = max_element(out_labels.begin(),out_labels.end(),[](const auto& v0,const auto& v1){ return v0.size()<v1.size();});
                output_nr = it->size();
            }

			int dims_3d[3] = {batch_size,output_nr,4};
			int dims_2d[2] = {batch_size,output_nr};
			int dims_1d[1] = {batch_size};
			TensorShape  outshape0;
			TensorShape  outshape1;
			TensorShape  outshape2;
			Tensor      *output_bbox   = NULL;
			Tensor      *output_labels = NULL;
			Tensor      *output_lens   = NULL;
			Tensor      *output_ids    = NULL;

			TensorShapeUtils::MakeShape(dims_3d, 3, &outshape0);
			TensorShapeUtils::MakeShape(dims_2d, 2, &outshape1);
			TensorShapeUtils::MakeShape(dims_1d, 1, &outshape2);

			OP_REQUIRES_OK(context, context->allocate_output(0, outshape0, &output_bbox));
			OP_REQUIRES_OK(context, context->allocate_output(1, outshape1, &output_labels));
			OP_REQUIRES_OK(context, context->allocate_output(2, outshape2, &output_lens));
			OP_REQUIRES_OK(context, context->allocate_output(3, outshape1, &output_ids));

			auto obbox   = output_bbox->template tensor<float,3>();
			auto olabels = output_labels->template tensor<T,2>();
			auto olens   = output_lens->template tensor<int32_t,1>();
			auto oids    = output_ids->template tensor<int32_t,2>();

            obbox.setZero();
            olabels.setZero();
            auto itb = out_bboxes.begin();
            auto itl = out_labels.begin();
            auto iti = out_ids.begin();

			for(int i=0; i<batch_size; ++i,++itb,++itl,++iti) {
                olens(i) = itl->size();
                for(auto j=0; j<olens(i); ++j) {
                    obbox(i,j,0) = std::get<0>((*itb)[j]);
                    obbox(i,j,1) = std::get<1>((*itb)[j]);
                    obbox(i,j,2) = std::get<2>((*itb)[j]);
                    obbox(i,j,3) = std::get<3>((*itb)[j]);
                    olabels(i,j) = (*itl)[j];
                    oids(i,j) = (*iti)[j];
                }
			}
		}
        /*
         * mask: [h,w]
         */
        vector<bbox_t> get_bboxes(const Eigen::Tensor<uint8_t,2,Eigen::RowMajor>& mask) {
            const auto h = mask.dimension(0);
            const auto w = mask.dimension(1);
            const auto y_delta = 1.0/h;
            const auto x_delta = 1.0/w;
            vector<bbox_t> res;
            res.reserve(256);

            for(auto i=0; i<h; ++i) {
                const auto ymin = i*y_delta;
                const auto ymax = (i+1)*y_delta;
                for(auto j=0; j<w; ++j) {
                    if(mask(i,j)<1) continue;
                    auto begin_j = j;
                    while((mask(i,j)>0) && (j<w))++j;
                    const auto xmin = begin_j*x_delta;
                    const auto xmax = (j==w)?1.0:j*x_delta;
                    res.emplace_back(ymin,xmin,ymax,xmax);
                }
            }
            return res;
        }
	private:
        int max_output_nr_ = 0;
};
REGISTER_KERNEL_BUILDER(Name("MaskLineBboxes").Device(DEVICE_CPU).TypeConstraint<int32_t>("T"), MaskLineBboxesOp<CPUDevice,int32_t>);
REGISTER_KERNEL_BUILDER(Name("MaskLineBboxes").Device(DEVICE_CPU).TypeConstraint<tensorflow::int64>("T"), MaskLineBboxesOp<CPUDevice,tensorflow::int64>);

/*
 *
 * masks:[Nr,h,w]
 * bboxes:[Nr,4]
 * size:[2]={H,W}
 * output_masks:[Nr,H,W]
 */
REGISTER_OP("FullSizeMask")
    .Attr("T: {float32,uint8}")
    .Input("mask: T")
    .Input("bboxes: float32")
    .Input("size: int32")
	.Output("output_masks:T")
	.SetShapeFn([](shape_inference::InferenceContext* c) {
            const auto nr = c->Value(c->Dim(c->input(0),0));
            const auto shape0     = c->MakeShape({nr,-1,-1});
			c->set_output(0, shape0);
			return Status::OK();
			});

template <typename Device,typename T>
class FullSizeMaskOp: public OpKernel {
    private:
        using bbox_t = tuple<float,float,float,float>;
        using type_to_int = bm::map<
              bm::pair<uint8_t,bm::int_<CV_8UC1>>
                  , bm::pair<float,bm::int_<CV_32FC1>>
             >;
        using tensor_t = Eigen::Tensor<T,2,Eigen::RowMajor>;
        using tensor_map_t = Eigen::TensorMap<tensor_t>;
	public:
		explicit FullSizeMaskOp(OpKernelConstruction* context) : OpKernel(context) {
		}

		void Compute(OpKernelContext* context) override
		{
			const Tensor &_mask= context->input(0);
			const Tensor &_bboxes = context->input(1);
			const Tensor &_size= context->input(2);

			OP_REQUIRES(context, _mask.dims() == 3, errors::InvalidArgument("mask data must be 3-dimensional"));
			OP_REQUIRES(context, _bboxes.dims() == 2, errors::InvalidArgument("bboxes data must be 2-dimensional"));
			OP_REQUIRES(context, _size.dims() == 1, errors::InvalidArgument("size data must be 1-dimensional"));

			auto         mask        = _mask.template flat<T>();
			auto         bboxes      = _bboxes.template tensor<float,2>();
			auto         size        = _size.template tensor<int,1>();
			const auto   mh          = _mask.dim_size(1);
			const auto   mw          = _mask.dim_size(2);
			const auto   H           = size(0);
			const auto   W           = size(1);
			const auto   data_nr     = _mask.dim_size(0);
			int          dims_3d[3]  = {data_nr,H,W};
			Tensor      *output_mask = NULL;
			TensorShape  outshape0;

			TensorShapeUtils::MakeShape(dims_3d, 3, &outshape0);

			OP_REQUIRES_OK(context, context->allocate_output(0, outshape0, &output_mask));

            auto o_tensor = output_mask->template tensor<T,3>();
            const auto H_max = H-1;
            const auto W_max = W-1;
            constexpr auto kMinSize = 1e-3;

            o_tensor.setZero();

            for(auto i=0; i<data_nr; ++i) {
                if((fabs(bboxes(i,3)-bboxes(i,1))<kMinSize)
                    || (fabs(bboxes(i,2)-bboxes(i,0))<kMinSize))
                    continue;

                long xmin = ba::clamp(bboxes(i,1)*W_max,0,W_max);
                long ymin = ba::clamp(bboxes(i,0)*H_max,0,H_max);
                long xmax = ba::clamp(bboxes(i,3)*W_max,0,W_max);
                long ymax = ba::clamp(bboxes(i,2)*H_max,0,H_max);
                const cv::Mat input_mask(mh,mw,bm::at<type_to_int,T>::type::value,(void*)(mask.data()+i*mh*mw));
                cv::Mat dst_mask(ymax-ymin+1,xmax-xmin+1,bm::at<type_to_int,T>::type::value);

                cv::resize(input_mask,dst_mask,cv::Size(xmax-xmin+1,ymax-ymin+1),0,0,cv::INTER_LINEAR);

                tensor_map_t src_map((T*)dst_mask.data,dst_mask.rows,dst_mask.cols);
                Eigen::array<long,2> offset = {ymin,xmin};
                Eigen::array<long,2> extents = {dst_mask.rows,dst_mask.cols};

                o_tensor.chip(i,0).slice(offset,extents) = src_map;

                /*if(((xmax-xmin>mw) || (ymax-ymin>mh)) && (xmax>xmin) && (ymax>ymin)) {
                    cv::Mat dst_mask(H,W,bm::at<type_to_int,T>::type::value,output_mask->template flat<T>().data()+H*W*i);
                    cv::Mat src_mask = dst_mask.clone();
                    const auto k = max<int>(3,sqrt((xmax-xmin)*(ymax-ymin)/(mh*mw))+1);
                    cv::medianBlur(src_mask,dst_mask,(k/2)*2+1);
                }*/

            }
		}
};
REGISTER_KERNEL_BUILDER(Name("FullSizeMask").Device(DEVICE_CPU).TypeConstraint<float>("T"), FullSizeMaskOp<CPUDevice,float>);
REGISTER_KERNEL_BUILDER(Name("FullSizeMask").Device(DEVICE_CPU).TypeConstraint<uint8_t>("T"), FullSizeMaskOp<CPUDevice,uint8_t>);

/*
 * 对mask [Nr,H,W] 旋转指定角度，同时返回相应instance的bbox
 * bbox [N,4],[ymin,xmin,ymax,xmax], 绝对坐标
 */
REGISTER_OP("MaskRotate")
    .Attr("T: {uint8,float}")
    .Input("mask: T")
    .Input("angle: float")
	.Output("o_image:T")
    .SetShapeFn([](shape_inference::InferenceContext* c){
            auto input_shape0 = c->input(0);
            c->set_output(0, input_shape0);
			return Status::OK();
            });

template <typename Device, typename T>
class MaskRotateOp: public OpKernel {
    private:
        using Tensor3D = Eigen::Tensor<T,3,Eigen::RowMajor>;
        using type_to_int = bm::map<
              bm::pair<uint8_t,bm::int_<CV_8UC1>>
                  , bm::pair<float,bm::int_<CV_32FC1>>
             >;
	public:
		explicit MaskRotateOp(OpKernelConstruction* context) : OpKernel(context) {
		}
		void Compute(OpKernelContext* context) override
		{
            TIME_THISV1("RotateMask");
			const Tensor &_input_img = context->input(0);
			const Tensor &_angle = context->input(1);

			OP_REQUIRES(context, _input_img.dims() == 3, errors::InvalidArgument("tensor must be a 3-dimensional tensor"));
			OP_REQUIRES(context, _angle.dims() == 0, errors::InvalidArgument("angle be a 0-dimensional tensor"));

            auto         input_img     = _input_img.template flat<T>().data();
            auto         angle         = _angle.template flat<float>().data()[0];
            const auto   img_channel   = _input_img.dim_size(0);
            const auto   img_height    = _input_img.dim_size(1);
            const auto   img_width     = _input_img.dim_size(2);
            Tensor      *output_tensor = nullptr;

			OP_REQUIRES_OK(context, context->allocate_output(0, _input_img.shape(), &output_tensor));

            auto          o_tensor     = output_tensor->template flat<T>().data();
            const cv::Point2f cp(img_width/2,img_height/2);
            const cv::Mat r            = cv::getRotationMatrix2D(cp,angle,1.0);
            const auto    cv_type      = bm::at<type_to_int,T>::type::value;
            auto fn  = [img_height,img_width,cv_type,r](T* i_data,T* o_data) {
                cv::Mat i_img(img_height,img_width,cv_type,i_data);
                cv::Mat o_img(img_height,img_width,cv_type,o_data);

                cv::warpAffine(i_img,o_img,r,cv::Size(img_width,img_height));
            };
            list<future<void>> futures;

            for(auto i=0; i<img_channel; ++i) {
                auto i_data = input_img+i *img_width *img_height;
                auto o_data = o_tensor+i *img_width *img_height;
                futures.emplace_back(async(launch::async,fn,(T*)i_data,o_data));
                if(futures.size()>8)
                    futures.pop_front();
            }
            futures.clear();
        }
};
REGISTER_KERNEL_BUILDER(Name("MaskRotate").Device(DEVICE_CPU).TypeConstraint<uint8_t>("T"), MaskRotateOp<CPUDevice, uint8_t>);
REGISTER_KERNEL_BUILDER(Name("MaskRotate").Device(DEVICE_CPU).TypeConstraint<float>("T"), MaskRotateOp<CPUDevice, float>);

/*
 * mask [Nr,H,W] 中查找目标instance的bbox
 * bbox [N,4],[ymin,xmin,ymax,xmax], 绝对坐标
 */
REGISTER_OP("GetBboxesFromMask")
    .Attr("T: {uint8,float}")
    .Input("mask: T")
	.Output("bbox:float")
    .SetShapeFn([](shape_inference::InferenceContext* c){
            auto input_shape0 = c->input(0);
            auto data_nr = c->Dim(input_shape0,0);
            auto output_shape = c->MakeShape({data_nr,4});
            c->set_output(0, output_shape);
			return Status::OK();
            });

template <typename Device, typename T>
class GetBboxesFromMaskOp: public OpKernel {
    private:
        using Tensor3D = Eigen::Tensor<T,3,Eigen::RowMajor>;
        using type_to_int = bm::map<
              bm::pair<uint8_t,bm::int_<CV_8UC1>>
                  , bm::pair<float,bm::int_<CV_32FC1>>
             >;
	public:
		explicit GetBboxesFromMaskOp(OpKernelConstruction* context) : OpKernel(context) {
		}
		void Compute(OpKernelContext* context) override
		{
            TIME_THISV1("GetBboxesFromMask");
			const Tensor &_input_img = context->input(0);

			OP_REQUIRES(context, _input_img.dims() == 3, errors::InvalidArgument("tensor must be a 3-dimensional tensor"));

            auto         input_img     = _input_img.template flat<T>().data();
            const auto   img_channel   = _input_img.dim_size(0);
            const auto   img_height    = _input_img.dim_size(1);
            const auto   img_width     = _input_img.dim_size(2);
            const int    dim2d[]       = {img_channel,4};
            TensorShape  output_shape;
            Tensor      *output_bbox   = nullptr;

            TensorShapeUtils::MakeShape(dim2d,2,&output_shape);

			OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output_bbox));

            auto          o_bbox       = output_bbox->template flat<float>().data();
            const auto    cv_type      = bm::at<type_to_int,T>::type::value;
            auto fn  = [img_height,img_width,cv_type](T* i_data,float* bbox) {
                cv::Mat i_img(img_height,img_width,cv_type,i_data);
                getBBox(i_img,bbox);
            };
            list<future<void>> futures;

            for(auto i=0; i<img_channel; ++i) {
                auto i_data = input_img+i *img_width *img_height;
                auto bbox   = o_bbox+i *4;
                futures.emplace_back(async(launch::async,fn,(T*)i_data,bbox));
                if(futures.size()>8)
                    futures.pop_front();
            }
            futures.clear();
        }

        static void getBBox(const cv::Mat& img,float* bbox)
        {
            vector<vector<cv::Point>> contours;
            vector<cv::Vec4i> hierarchy;
            vector<cv::Point> points;
            const auto    cv_type      = bm::at<type_to_int,T>::type::value;
            cv::Mat dst_img(img.rows,img.cols,cv_type);

            if(cv_type == CV_32FC1) {
                cv::threshold(img,dst_img,0.5,255,cv::THRESH_BINARY);
                cv::Mat dst_img1;
                dst_img.convertTo(dst_img1,CV_8UC1);
                cv::findContours(dst_img1, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE, cv::Point(0,0));
            } else {
                cv::threshold(img,dst_img,127,255,cv::THRESH_BINARY);
                cv::findContours(dst_img, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE, cv::Point(0,0));
            }

            for (auto &cont:contours)
                points.insert(points.end(),cont.begin(),cont.end());

            if(points.size()<2) {
               memset(bbox,0,sizeof(float)*4);
               return;
            }

            const auto rect = cv::boundingRect(points);

            bbox[0] = rect.y;
            bbox[1] = rect.x;
            bbox[2] = rect.y+rect.height;
            bbox[3] = rect.x+rect.width;
        }
};
REGISTER_KERNEL_BUILDER(Name("GetBboxesFromMask").Device(DEVICE_CPU).TypeConstraint<uint8_t>("T"), GetBboxesFromMaskOp<CPUDevice, uint8_t>);
REGISTER_KERNEL_BUILDER(Name("GetBboxesFromMask").Device(DEVICE_CPU).TypeConstraint<float>("T"), GetBboxesFromMaskOp<CPUDevice, float>);
/*
 * 输入mask[H,W,N]
 * 输入labels[N]
 * set_background:如果一个位置没有标签，则默认为背景
 * attr:num_classes
 * 输出mask[W,H,num_classes]
 */
REGISTER_OP("SparseMaskToDense")
    .Attr("T: {int32,bool,int8,uint8}")
	.Attr("num_classes:int")
	.Attr("set_background:bool")
    .Input("mask: T")
    .Input("labels: int32")
	.Output("data:T")
	.SetShapeFn([](shape_inference::InferenceContext* c) {
			int num_classes = 0;
			c->GetAttr("num_classes",&num_classes);
			auto w = c->Value(c->Dim(c->input(0),0));
			auto h = c->Value(c->Dim(c->input(0),1));
            auto output_shape = c->MakeShape({h, w, num_classes});
			c->set_output(0,output_shape);
			return Status::OK();
			});

template <typename Device, typename T>
class SparseMaskToDenseOp: public OpKernel {
	public:
		explicit SparseMaskToDenseOp(OpKernelConstruction* context) : OpKernel(context) {
			OP_REQUIRES_OK(context,
					context->GetAttr("num_classes", &num_classes_));
			OP_REQUIRES_OK(context,
					context->GetAttr("set_background", &set_background_));
		}
		void Compute(OpKernelContext* context) override
        {
            const Tensor &_mask   = context->input(0);
            const Tensor &_labels = context->input(1);
            auto          mask    = _mask.template tensor<T,3>();
            auto          labels  = _labels.template flat<int>().data();
            auto          h       = _mask.dim_size(0);
            auto          w       = _mask.dim_size(1);
            auto          nr      = _mask.dim_size(2);
            auto          nr1     = _labels.dim_size(0);

            OP_REQUIRES(context, _labels.dims()==1, errors::InvalidArgument("labels must be 1-dimensional"));
            OP_REQUIRES(context, _mask.dims()==3, errors::InvalidArgument("mask must be 3-dimensional"));
            OP_REQUIRES(context, nr==nr1, errors::InvalidArgument("size unmatch."));

            int          dims3d[]     = {int(h),int(w),num_classes_};
            Tensor      *output_data  = NULL;
            TensorShape  output_shape;

            TensorShapeUtils::MakeShape(dims3d, 3, &output_shape);

            OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output_data));

            auto      oq_tensor = output_data->template tensor<T,3>();
            oq_tensor.setZero();
			using tensor_t = Eigen::Tensor<T,2,Eigen::RowMajor>;

            for(auto i=0; i<nr; ++i) {
				const auto label = labels[i];
				if((label<0) || (label>=num_classes_)) {
					cout<<"Error label "<<label<<", not in range [0,"<<num_classes_<<")"<<endl;
					continue;
				}
                auto t = oq_tensor.chip(label,2);
                t = (t||mask.chip(i,2)).template cast<T>();
            }

			if(!set_background_) return;

			for(auto i=0; i<h; ++i) {
				for(auto j=0; j<w; ++j) {
					bool have_label = false;
					for(auto k=0; k<nr; ++k) {
						const auto label = labels[k];

						if((label<0) || (label>=num_classes_)) continue;

						if(mask(i,j,k) != 0) {
							have_label = true;
							break;
						}
					}
					if(have_label) continue;

					oq_tensor(i,j,0) = T(true);
				}
			}
        }
	private:
		int num_classes_ = 0;
		bool set_background_ = false;
};
REGISTER_KERNEL_BUILDER(Name("SparseMaskToDense").Device(DEVICE_CPU).TypeConstraint<int32_t>("T"), SparseMaskToDenseOp<CPUDevice, int32_t>);
REGISTER_KERNEL_BUILDER(Name("SparseMaskToDense").Device(DEVICE_CPU).TypeConstraint<bool>("T"), SparseMaskToDenseOp<CPUDevice, bool>);
REGISTER_KERNEL_BUILDER(Name("SparseMaskToDense").Device(DEVICE_CPU).TypeConstraint<int8_t>("T"), SparseMaskToDenseOp<CPUDevice, int8_t>);
REGISTER_KERNEL_BUILDER(Name("SparseMaskToDense").Device(DEVICE_CPU).TypeConstraint<uint8_t>("T"), SparseMaskToDenseOp<CPUDevice, uint8_t>);

/*
 * mask_size: 输出的mask像素大小
 * masks:[Nr,h,w]
 * test_threshold: 距离小于test_threshold的bbox才进行是否可以合并的检测
 * kerner_size: 闭运算kernel大小
 * bboxes:[Nr,4] absolute coordinate
 * labels:[Nr]
 * mask:[Nr,H,W]
 * output:
 * out_bboxes: [out_Nr,4]
 * out_labels: [out_Nr]
 * out_masks: [out_Nr,mask_size,mask_size]
 * out_indices: [Nr] 相同的值表示相应位置的输入合并在了一起
 */
REGISTER_OP("MergeInstanceByMask")
    .Attr("T: {float32,uint8}")
    .Attr("mask_size:int=63")
    .Attr("test_threshold:int=8")
    .Attr("kernel_size:int=3")
    .Input("bboxes: T")
    .Input("labels: int32")
    .Input("probability: T")
    .Input("mask: uint8")
	.Output("out_bboxes:T")
	.Output("out_labels:int32")
	.Output("out_probability:T")
	.Output("out_masks:uint8")
	.Output("out_indices:int32")
	.SetShapeFn([](shape_inference::InferenceContext* c) {
            int mask_size = 0;
            c->GetAttr("mask_size",&mask_size);
            const auto nr = c->Value(c->Dim(c->input(0),0));
            const auto shape0     = c->MakeShape({-1,4});
            const auto shape1     = c->MakeShape({-1});
            const auto shape2     = c->MakeShape({nr});
            const auto shape3     = c->MakeShape({-1,mask_size,mask_size});
			c->set_output(0, shape0);
			c->set_output(1, shape1);
			c->set_output(2, shape1);
			c->set_output(3, shape3);
			c->set_output(4, shape2);
			return Status::OK();
			});

template <typename Device,typename T>
class MergeInstanceByMaskOp: public OpKernel {
    private:
        using bbox_t = Eigen::Tensor<T,1,Eigen::RowMajor>;
        using mask_t = Eigen::Tensor<uint8_t,2,Eigen::RowMajor>;
    public:
        explicit MergeInstanceByMaskOp(OpKernelConstruction* context) : OpKernel(context) {
            OP_REQUIRES_OK(context, context->GetAttr("mask_size", &mask_size_));
            OP_REQUIRES_OK(context, context->GetAttr("test_threshold", &test_threshold_));
            OP_REQUIRES_OK(context, context->GetAttr("kernel_size", &kerner_size_));
        }

        void Compute(OpKernelContext* context) override
        {
            const Tensor &_bboxes = context->input(0);
            const Tensor &_labels= context->input(1);
            const Tensor &_probability = context->input(2);
            const Tensor &_mask= context->input(3);

            OP_REQUIRES(context, _mask.dims() == 3, errors::InvalidArgument("mask data must be 3-dimensional"));
            OP_REQUIRES(context, _bboxes.dims() == 2, errors::InvalidArgument("bboxes data must be 2-dimensional"));
            OP_REQUIRES(context, _labels.dims() == 1, errors::InvalidArgument("labels data must be 1-dimensional"));
            OP_REQUIRES(context, _probability.dims() == 1, errors::InvalidArgument("probability data must be 1-dimensional"));

            auto       mask        = _mask.template tensor<uint8_t,3>();
            auto       bboxes      = _bboxes.template tensor<T,2>();
            auto       labels      = _labels.template tensor<int,1>();
            auto       probability = _probability.template tensor<T,1>();
            const auto data_nr     = _mask.dim_size(0);

            bool is_merged = false;
            vector<bbox_t> in_bboxes_data;
            vector<int> in_labels_data;
            vector<mask_t> in_mask_data;
            vector<int> in_indices_data;
            vector<T> in_probability;
            vector<bool> need_remove(data_nr,false);

            for(auto i=0; i<data_nr; ++i) {
                in_bboxes_data.push_back(bboxes.chip(i,0));
                in_labels_data.push_back(labels(i));
                in_mask_data.push_back(mask.chip(i,0));
                in_indices_data.push_back(i);
                in_probability.push_back(probability(i));
            }

            do {
                is_merged = false;
                if(in_bboxes_data.size()>1) {
                    for(auto i=0; i<in_bboxes_data.size()-1; ++i) {
                        if(need_remove[i])
                            continue;
                        auto& bbox0 = in_bboxes_data[i];
                        for(auto j=i+1; j<in_bboxes_data.size(); ++j) {
                            if(need_remove[j])
                                continue;
                            if(in_labels_data[i] != in_labels_data[j])
                                continue;

                            auto &bbox1 = in_bboxes_data[j];
                            auto  dis   = bboxes_distance(in_bboxes_data[i],in_bboxes_data[j]);

                            if(dis > test_threshold_)
                                continue;

                            try {
                                auto res = merge_bboxes({bbox0,bbox1},{in_mask_data[i],in_mask_data[j]});

                                in_mask_data[i] = get<1>(res);
                                in_bboxes_data[i] = get<0>(res);
                                in_indices_data[j] = in_indices_data[i];
                                in_probability[i] = max(in_probability[i],in_probability[j]);
                                need_remove[j] = true;
                                is_merged = true;
                            }catch(...) {
                            }
                        }
                    }
                }
            }while(is_merged);

            auto         total_nr           = count(need_remove.begin(),need_remove.end(),false);
            int          dims_1d[1]         = {total_nr};
            int          dims_2d[2]         = {total_nr,4};
            int          dims_1d2[1]        = {data_nr};
            int          dims_3d[3]         = {total_nr,mask_size_,mask_size_};
            TensorShape  outshape0;
            TensorShape  outshape1;
            TensorShape  outshape2;
            TensorShape  outshape3;
            Tensor      *output_bboxes      = NULL;
            Tensor      *output_labels      = NULL;
            Tensor      *output_probability = NULL;
            Tensor      *output_mask        = NULL;
            Tensor      *output_indices     = NULL;

            TensorShapeUtils::MakeShape(dims_1d, 1, &outshape0);
            TensorShapeUtils::MakeShape(dims_2d, 2, &outshape1);
            TensorShapeUtils::MakeShape(dims_1d2, 1, &outshape2);
            TensorShapeUtils::MakeShape(dims_3d, 3, &outshape3);

            OP_REQUIRES_OK(context, context->allocate_output(0, outshape1, &output_bboxes));
            OP_REQUIRES_OK(context, context->allocate_output(1, outshape0, &output_labels));
            OP_REQUIRES_OK(context, context->allocate_output(2, outshape0, &output_probability));
            OP_REQUIRES_OK(context, context->allocate_output(3, outshape3, &output_mask));
            OP_REQUIRES_OK(context, context->allocate_output(4, outshape2, &output_indices));

            auto o_mask        = output_mask->template tensor<uint8_t,3>();
            auto o_bboxes      = output_bboxes->template tensor<T,2>();
            auto o_labels      = output_labels->template tensor<int,1>();
            auto o_probability = output_probability->template tensor<T,1>();
            auto o_indices     = output_indices->template tensor<int,1>();

            for(auto i=0,j=0; i<data_nr; ++i) {
                o_indices(i) = in_indices_data[i];
                if(need_remove[i])
                    continue;
                o_bboxes.chip(j,0) = in_bboxes_data[i];
                o_labels(j) = in_labels_data[i];
                o_probability(j) = in_probability[i];
                if(in_mask_data[i].dimension(0) == mask_size_) {
                    o_mask.chip(j,0) = in_mask_data[i];
                } else {
                    o_mask.chip(j,0) = resize_mask(in_mask_data[i],mask_size_);
                }
                ++j;
            }
        }
        mask_t resize_mask(const mask_t& msk,size_t size) {
            const cv::Mat input_mask(msk.dimension(0),msk.dimension(1),CV_8UC1,(void*)(msk.data()));
            cv::Mat dst_mask(size,size,CV_8UC1);
            mask_t res_mask(size,size);

            cv::resize(input_mask,dst_mask,cv::Size(size,size),0,0,cv::INTER_LINEAR);
            memcpy(res_mask.data(),dst_mask.data,size*size);

            return res_mask;
        }
        static vector<bbox_t> make_around_bboxes(const bbox_t& box0,const bbox_t& box1) {

            bbox_t res_box0(4);
            bbox_t res_box1(4);
            bbox_t res_box2(4);
            bbox_t res_box3(4);

            res_box0(0) = box0(0)-(box1(2)-box1(0));
            res_box0(1) = box0(1);
            res_box0(2) = box0(0);
            res_box0(3) = box0(3);

            res_box1(0) = box0(0);
            res_box1(1) = box0(3);
            res_box1(2) = box0(2);
            res_box1(3) = box0(3)+(box1(3)-box1(1));

            res_box2(0) = box0(2);
            res_box2(1) = box0(1);
            res_box2(2) = box0(2)+(box1(2)-box1(0));;
            res_box2(3) = box0(3);

            res_box3(0) = box0(0);
            res_box3(1) = box0(1)-(box1(3)-box1(1));
            res_box3(2) = box0(2);
            res_box3(3) = box0(1);
            return {box0,res_box0,res_box1,res_box2,res_box3};

        }
        float bboxes_distance(const bbox_t& box0,const bbox_t& box1) {
            auto  bboxes  = make_around_bboxes(box0,box1);
            int   index   = -1;
            float max_iou = -1.0;

            for(auto i=0; i<bboxes.size(); ++i) {
                auto v0 = bboxes_jaccard_of_box0v1(bboxes[i],box1);
                auto v1 = bboxes_jaccard_of_box0v1(box1,bboxes[i]);
                auto v = max(v0,v1);

                if(v>max_iou) {
                    max_iou = v;
                    index = i;
                }
            }
            if(max_iou<0.4)
                return 1e8;
            switch(index) {
                case 0:
                    return 0;
                case 1:
                    return fabs(box1(2)-box0(0));
                case 2:
                    return fabs(box1(1)-box0(3));
                case 3:
                    return fabs(box1(0)-box0(2));
                case 4:
                    return fabs(box1(3)-box0(1));
                default:
                    cout<<"Error type."<<endl;
                    return 1e8;
            }
        }
        void show_box(const bbox_t& box) {
            cout<<box(0)<<","<<box(1)<<","<<box(2)<<","<<box(3)<<endl;
        }
        static cv::Mat to_mat(const mask_t& msk)
        {
            cv::Mat res(msk.dimension(0),msk.dimension(1),CV_8UC1,(void*)msk.data());
            return res.clone();
        }
        tuple<bbox_t,mask_t> merge_bboxes(const vector<bbox_t>& boxes,const vector<mask_t>& masks) noexcept(false)
        {
            bbox_t  env_bbox(4);
            cv::Mat res_mask = cv::Mat::zeros(mask_size_,mask_size_,CV_8UC1);
            cv::Mat res_mask1 = cv::Mat::zeros(mask_size_,mask_size_,CV_8UC1);
            int total_nr = 0;

            bboxes_envelope(boxes[0],boxes[1],env_bbox);

            const auto env_bbox_w = env_bbox(3)-env_bbox(1);
            const auto env_bbox_h = env_bbox(2)-env_bbox(0);

            for(auto i=0; i<boxes.size(); ++i) {
                vector<vector<cv::Point>> contours;
                vector<vector<cv::Point>> new_contours;
                vector<cv::Vec4i> hierarchy;
                vector<cv::Point> points;
                cv::Mat dst_img = to_mat(masks[i]);
                const auto cur_bbox = boxes[i];
                const auto cur_mask_h = masks[i].dimension(0);
                const auto cur_mask_w = masks[i].dimension(1);
                const auto cur_bbox_w = cur_bbox(3)-cur_bbox(1);
                const auto cur_bbox_h = cur_bbox(2)-cur_bbox(0);

                cv::findContours(dst_img, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE, cv::Point(0,0));

                if(contours.size()==0) {
                    cout<<"Error contours size."<<endl;
                    throw std::runtime_error("error contours size.");
                }

                for(auto& points:contours) {
                    vector<cv::Point> new_points;
                    for(auto& p:points) {
                        auto x = (p.x*cur_bbox_w/cur_mask_w+cur_bbox(1)-env_bbox(1))*mask_size_/env_bbox_w;
                        auto y = (p.y*cur_bbox_h/cur_mask_h+cur_bbox(0)-env_bbox(0))*mask_size_/env_bbox_h;
                        new_points.push_back(cv::Point(x,y));
                    }
                    new_contours.push_back(new_points);
                }
                cv::drawContours(res_mask,new_contours,-1,cv::Scalar(1),cv::FILLED,8,hierarchy);

                total_nr += contours.size();
            }

            vector<vector<cv::Point>> contours;
            vector<vector<cv::Point>> new_contours;
            vector<cv::Vec4i> hierarchy;
            cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(kerner_size_, kerner_size_), cv::Point(-1, -1));

            cv::morphologyEx(res_mask, res_mask1, cv::MORPH_CLOSE, kernel);
            cv::findContours(res_mask1.clone(), contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE, cv::Point(0,0));

            if((contours.size()==1)
                    || (contours.size()<total_nr-2)) {
                mask_t r_mask(mask_size_,mask_size_);
                memcpy(r_mask.data(),res_mask1.data,mask_size_*mask_size_);
                return make_tuple(env_bbox,r_mask);
            }
            throw std::runtime_error("merge faild.");
        }
    private:
        int mask_size_      = 63;
        int test_threshold_ = 8;
        int kerner_size_    = 3;
};
REGISTER_KERNEL_BUILDER(Name("MergeInstanceByMask").Device(DEVICE_CPU).TypeConstraint<float>("T"), MergeInstanceByMaskOp<CPUDevice,float>);


#ifndef BOOST_GEOMETRY_PROJECTIONS_IMW_P_HPP
#define BOOST_GEOMETRY_PROJECTIONS_IMW_P_HPP

// Boost.Geometry - extensions-gis-projections (based on PROJ4)
// This file is automatically generated. DO NOT EDIT.

// Copyright (c) 2008-2015 Barend Gehrels, 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)

// This file is converted from PROJ4, http://trac.osgeo.org/proj
// PROJ4 is originally written by Gerald Evenden (then of the USGS)
// PROJ4 is maintained by Frank Warmerdam
// PROJ4 is converted to Boost.Geometry by Barend Gehrels

// Last updated version of proj: 4.9.1

// Original copyright notice:

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

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

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

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

#include <boost/geometry/extensions/gis/projections/impl/base_static.hpp>
#include <boost/geometry/extensions/gis/projections/impl/base_dynamic.hpp>
#include <boost/geometry/extensions/gis/projections/impl/projects.hpp>
#include <boost/geometry/extensions/gis/projections/impl/factory_entry.hpp>
#include <boost/geometry/extensions/gis/projections/impl/pj_mlfn.hpp>

namespace boost { namespace geometry { namespace projections
{
    #ifndef DOXYGEN_NO_DETAIL
    namespace detail { namespace imw_p
    {

            static const double TOL = 1e-10;
            static const double EPS = 1e-10;

            struct XY { double x, y; }; // specific for IMW_P

            struct par_imw_p
            {
                double    P, Pp, Q, Qp, R_1, R_2, sphi_1, sphi_2, C2;
                double    phi_1, phi_2, lam_1;
                double    en[EN_SIZE];
                int    mode; /* = 0, phi_1 and phi_2 != 0, = 1, phi_1 = 0, = -1 phi_2 = 0 */
            };

            template <typename Parameters>
                static int
            phi12(Parameters& par, par_imw_p& proj_parm, double *del, double *sig) {
                int err = 0;

                if (!pj_param(par.params, "tlat_1").i ||
                    !pj_param(par.params, "tlat_2").i) {
                    err = -41;
                } else {
                    proj_parm.phi_1 = pj_param(par.params, "rlat_1").f;
                    proj_parm.phi_2 = pj_param(par.params, "rlat_2").f;
                    *del = 0.5 * (proj_parm.phi_2 - proj_parm.phi_1);
                    *sig = 0.5 * (proj_parm.phi_2 + proj_parm.phi_1);
                    err = (fabs(*del) < EPS || fabs(*sig) < EPS) ? -42 : 0;
                }
                return err;
            }
            template <typename Parameters>
                static XY
            loc_for(double const& lp_lam, double const& lp_phi, Parameters const& par, par_imw_p const& proj_parm, double *yc) {
                XY xy;

                if (! lp_phi) {
                    xy.x = lp_lam;
                    xy.y = 0.;
                } else {
                    double xa, ya, xb, yb, xc, D, B, m, sp, t, R, C;

                    sp = sin(lp_phi);
                    m = pj_mlfn(lp_phi, sp, cos(lp_phi), proj_parm.en);
                    xa = proj_parm.Pp + proj_parm.Qp * m;
                    ya = proj_parm.P + proj_parm.Q * m;
                    R = 1. / (tan(lp_phi) * sqrt(1. - par.es * sp * sp));
                    C = sqrt(R * R - xa * xa);
                    if (lp_phi < 0.) C = - C;
                    C += ya - R;
                    if (proj_parm.mode < 0) {
                        xb = lp_lam;
                        yb = proj_parm.C2;
                    } else {
                        t = lp_lam * proj_parm.sphi_2;
                        xb = proj_parm.R_2 * sin(t);
                        yb = proj_parm.C2 + proj_parm.R_2 * (1. - cos(t));
                    }
                    if (proj_parm.mode > 0) {
                        xc = lp_lam;
                        *yc = 0.;
                    } else {
                        t = lp_lam * proj_parm.sphi_1;
                        xc = proj_parm.R_1 * sin(t);
                        *yc = proj_parm.R_1 * (1. - cos(t));
                    }
                    D = (xb - xc)/(yb - *yc);
                    B = xc + D * (C + R - *yc);
                    xy.x = D * sqrt(R * R * (1 + D * D) - B * B);
                    if (lp_phi > 0)
                        xy.x = - xy.x;
                    xy.x = (B + xy.x) / (1. + D * D);
                    xy.y = sqrt(R * R - xy.x * xy.x);
                    if (lp_phi > 0)
                        xy.y = - xy.y;
                    xy.y += C + R;
                }
                return (xy);
            }
            template <typename Parameters>
                static void
            xy(Parameters const& par, par_imw_p const& proj_parm, double phi, double *x, double *y, double *sp, double *R) {
                double F;

                *sp = sin(phi);
                *R = 1./(tan(phi) * sqrt(1. - par.es * *sp * *sp ));
                F = proj_parm.lam_1 * *sp;
                *y = *R * (1 - cos(F));
                *x = *R * sin(F);
            }

            // template class, using CRTP to implement forward/inverse
            template <typename Geographic, typename Cartesian, typename Parameters>
            struct base_imw_p_ellipsoid : public base_t_fi<base_imw_p_ellipsoid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>
            {

                 typedef double geographic_type;
                 typedef double cartesian_type;

                par_imw_p m_proj_parm;

                inline base_imw_p_ellipsoid(const Parameters& par)
                    : base_t_fi<base_imw_p_ellipsoid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>(*this, par) {}

                // FORWARD(e_forward)  ellipsoid
                // Project coordinates from geographic (lon, lat) to cartesian (x, y)
                inline void fwd(geographic_type& lp_lon, geographic_type& lp_lat, cartesian_type& xy_x, cartesian_type& xy_y) const
                {
                    double yc = 0;
                    XY xy = loc_for(lp_lon, lp_lat, this->m_par, m_proj_parm, &yc);
                    xy_x = xy.x; xy_y = xy.y;
                }

                // INVERSE(e_inverse)  ellipsoid
                // Project coordinates from cartesian (x, y) to geographic (lon, lat)
                inline void inv(cartesian_type& xy_x, cartesian_type& xy_y, geographic_type& lp_lon, geographic_type& lp_lat) const
                {
                    XY t;
                    double yc = 0;

                    lp_lat = this->m_proj_parm.phi_2;
                    lp_lon = xy_x / cos(lp_lat);
                    do {
                        t = loc_for(lp_lon, lp_lat, this->m_par, m_proj_parm, &yc);
                        lp_lat = ((lp_lat - this->m_proj_parm.phi_1) * (xy_y - yc) / (t.y - yc)) + this->m_proj_parm.phi_1;
                        lp_lon = lp_lon * xy_x / t.x;
                    } while (fabs(t.x - xy_x) > TOL || fabs(t.y - xy_y) > TOL);
                }

                static inline std::string get_name()
                {
                    return "imw_p_ellipsoid";
                }

            };

            // International Map of the World Polyconic
            template <typename Parameters>
            void setup_imw_p(Parameters& par, par_imw_p& proj_parm)
            {
                double del, sig, s, t, x1, x2, T2, y1, m1, m2, y2;
                int i;

                if (!pj_enfn(par.es, proj_parm.en)) throw proj_exception(0);
                if( (i = phi12(par, proj_parm, &del, &sig)) != 0)
                    throw proj_exception(i);
                if (proj_parm.phi_2 < proj_parm.phi_1) { /* make sure proj_parm.phi_1 most southerly */
                    del = proj_parm.phi_1;
                    proj_parm.phi_1 = proj_parm.phi_2;
                    proj_parm.phi_2 = del;
                }
                if (pj_param(par.params, "tlon_1").i)
                    proj_parm.lam_1 = pj_param(par.params, "rlon_1").f;
                else { /* use predefined based upon latitude */
                    sig = fabs(sig * geometry::math::r2d<double>());
                    if (sig <= 60)        sig = 2.;
                    else if (sig <= 76) sig = 4.;
                    else                sig = 8.;
                    proj_parm.lam_1 = sig * geometry::math::d2r<double>();
                }
                proj_parm.mode = 0;
                if (proj_parm.phi_1) xy(par, proj_parm, proj_parm.phi_1, &x1, &y1, &proj_parm.sphi_1, &proj_parm.R_1);
                else {
                    proj_parm.mode = 1;
                    y1 = 0.;
                    x1 = proj_parm.lam_1;
                }
                if (proj_parm.phi_2) xy(par, proj_parm, proj_parm.phi_2, &x2, &T2, &proj_parm.sphi_2, &proj_parm.R_2);
                else {
                    proj_parm.mode = -1;
                    T2 = 0.;
                    x2 = proj_parm.lam_1;
                }
                m1 = pj_mlfn(proj_parm.phi_1, proj_parm.sphi_1, cos(proj_parm.phi_1), proj_parm.en);
                m2 = pj_mlfn(proj_parm.phi_2, proj_parm.sphi_2, cos(proj_parm.phi_2), proj_parm.en);
                t = m2 - m1;
                s = x2 - x1;
                y2 = sqrt(t * t - s * s) + y1;
                proj_parm.C2 = y2 - T2;
                t = 1. / t;
                proj_parm.P = (m2 * y1 - m1 * y2) * t;
                proj_parm.Q = (y2 - y1) * t;
                proj_parm.Pp = (m2 * x1 - m1 * x2) * t;
                proj_parm.Qp = (x2 - x1) * t;
            }

        }} // namespace detail::imw_p
    #endif // doxygen

    /*!
        \brief International Map of the World Polyconic projection
        \ingroup projections
        \tparam Geographic latlong point type
        \tparam Cartesian xy point type
        \tparam Parameters parameter type
        \par Projection characteristics
         - Mod. Polyconic
         - Ellipsoid
        \par Projection parameters
         - lat_1: Latitude of first standard parallel
         - lat_2: Latitude of second standard parallel
         - lon_1 (degrees)
        \par Example
        \image html ex_imw_p.gif
    */
    template <typename Geographic, typename Cartesian, typename Parameters = parameters>
    struct imw_p_ellipsoid : public detail::imw_p::base_imw_p_ellipsoid<Geographic, Cartesian, Parameters>
    {
        inline imw_p_ellipsoid(const Parameters& par) : detail::imw_p::base_imw_p_ellipsoid<Geographic, Cartesian, Parameters>(par)
        {
            detail::imw_p::setup_imw_p(this->m_par, this->m_proj_parm);
        }
    };

    #ifndef DOXYGEN_NO_DETAIL
    namespace detail
    {

        // Factory entry(s)
        template <typename Geographic, typename Cartesian, typename Parameters>
        class imw_p_entry : public detail::factory_entry<Geographic, Cartesian, Parameters>
        {
            public :
                virtual projection<Geographic, Cartesian>* create_new(const Parameters& par) const
                {
                    return new base_v_fi<imw_p_ellipsoid<Geographic, Cartesian, Parameters>, Geographic, Cartesian, Parameters>(par);
                }
        };

        template <typename Geographic, typename Cartesian, typename Parameters>
        inline void imw_p_init(detail::base_factory<Geographic, Cartesian, Parameters>& factory)
        {
            factory.add_to_factory("imw_p", new imw_p_entry<Geographic, Cartesian, Parameters>);
        }

    } // namespace detail
    #endif // doxygen

}}} // namespace boost::geometry::projections

#endif // BOOST_GEOMETRY_PROJECTIONS_IMW_P_HPP


//---------------------------------------------------------------------------//
//  MIT License
//
//  Copyright (c) 2020-2021 Mikhail Komarov <nemo@nil.foundation>
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.
//---------------------------------------------------------------------------//

#ifndef FILECOIN_SEAL_API_POST_HPP
#define FILECOIN_SEAL_API_POST_HPP

#include <iostream>

#include <boost/filesystem/file_status.hpp>

#include <nil/filecoin/storage/proofs/core/btree/map.hpp>
#include <nil/filecoin/storage/proofs/core/proof/compound_proof.hpp>
#include <nil/filecoin/storage/proofs/core/cache_key.hpp>
#include <nil/filecoin/storage/proofs/core/sector.hpp>

#include <nil/filecoin/proofs/api/utilities.hpp>

#include <nil/filecoin/proofs/types/post_config.hpp>

#include <nil/filecoin/proofs/parameters.hpp>

namespace nil {
    namespace filecoin {
        /// The minimal information required about a replica, in order to be able to generate
        /// a PoSt over it.
        template<typename MerkleTreeType>
        struct PrivateReplicaInfo {
            PrivateReplicaInfo(const boost::filesystem::path &replica, const commitment_type &comm_r,
                               const boost::filesystem::path &cache_dir) :
                cache_dir(cache_dir),
                replica(replica), comm_r(comm_r) {
                assert(("Invalid all zero commitment (comm_r)",
                        !std::accumulate(comm_r.begin(), comm_r.end(), false,
                                         [&](bool state, typename commitment_type::value_type &v) -> bool {
                                             return state * (v != 0);
                                         })));

                boost::filesystem::path f_aux_path = cache_dir / std::to_string(cache_key::PAux);
                std::ifstream file;
                file.open(f_aux_path.string(), std::ios::binary);
                std::streamsize size = file.tellg();
                file.seekg(0, std::ios::beg);

                std::vector<char> aux_bytes(size);
                if (file.read(aux_bytes.data(), size)) {
                    aux = deserialize(aux_bytes);
                }

                assert(("Sealed replica does not exist", boost::filesystem::exists(replica.status())));
            }

            boost::filesystem::path cache_dir_path() const {
                return cache_dir;
            }

            boost::filesystem::path replica_path() const {
                return replica;
            }

            typename MerkleTreeType::hash_type::digest_type safe_comm_r() const {
                return as_safe_commitment(comm_r, "comm_r");
            }

            typename MerkleTreeType::hash_type::digest_type safe_comm_c() const {
                return aux.comm_c;
            }

            typename MerkleTreeType::hash_type::digest_type safe_comm_r_last() const {
                return aux.comm_r_last;
            }

            /// Generate the merkle tree of this particular replica.
            MerkleTreeWrapper<typename MerkleTreeType::hash_type, MerkleTreeType::Store, MerkleTreeType::base_arity,
                              MerkleTreeType::sub_tree_arity, MerkleTreeType::top_tree_arity>
                merkle_tree(sector_size_type sector_size) {
                std::size_t base_tree_size = get_base_tree_size<MerkleTreeType>(sector_size);
                std::size_t base_tree_leafs = get_base_tree_leafs<MerkleTreeType>(base_tree_size);

                StoreConfg config(cache_dir_path(), std::to_string(cache_key::CommRLastTree),
                                  default_rows_to_discard(base_tree_leafs, MerkleTreeType::base_arity));
                config.size = base_tree_size;

                std::size_t tree_count = get_base_tree_count<MerkleTreeType>();
                let(configs, replica_config) =
                    split_config_and_replica(config, replica_path().to_path_buf(), base_tree_leafs, tree_count);

                return create_tree<MerkleTreeType>(base_tree_size, configs, replica_config);
            }

            /// Path to the replica.
            boost::filesystem::path replica;
            /// The replica commitment.
            commitment_type comm_r;
            /// Persistent Aux.
            PersistentAux<typename MerkleTreeType::hash_type::digest_type> aux;
            /// Contains sector-specific (e.g. merkle trees) assets
            boost::filesystem::path cache_dir;
        };    // namespace filecoin
    }         // namespace filecoin
}    // namespace nil

namespace std {
    template<typename MerkleTreeType>
    struct hash<nil::filecoin::PrivateReplicaInfo<MerkleTreeType>> {
        int operator()(const nil::filecoin::PrivateReplicaInfo<MerkleTreeType> &v) {
            return hash()(v.replica) ^ hash()(v.comm_r) ^ hash()(v.aux) ^ hash()(v.cache_dir);
        }
    };
}    // namespace std

namespace nil {
    namespace filecoin {

        /// The minimal information required about a replica, in order to be able to verify
        /// a PoSt over it.
        struct PublicReplicaInfo {
            PublicReplicaInfo(const commitment_type &comm_r);

            template<typename Domain>
            Domain safe_comm_r() const {
                return as_safe_commitment(comm_r, "comm_r");
            }

            /// The replica commitment.
            commitment_type comm_r;
        };

        // Ensure that any associated cached data persisted is discarded.
        template<typename MerkleTreeType>
        void clear_cache(const boost::filesystem::path &cache_dir) {
            info !("clear_cache:start");

            TemporaryAux<MerkleTreeType> t_aux;
            boost::filesystem::path f_aux_path = cache_dir / std::to_string(cache_key::TAux);
            std::vector<std::uint8_t> aux_bytes =
                std::fs::read(&f_aux_path).with_context(|| format !("could not read from path={:?}", f_aux_path));

            deserialize(aux_bytes);

            TemporaryAux<MerkleTreeType> result = TemporaryAux<MerkleTreeType, DefaultPieceHasher>::clear_temp(t_aux);

            return result;
        }    // namespace filecoin

        // Ensure that any associated cached data persisted is discarded.
        template<typename MerkleTreeType>
        void clear_caches(const btree::map<sector_id_type, PrivateReplicaInfo<MerkleTreeType>> &replicas) {
            for (const typename btree::map<sector_id_type, PrivateReplicaInfo<MerkleTreeType>>::value_type &replica :
                 replicas) {
                clear_cache<MerkleTreeType>(replica.second().cache_dir.as_path());
            }
        }

        typedef std::vector<std::uint8_t> SnarkProof;

        /// Generates a Winning proof-of-spacetime.
        template<typename MerkleTreeType>
        SnarkProof generate_winning_post(const post_config &config, const challenge_seed_type &randomness,
                                         const btree::map<sector_id_type, PrivateReplicaInfo<MerkleTreeType>> &replicas,
                                         prover_id_type prover_id) {
            assert(("invalid post config type", config.typ == post_type::Winning));
            assert(("invalid amount of replicas", replicas.size() == post_config.sector_count));

            typename MerkleTreeType::hash_type::digest_type randomness_safe =
                as_safe_commitment(randomness, "randomness");
            typename MerkleTreeType::hash_type::digest_type prover_id_safe =
                as_safe_commitment(&prover_id, "prover_id");

            WinningPostSetupParams vanilla_params = winning_post_setup_params(config);
            std::size_t param_sector_count = vanilla_params.sector_count;

            compound_proof::SetupParams setup_params {vanilla_params, partitions : None, config.priority};

            compound_proof::PublicParams<fallback::FallbackPoSt<MerkleTreeType>> pub_params =
                fallback::FallbackPoStCompound::setup(setup_params);
            let groth_params = get_post_params<MerkleTreeType>(config);

            for (typename btree::map<sector_id_type, PrivateReplicaInfo<MerkleTreeType>>)

                let trees = replicas.iter()
                                .map(| (_, replica) | replica.merkle_tree(config.sector_size))
                                .collect::<Result<Vec<_>>>();

            std::vector<fallback::PublicSector> pub_sectors(param_sector_count);
            std::vector<fallback::PrivateSector> priv_sectors(param_sector_count);

            for (int i = 0; i < param_sector_count; i++) {
                for (((id, replica), tree) : replicas.iter().zip(trees.iter())) {
                    typename MerkleTreeType::hash_type::digest_type comm_r = replica.safe_comm_r();
                    typename MerkleTreeType::hash_type::digest_type comm_c = replica.safe_comm_c();
                    typename MerkleTreeType::hash_type::digest_type comm_r_last = replica.safe_comm_r_last();

                    pub_sectors.push_back(
                        fallback::PublicSector<typename MerkleTreeType::hash_type::digest_type> {id : *id, comm_r});
                    priv_sectors.push_back(fallback::PrivateSector {tree, comm_c, comm_r_last});
                }
            }

            fallback::PublicInputs<typename MerkleTreeType::hash_type::digest_type> pub_inputs =
                {randomness_safe, prover_id_safe, pub_sectors, k : None};

            fallback::PrivateInputs<MerkleTreeType> priv_inputs = {priv_sectors};

            let proof = fallback::FallbackPoStCompound<MerkleTreeType>::prove(pub_params, pub_inputs, priv_inputs,
                                                                              groth_params);
            let proof = proof.to_vec();

            return proof;
        }

        /// Given some randomness and a the length of available sectors, generates the challenged sector.
        ///
        /// The returned values are indicies in the range of `0..sector_set_size`, requiring the caller
        /// to match the index to the correct sector.
        template<typename MerkleTreeType>
        std::vector<std::uint64_t>
            generate_winning_post_sector_challenge(const post_config &config, const challenge_seed_type &randomness,
                                                   std::uint64_t sector_set_size, const commitment_type &prover_id) {
            assert(sector_set_size != 0, "empty sector set is invalid");
            assert(post_config.typ == PoStType::Winning, "invalid post config type");

            typename MerkleTreeType::hash_type::digest_type prover_id_safe = as_safe_commitment(prover_id, "prover_id");

            typename MerkleTreeType::hash_type::digest_type randomness_safe =
                as_safe_commitment(randomness, "randomness");
            std::vector<std::uint64_t> result = fallback::generate_sector_challenges(
                randomness_safe, config.sector_count, sector_set_size, prover_id_safe);

            return result;
        }

        /// Verifies a winning proof-of-spacetime.
        ///
        /// The provided `replicas` must be the same ones as passed to `generate_winning_post`, and be based on
        /// the indices generated by `generate_winning_post_sector_challenge`. It is the responsibility of the
        /// caller to ensure this.
        template<typename MerkleTreeType>
        bool verify_winning_post(const post_config &config, const challenge_seed_type &randomness,
                                 const btree::map<sector_id_type, PublicReplicaInfo> &replicas,
                                 prover_id_type prover_id, const std::vector<std::uint8_t> &proof) {
            assert(("invalid post config type", config.typ == PoStType::Winning));
            assert(("invalid amount of replicas provided", config.sector_count == replicas.size()));

            typename MerkleTreeType::hash_type::digest_type randomness_safe =
                as_safe_commitment(randomness, "randomness");
            typename MerkleTreeType::hash_type::digest_type prover_id_safe = as_safe_commitment(prover_id, "prover_id");

            WinningPostSetupParams vanilla_params = winning_post_setup_params(config);
            std::size_t param_sector_count = vanilla_params.sector_count;

            compound_proof::SetupParams setup_params = {vanilla_params, partitions : None, priority : false};
            compound_proof::PublicParams<fallback::FallbackPoSt<MerkleTreeType>> pub_params =
                fallback::FallbackPoStCompound::setup(&setup_params);

            let verifying_key = get_post_verifying_key<MerkleTreeType>(config);

            MultiProof proof = MultiProof::new_from_reader(None, &proof[..], &verifying_key);
            if (proof.size() != 1) {
                return false;
            }

            std::vector<fallback::PublicSector> pub_sectors(param_sector_count);
            for (int i = 0; i < param_sector_count; i++) {
                for ((id, replica) : replicas.iter()) {
                    typename MerkleTreeType::hash_type::digest_type comm_r = replica.safe_comm_r();
                    pub_sectors.push_back({*id, comm_r});
                }
            }

            fallback::PublicInputs pub_inputs =
                {randomness : randomness_safe, prover_id : prover_id_safe, sectors : pub_sectors, k : None};

            bool is_valid = fallback::FallbackPoStCompound::verify(pub_params, pub_inputs, proof,
                                                                   {config.challenge_count * config.sector_count});

            if (!is_valid) {
                return false;
            }

            return true;
        }

        /// Generates a Window proof-of-spacetime.
        template<typename MerkleTreeType>
        SnarkProof generate_window_post(const post_config &config, const challenge_seed_type &randomness,
                                        const btree::map<sector_id_type, PrivateReplicaInfo<MerkleTreeType>> &replicas,
                                        prover_id_type prover_id) {
            assert(("invalid post config type", post_config.typ == PoStType::Window));

            typename MerkleTreeType::hash_type::digest_type randomness_safe =
                as_safe_commitment(randomness, "randomness");
            typename MerkleTreeType::hash_type::digest_type prover_id_safe = as_safe_commitment(prover_id, "prover_id");

            let vanilla_params = window_post_setup_params(config);
            let partitions = get_partitions_for_window_post(replicas.size(), config);

            std::size_t sector_count = vanilla_params.sector_count;
            compound_proof::SetupParams setup_params = {vanilla_params, partitions, priority : config.priority};

            compound_proof::PublicParams<fallback::FallbackPoSt<MerkleTreeType>> pub_params =
                fallback::FallbackPoStCompound::setup(setup_params);
            let groth_params = get_post_params<MerkleTreeType>(config);

            std::vector<MerkleTreeType> trees =
                replicas.iter().map(| (_id, replica) | replica.merkle_tree(config.sector_size)).collect::<Result<_>>();

            std::vector<fallback::PublicSector> pub_sectors(sector_count);
            std::vector<fallback::PrivateSector> priv_sectors(sector_count);

            for (((sector_id, replica), tree) : replicas.iter().zip(trees.iter())) {
                typename MerkleTreeType::hash_type::digest_type comm_r = replica.safe_comm_r();
                typename MerkleTreeType::hash_type::digest_type comm_c = replica.safe_comm_c();
                typename MerkleTreeType::hash_type::digest_type comm_r_last = replica.safe_comm_r_last();

                pub_sectors.push_back(fallback::PublicSector {id : *sector_id, comm_r});
                priv_sectors.push_back(fallback::PrivateSector {tree, comm_c, comm_r_last});
            }

            fallback::PublicInputs pub_inputs =
                {randomness : randomness_safe, prover_id : prover_id_safe, sectors : pub_sectors, k : None};

            fallback::PrivateInputs<MerkleTreeType> priv_inputs = {sectors : priv_sectors};

            let proof = fallback::FallbackPoStCompound::prove(&pub_params, &pub_inputs, &priv_inputs, &groth_params, );

            return proof.to_vec();
        }

        /// Verifies a window proof-of-spacetime.
        template<typename MerkleTreeType>
        bool verify_window_post(const post_config &config,
                                const challenge_seed_type &randomness,
                                const btree::map<sector_id_type, PublicReplicaInfo> &replicas,
                                prover_id_type prover_id,
                                const std::vector<std::uint8_t> &proof) {
            assert(("invalid post config type", post_config.typ == PoStType::Window));

            typename MerkleTreeType::hash_type::digest_type randomness_safe =
                as_safe_commitment(randomness, "randomness");
            typename MerkleTreeType::hash_type::digest_type prover_id_safe = as_safe_commitment(prover_id, "prover_id");

            WindowPostSetupParams vanilla_params = window_post_setup_params(config);
            let partitions = get_partitions_for_window_post(replicas.size(), config);

            compound_proof::SetupParams setup_params = {vanilla_params, partitions, priority : false};
            compound_proof::PublicParams<fallback::FallbackPoSt<Tree>> pub_params =
                fallback::FallbackPoStCompound::setup(setup_params);

            let verifying_key = get_post_verifying_key<MerkleTreeType>(config);

            MultiProof proof = MultiProof::new_from_reader(partitions, &proof[..], &verifying_key);

            std::vector<PublicSector> pub_sectors = replicas.iter()
                                                        .map(| (sector_id, replica) |
                                                             {
                                                                 let comm_r = replica.safe_comm_r() ? ;
                                                                 Ok(fallback::PublicSector {
                                                                     id : *sector_id,
                                                                     comm_r,
                                                                 })
                                                             })
                                                        .collect::<Result<_>>();

            fallback::PublicInputs pub_inputs = {
                .randomness = randomness_safe, .prover_id = prover_id_safe, .sectors = pub_sectors, .k = None};

            bool is_valid =
                fallback::FallbackPoStCompound::verify(pub_params, pub_inputs, proof, fallback::ChallengeRequirements {
                    minimum_challenge_count : post_config.challenge_count * post_config.sector_count
                });

            if (!is_valid) {
                return false;
            }

            return true;
        }

        boost::optional<std::size_t> get_partitions_for_window_post(std::size_t total_sector_count,
                                                                    const post_config &config);
    }    // namespace filecoin
}    // namespace nil

#endif    // FILECOIN_SEAL_HPP


#include <Eigen/Dense>
#include "space.h"
#include "reactioncoordinate.h"

namespace Faunus {
namespace ReactionCoordinate {

void ReactionCoordinateBase::_to_json(json &) const {}

double ReactionCoordinateBase::normalize(double) const { return 1.; }

double ReactionCoordinateBase::operator()() {
    assert(f != nullptr);
    return f();
}

bool ReactionCoordinateBase::inRange(double coord) const { return (coord >= min && coord <= max); }

void to_json(json &j, const ReactionCoordinateBase &r) {
    j = {{"range", {r.min, r.max}}, {"resolution", r.binwidth}};
    r._to_json(j);
}

void from_json(const json &j, ReactionCoordinateBase &r) {
    r.binwidth = j.value("resolution", 0.5);
    auto range = j.value("range", std::vector<double>({0, 0}));
    if (range.size() == 2)
        if (range[0] <= range[1]) {
            r.min = range[0];
            r.max = range[1];
            return;
        }
    throw std::runtime_error(r.name + ": 'range' require two numbers: [min, max>=min]");
}

void SystemProperty::_to_json(json &j) const { j["property"] = property; }
SystemProperty::SystemProperty(const json &j, Space &spc) {
    name = "system";
    from_json(j, *this);
    property = j.at("property").get<std::string>();
    if (property == "V")
        f = [&g = spc.geo]() { return g.getVolume(); };
    else if (property == "Lx")
        f = [&g = spc.geo]() { return g.getLength().x(); };
    else if (property == "Ly")
        f = [&g = spc.geo]() { return g.getLength().y(); };
    else if (property == "Lz" or property == "height")
        f = [&g = spc.geo]() { return g.getLength().z(); };
    else if (property == "radius") {
        if (spc.geo.type == Geometry::CUBOID or spc.geo.type == Geometry::SLIT)
            std::cerr << "`radius` coordinate unavailable for geometry" << endl;
        else
            f = [&g = spc.geo]() { return 0.5 * g.getLength().x(); };
    } else if (property == "Q") // system net charge
        f = [&groups = spc.groups]() {
            double charge_sum = 0;
            for (auto &g : groups) // loops over groups
                for (auto &p : g)  // loop over particles
                    charge_sum += p.charge;
            return charge_sum;
        };
    else if (property == "N") // number of particles
        f = [&groups = spc.groups]() {
            int N_sum = 0;
            for (auto &g : groups) // loops over groups
                N_sum += g.size();
            return N_sum;
        };
    if (f == nullptr)
        throw std::runtime_error(name + ": unknown property '" + property + "'" + usageTip["coords=[system]"]);
}

void AtomProperty::_to_json(json &j) const {
    j["property"] = property;
    j["index"] = index;
    if (dir.squaredNorm() > 1e-9)
        j["dir"] = dir;
}
AtomProperty::AtomProperty(const json &j, Space &spc) {
    name = "atom";
    from_json(j, *this);
    index = j.at("index");
    property = j.at("property").get<std::string>();
    if (property == "x")
        f = [&p = spc.p, i = index]() { return p[i].pos.x(); };
    else if (property == "y")
        f = [&p = spc.p, i = index]() { return p[i].pos.y(); };
    else if (property == "z")
        f = [&p = spc.p, i = index]() { return p[i].pos.z(); };
    else if (property == "R")
        f = [&p = spc.p, i = index]() { return p[i].pos.norm(); };
    else if (property == "q")
        f = [&p = spc.p, i = index]() { return p[i].charge; };
    else if (property == "N") // number of atom of id=index
        f = [&groups = spc.groups, i = index]() {
            int N_sum = 0;
            for (auto &g : groups) // loops over groups
                for (auto &p : g)  // loops over particles
                    if (p.id == i)
                        N_sum++;
            return N_sum;
        };
    if (f == nullptr)
        throw std::runtime_error(name + ": unknown property '" + property + "'" + usageTip["coords=[atom]"]);
}

void MoleculeProperty::_to_json(json &j) const {
    j["property"] = property;
    j["index"] = index;
    if (dir.squaredNorm() > 1e-9)
        j["dir"] = dir;
    if (indexes.size() >= 2)
        j["indexes"] = indexes;
}
MoleculeProperty::MoleculeProperty(const json &j, Space &spc) {
    typedef typename Space::Tparticle Tparticle;
    name = "molecule";
    from_json(j, *this);
    index = j.value("index", 0);
    auto b = spc.geo.getBoundaryFunc();
    property = j.at("property").get<std::string>();

    if (property == "confid")
        f = [&g = spc.groups, i = index]() { return g[i].confid; };
    else if (property == "com_x")
        f = [&g = spc.groups, i = index]() { return g[i].cm.x(); };
    else if (property == "com_y")
        f = [&g = spc.groups, i = index]() { return g[i].cm.y(); };
    else if (property == "com_z")
        f = [&g = spc.groups, i = index]() { return g[i].cm.z(); };
    else if (property == "N")
        f = [&g = spc.groups, i = index]() { return g[i].size(); };
    else if (property == "Q")
        f = [&g = spc.groups, i = index]() { return Geometry::monopoleMoment(g[i].begin(), g[i].end()); };

    else if (property == "mu_x")
        f = [&g = spc.groups, i = index, b]() { return Geometry::dipoleMoment(g[i].begin(), g[i].end(), b).x(); };

    else if (property == "mu_y")
        f = [&g = spc.groups, i = index, b]() { return Geometry::dipoleMoment(g[i].begin(), g[i].end(), b).y(); };

    else if (property == "mu_z")
        f = [&g = spc.groups, i = index, b]() { return Geometry::dipoleMoment(g[i].begin(), g[i].end(), b).z(); };

    else if (property == "mu")
        f = [&g = spc.groups, i = index, b]() { return Geometry::dipoleMoment(g[i].begin(), g[i].end(), b).norm(); };

    else if (property == "end2end")
        f = [&spc, i = index]() {
            assert(spc.groups[i].size() > 1);
            return std::sqrt(spc.geo.sqdist(spc.groups[i].begin()->pos, (spc.groups[i].end() - 1)->pos));
        };

    else if (property == "muangle") {
        dir = j.at("dir").get<Point>().normalized();
        if (not spc.groups.at(index).atomic)
            f = [&g = spc.groups, i = index, b, &dir = dir]() {
                Point mu = Geometry::dipoleMoment(g[i].begin(), g[i].end(), b);
                return std::acos(mu.dot(dir)) * 180 / pc::pi;
            };
    }

    else if (property == "atomatom") {
        dir = j.at("dir");
        indexes = j.at("indexes").get<decltype(indexes)>();
        if (indexes.size() != 2)
            throw std::runtime_error("exactly two indices expected");
        f = [&spc, &dir = dir, i = indexes[0], j = indexes[1]]() {
            auto &pos1 = spc.p.at(i).pos;
            auto &pos2 = spc.p.at(j).pos;
            return spc.geo.vdist(pos1, pos2).cwiseProduct(dir.cast<double>()).norm();
        };
    }

    else if (property == "cmcm_z") {
        indexes = j.value("indexes", decltype(indexes)());
        if (indexes.size() != 4)
            throw std::runtime_error("exactly four indices expected");
        f = [&spc, dir = dir, i = indexes[0], j = indexes[1] + 1, k = indexes[2], l = indexes[3] + 1]() {
            auto cm1 = Geometry::massCenter(spc.p.begin() + i, spc.p.begin() + j, spc.geo.getBoundaryFunc());
            auto cm2 = Geometry::massCenter(spc.p.begin() + k, spc.p.begin() + l, spc.geo.getBoundaryFunc());
            return spc.geo.vdist(cm1, cm2).z();
        };
    }

    else if (property == "cmcm") {
        dir = j.at("dir");
        indexes = j.value("indexes", decltype(indexes)());
        assert(indexes.size() == 4 && "An array of 4 indexes should be specified.");
        f = [&spc, dir = dir, i = indexes[0], j = indexes[1] + 1, k = indexes[2], l = indexes[3] + 1]() {
            auto cm1 = Geometry::massCenter(spc.p.begin() + i, spc.p.begin() + j, spc.geo.getBoundaryFunc());
            auto cm2 = Geometry::massCenter(spc.p.begin() + k, spc.p.begin() + l, spc.geo.getBoundaryFunc());
            return spc.geo.vdist(cm1, cm2).cwiseProduct(dir.cast<double>()).norm();
        };
    }

    else if (property == "L/R") {
        dir = j.at("dir");
        indexes = j.value("indexes", decltype(indexes)());
        assert(indexes.size() == 2 && "An array of 2 indexes should be specified.");
        f = [&spc, &dir = dir, i = indexes[0], j = indexes[1]]() {
            Average<double> Rj, Rin, Rout;
            Group<Tparticle> g(spc.p.begin(), spc.p.end());
            auto slicei = g.find_id(i);
            auto cm = Geometry::massCenter(slicei.begin(), slicei.end(), spc.geo.getBoundaryFunc());
            auto slicej = g.find_id(j);
            for (auto p : slicej)
                Rj += spc.geo.vdist(p.pos, cm).cwiseProduct(dir.cast<double>()).norm();
            double Rjavg = Rj.avg();
            for (auto p : slicei) {
                double d = spc.geo.vdist(p.pos, cm).cwiseProduct(dir.cast<double>()).norm();
                if (d < Rjavg)
                    Rin += d;
                else if (d > Rjavg)
                    Rout += d;
            }
            return 2 * spc.geo.getLength().z() / (Rin.avg() + Rout.avg());
        };
    }

    else if (property == "angle") {
        dir = j.at("dir").get<Point>().normalized();
        if (not spc.groups.at(index).atomic) {
            f = [&spc, &dir = dir, i = index]() {
                auto &cm = spc.groups[i].cm;
                auto S = Geometry::gyration(spc.groups[i].begin(), spc.groups[i].end(), spc.geo.getBoundaryFunc(), cm);
                Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> esf(S);
                Point eivals = esf.eigenvalues();
                std::ptrdiff_t i_eival;
                eivals.minCoeff(&i_eival);
                Point vec = esf.eigenvectors().col(i_eival).real();
                double cosine = vec.dot(dir);
                double angle = std::acos(std::fabs(cosine)) * 180. / pc::pi;
                return angle;
            };
        }
    }

    if (f == nullptr)
        throw std::runtime_error(name + ": unknown or impossible property '" + property + "'" +
                                 usageTip["coords=[molecule]"]);
}

double MassCenterSeparation::normalize(double coord) const {
    int dim = dir.sum();
    if (dim == 2)
        return 1 / (2 * pc::pi * coord);
    if (dim == 3)
        return 1 / (4 * pc::pi * coord * coord);
    return 1.0;
}

void MassCenterSeparation::_to_json(json &j) const {
    j["dir"] = dir;
    j["indexes"] = indexes;
    j["type"] = type;
}
MassCenterSeparation::MassCenterSeparation(const json &j, Space &spc) {
    typedef typename Space::Tparticle Tparticle;
    name = "cmcm";
    from_json(j, *this);
    dir = j.value("dir", dir);
    indexes = j.value("indexes", decltype(indexes)());
    type = j.value("type", decltype(type)());
    if (indexes.size() == 4) {
        f = [&spc, dir = dir, i = indexes[0], j = indexes[1] + 1, k = indexes[2], l = indexes[3] + 1]() {
            auto cm1 = Geometry::massCenter(spc.p.begin() + i, spc.p.begin() + j, spc.geo.getBoundaryFunc());
            auto cm2 = Geometry::massCenter(spc.p.begin() + k, spc.p.begin() + l, spc.geo.getBoundaryFunc());
            return spc.geo.vdist(cm1, cm2).cwiseProduct(dir.cast<double>()).norm();
        };
    } else if (type.size() == 2) {
        f = [&spc, dir = dir, type1 = type[0], type2 = type[1]]() {
            Group<Tparticle> g(spc.p.begin(), spc.p.end());
            auto slice1 = g.find_id(findName(atoms, type1)->id());
            auto slice2 = g.find_id(findName(atoms, type2)->id());
            auto cm1 = Geometry::massCenter(slice1.begin(), slice1.end(), spc.geo.getBoundaryFunc());
            auto cm2 = Geometry::massCenter(slice2.begin(), slice2.end(), spc.geo.getBoundaryFunc());
            return spc.geo.vdist(cm1, cm2).cwiseProduct(dir.cast<double>()).norm();
        };
    } else
        throw std::runtime_error(name + ": specify 4 indexes or two atom types");
}

} // namespace ReactionCoordinate
} // namespace Faunus


/**
 * Copyright (c) 2019, Arjan van der Velde, Weng Lab
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#pragma once

#include "Exception.hpp"
#include "TransformLigand.hpp"
#include "TransformMultimer.hpp"
#include "ZDOCK.hpp"
#include <Eigen/Dense>
#include <string>

namespace zdock {

/**
 * @brief Perform RMSD based pruning on (M-)ZDOCK output
 */
class Pruning {
private:
  typedef Eigen::Transform<double, 3, Eigen::Affine> Transform;
  typedef Eigen::Matrix<double, 3, Eigen::Dynamic> Matrix;

  ZDOCK zdock_;                 // zdock output
  const double cutoff_;         // cutoff
  const TransformLigand txl_;   // ligand tranfomation class
  const TransformMultimer txm_; // multimertranfomation class
  std::string strucfn_;         // receptor and ligand filenames
  const bool getclusters_;      // return all w/ cluster number in score

  // results
  std::vector<int> clusters_; // cluster assignments
  size_t strucsize_;          // structure size
  int nclusters_;             // number of clusters

public:
  /**
   * @brief Constructor
   *
   * @param zdockoutput ZDOCK or M-ZDOCK output file name
   * @param cutoff RMSD cutoff
   * @param structurefn Structure PDB file name
   * @param getclusters Toggle return for full (M-)ZDOCK output with cluster numbers for scores
   */
  Pruning(
      const std::string &zdockoutput, const double cutoff,
      const std::string &structurefn = "", // or grab from zdock.out
      const bool getclusters = false // return all w/ cluster number in score
  );

  /**
   * @brief Actually perform pruning
   */
  void prune();
  /**
   * @brief Get cluster assignments
   *
   * @return vector of cluster numbers, one for each prediction
   */
  const std::vector<int> &clusters() const { return clusters_; }
  /**
   * @brief Get number of clusters
   *
   * @return number of clusters found
   */
  int nclusters() const { return nclusters_; }
  /**
   * @brief Get ZDOCK output with cluster numbers for scores
   *
   * @return ZDOCK output with cluster numbers for scores
   */
  const ZDOCK &zdock() const { return zdock_; }
};

/**
 * @brief General exception during pruning
 */
class PruningException : public Exception {
public:
  PruningException(const std::string &msg) : Exception(msg) {}
};

} // namespace zdock


#pragma once

#include <Epetra_CrsMatrix.h>
#include <Epetra_MultiVector.h>
#include <Epetra_Operator.h>
#include <deal.II/base/conditional_ostream.h>
#include "aux/message.hpp"
#include "aux/timer.hpp"
#include "boundary_conditions.hpp"

namespace boltzmann {
namespace otf_bc {

template <typename METHOD, typename APP>
class SystemMatrix : public ::Epetra_Operator
{
 private:
  typedef ::boltzmann::BoundaryConditions<METHOD, APP, impl::BdFacesManager> B_t;

 public:
  SystemMatrix(const ::Epetra_CrsMatrix& A, const B_t& B)
      : A_(A)
      , B_(B)
      , pcout(std::cout, ::dealii::Utilities::MPI::this_mpi_process(MPI_COMM_WORLD) == 0)
  {
    /* empty */
  }

  int Apply(const ::Epetra_MultiVector& X, ::Epetra_MultiVector& Y) const;
  int ApplyInverse(const ::Epetra_MultiVector& X, ::Epetra_MultiVector& Y) const;
  double NormInf() const;
  const char* Label() const;
  int SetUseTranspose(bool f);
  bool UseTranspose() const;
  bool HasNormInf() const;
  const ::Epetra_Comm& Comm() const;
  const ::Epetra_Map& OperatorDomainMap() const;
  const ::Epetra_Map& OperatorRangeMap() const;

 private:
  const ::Epetra_CrsMatrix& A_;
  const B_t& B_;

  mutable Timer<> timer;
  mutable ::dealii::ConditionalOStream pcout;
};

template <typename METHOD, typename APP>
int
SystemMatrix<METHOD, APP>::Apply(const ::Epetra_MultiVector& X, ::Epetra_MultiVector& Y) const
{
  //  timer.start();
  int retA = A_.Apply(X, Y);
  //  print_timer(timer.stop(), "A*x", pcout);

  // timer.start();
  B_.apply(Y, X);
  // print_timer(timer.stop(), "B*x", pcout);

  return retA;
}

template <typename METHOD, typename APP>
int
SystemMatrix<METHOD, APP>::ApplyInverse(const ::Epetra_MultiVector& X,
                                        ::Epetra_MultiVector& Y) const
{
  throw std::runtime_error("not implemented");
}

template <typename METHOD, typename APP>
double
SystemMatrix<METHOD, APP>::NormInf() const
{
  throw std::runtime_error("not implemented");
}

template <typename METHOD, typename APP>
bool
SystemMatrix<METHOD, APP>::HasNormInf() const
{
  return false;
}

template <typename METHOD, typename APP>
const char*
SystemMatrix<METHOD, APP>::Label() const
{
  return A_.Label();
}

template <typename METHOD, typename APP>
int
SystemMatrix<METHOD, APP>::SetUseTranspose(bool f)
{
  std::runtime_error("not implemented");
  return -1;
}

template <typename METHOD, typename APP>
bool
SystemMatrix<METHOD, APP>::UseTranspose() const
{
  return false;
}

template <typename METHOD, typename APP>
const ::Epetra_Comm&
SystemMatrix<METHOD, APP>::Comm() const
{
  return A_.Comm();
}

template <typename METHOD, typename APP>
const ::Epetra_Map&
SystemMatrix<METHOD, APP>::OperatorDomainMap() const
{
  return A_.OperatorDomainMap();
}

template <typename METHOD, typename APP>
const ::Epetra_Map&
SystemMatrix<METHOD, APP>::OperatorRangeMap() const
{
  return A_.OperatorRangeMap();
}

}  // end namespace otf_bc
}  // end namespace boltzmann


#include <Eigen/Core>
#include <memory>
#include <unordered_map>
#include <vector>

#include "MarchingCubesConst.h"
#include "VDBVolume.h"

namespace openvdb {
static const openvdb::Coord shift[8] = {
    openvdb::Coord(0, 0, 0), openvdb::Coord(1, 0, 0), openvdb::Coord(1, 1, 0),
    openvdb::Coord(0, 1, 0), openvdb::Coord(0, 0, 1), openvdb::Coord(1, 0, 1),
    openvdb::Coord(1, 1, 1), openvdb::Coord(0, 1, 1),
};
}

// Taken from <open3d/utility/Eigen.h>
namespace {
template <typename T>
struct hash_eigen {
    std::size_t operator()(T const& matrix) const {
        size_t seed = 0;
        for (int i = 0; i < (int)matrix.size(); i++) {
            auto elem = *(matrix.data() + i);
            seed ^= std::hash<typename T::Scalar>()(elem) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
        }
        return seed;
    }
};
}  // namespace

namespace vdbfusion {
std::tuple<std::vector<Eigen::Vector3d>, std::vector<Eigen::Vector3i>>
VDBVolume::ExtractTriangleMesh(bool fill_holes, float min_weight) const {
    // implementation of marching cubes, based on Open3D
    std::vector<Eigen::Vector3d> vertices;
    std::vector<Eigen::Vector3i> triangles;

    double half_voxel_length = voxel_size_ * 0.5;
    // Map of "edge_index = (x, y, z, 0) + edge_shift" to "global vertex index"
    std::unordered_map<Eigen::Vector4i, int, hash_eigen<Eigen::Vector4i>, std::equal_to<>,
                       Eigen::aligned_allocator<std::pair<const Eigen::Vector4i, int>>>
        edgeindex_to_vertexindex;
    int edge_to_index[12];

    auto tsdf_acc = tsdf_->getAccessor();
    auto weights_acc = weights_->getAccessor();
    for (auto iter = tsdf_->beginValueOn(); iter; ++iter) {
        int cube_index = 0;
        float f[8];
        const openvdb::Coord& voxel = iter.getCoord();
        const int32_t x = voxel.x();
        const int32_t y = voxel.y();
        const int32_t z = voxel.z();
        for (int i = 0; i < 8; i++) {
            openvdb::Coord idx = voxel + openvdb::shift[i];
            if (!fill_holes) {
                if (weights_acc.getValue(idx) == 0.0f) {
                    cube_index = 0;
                    break;
                }
            }
            if (weights_acc.getValue(idx) < min_weight) {
                cube_index = 0;
                break;
            }
            f[i] = tsdf_acc.getValue(idx);
            if (f[i] < 0.0f) {
                cube_index |= (1 << i);
            }
        }
        if (cube_index == 0 || cube_index == 255) {
            continue;
        }
        for (int i = 0; i < 12; i++) {
            if ((edge_table[cube_index] & (1 << i)) != 0) {
                Eigen::Vector4i edge_index = Eigen::Vector4i(x, y, z, 0) + edge_shift[i];
                if (edgeindex_to_vertexindex.find(edge_index) == edgeindex_to_vertexindex.end()) {
                    edge_to_index[i] = (int)vertices.size();
                    edgeindex_to_vertexindex[edge_index] = (int)vertices.size();
                    Eigen::Vector3d pt(half_voxel_length + voxel_size_ * edge_index(0),
                                       half_voxel_length + voxel_size_ * edge_index(1),
                                       half_voxel_length + voxel_size_ * edge_index(2));
                    double f0 = std::abs((double)f[edge_to_vert[i][0]]);
                    double f1 = std::abs((double)f[edge_to_vert[i][1]]);
                    pt(edge_index(3)) += f0 * voxel_size_ / (f0 + f1);
                    vertices.push_back(pt /* + origin_*/);
                } else {
                    edge_to_index[i] = edgeindex_to_vertexindex.find(edge_index)->second;
                }
            }
        }
        for (int i = 0; tri_table[cube_index][i] != -1; i += 3) {
            triangles.emplace_back(edge_to_index[tri_table[cube_index][i]],
                                   edge_to_index[tri_table[cube_index][i + 2]],
                                   edge_to_index[tri_table[cube_index][i + 1]]);
        }
    }
    return std::make_tuple(vertices, triangles);
}

}  // namespace vdbfusion