#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

namespace tensorflow {

using shape_inference::DimensionHandle;
using shape_inference::InferenceContext;
using shape_inference::ShapeHandle;
using shape_inference::UnchangedShape;

REGISTER_OP("Zhj")
    .Input("input: float")
    .Output("output: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
        DimensionHandle batch = c->Dim(c->input(0), 0);
        ShapeHandle output_shape = c->MakeShape({batch,c->MakeDim(512),c->MakeDim(2)});
        c->set_output(0, output_shape);
        return Status::OK();
    });

class PyFuncZhj : public tensorflow::OpKernel {
 public:
  explicit PyFuncZhj(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& tensor_in = context->input(0);
    const TensorShape& input_shape = tensor_in.shape();
    auto inputEigen = tensor_in.shaped<float, 2>({input_shape.dim_size(0), input_shape.dim_size(1)});
    // Create an output tensor
    Tensor* output_tensor;
    OP_REQUIRES_OK(context, context->allocate_output(
                                0, TensorShape({input_shape.dim_size(0), 512,2}),
                                &output_tensor));
    auto outputEigen = output_tensor->shaped<int32, 3>({input_shape.dim_size(0), 512, 2});
    std::vector<std::vector<int>> indicesVec;
    for (int i = 0; i < input_shape.dim_size(0); i++)
    {
      std::vector<int> tmpIndices;
      for (int j = 0; j < input_shape.dim_size(1); j++)//output_shape.dim_size(1)
      {
        if(inputEigen(i,j) > 0.1)
        {
          tmpIndices.push_back(j);
        }
      }
      int lengthIndices = tmpIndices.size();
      if(lengthIndices > 0)
      {
        while(tmpIndices.size() < 512)
        {
          int rand_num = rand() % lengthIndices;
          tmpIndices.push_back(tmpIndices[rand_num]);
        }
        while (tmpIndices.size() > 512)
        {
          int rand_num = rand() % tmpIndices.size();
          tmpIndices.erase(tmpIndices.begin() + rand_num);
        }
      }
      indicesVec.push_back(tmpIndices);
    }
    for(int i = 0; i < indicesVec.size(); i++)
    {
      if(0 == indicesVec[i].size())
      {
        for(int j = 0; j < 512; j++)
        {
          outputEigen(i, j, 0) = i;  // batchsize
          outputEigen(i, j, 1) = 0;
        }
      }
      else
      {
        for(int j = 0; j < indicesVec[i].size(); j++)
        {
          outputEigen(i, j, 0) = i;  // batchsize
          outputEigen(i, j, 1) = indicesVec[i][j];
        }
      }
    }
  }
};
REGISTER_KERNEL_BUILDER(Name("Zhj").Device(DEVICE_CPU), PyFuncZhj);
}

Zhj = tf.load_op_library("/home/zhj/Py/frustum-pointnets-master/train/log_v1/zhj.so")
def tf_gather_object_pc(point_cloud, mask, npoints=512):
    ''' Gather object point clouds according to predicted masks.
    Input:
        point_cloud: TF tensor in shape (B,N,C)
        mask: TF tensor in shape (B,N) of 0 (not pick) or 1 (pick)
        npoints: int scalar, maximum number of points to keep (default: 512)
    Output:
        object_pc: TF tensor in shape (B,npoint,C)
        indices: TF int tensor in shape (B,npoint,2)
    '''
    def mask_to_indices(mask):
        indices = np.zeros((mask.shape[0], npoints, 2), dtype=np.int32)
        for i in range(mask.shape[0]):
            pos_indices = np.where(mask[i,:]>0.5)[0]
            # skip cases when pos_indices is empty
            if len(pos_indices) > 0: 
                if len(pos_indices) > npoints:
                    choice = np.random.choice(len(pos_indices),
                        npoints, replace=False)
                else:
                    choice = np.random.choice(len(pos_indices),
                        npoints-len(pos_indices), replace=True)
                    choice = np.concatenate((np.arange(len(pos_indices)), choice))
                np.random.shuffle(choice)
                indices[i,:,1] = pos_indices[choice]
            indices[i,:,0] = i
        return indices

    print("mask.shape----------------------------", mask.shape)  # mask.shape: (1,1024)
    # indices = tf.py_func(mask_to_indices, [mask], tf.int32)
    indices = Zhj.zhj([mask])
    object_pc = tf.gather_nd(point_cloud, indices)
    return object_pc, indices