/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

/** \file
 * \ingroup mathutils
 *
 * This file defines the 'mathutils.kdtree' module, a general purpose module to access
 * blenders kdtree for 3d spatial lookups.
 */

#include <Python.h>

#include "MEM_guardedalloc.h"

#include "BLI_kdtree.hh"
#include "BLI_utildefines.h"

#include "../generic/py_capi_utils.hh"
#include "../generic/python_utildefines.hh"

#include "mathutils.hh"
#include "mathutils_kdtree.hh" /* own include */

#include "BLI_strict_flags.h" /* IWYU pragma: keep. Keep last. */

namespace blender {

struct PyKDTree {
  PyObject_HEAD
  KDTree<float3> *obj;
  uint maxsize;
  uint count;
  uint count_balance; /* size when we last balanced */
};

/* -------------------------------------------------------------------- */
/* Utility helper functions */

static void kdtree_nearest_to_py_tuple(const KDTreeNearest<float3> *nearest, PyObject *py_retval)
{
  BLI_assert(nearest->index >= 0);
  BLI_assert(PyTuple_GET_SIZE(py_retval) == 3);

  PyTuple_SET_ITEMS(py_retval,
                    Vector_CreatePyObject(nearest->co, 3, nullptr),
                    PyLong_FromLong(nearest->index),
                    PyFloat_FromDouble(nearest->dist));
}

static PyObject *kdtree_nearest_to_py(const KDTreeNearest<float3> *nearest)
{
  PyObject *py_retval;

  py_retval = PyTuple_New(3);

  kdtree_nearest_to_py_tuple(nearest, py_retval);

  return py_retval;
}

static PyObject *kdtree_nearest_to_py_and_check(const KDTreeNearest<float3> *nearest)
{
  PyObject *py_retval;

  py_retval = PyTuple_New(3);

  if (nearest->index != -1) {
    kdtree_nearest_to_py_tuple(nearest, py_retval);
  }
  else {
    PyC_Tuple_Fill(py_retval, Py_None);
  }

  return py_retval;
}

/* -------------------------------------------------------------------- */
/* KDTree */

/* annoying since arg parsing won't check overflow */
#define UINT_IS_NEG(n) ((n) > INT_MAX)

static int PyKDTree__tp_init(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
  uint maxsize;
  const char *keywords[] = {"size", nullptr};

  if (!PyArg_ParseTupleAndKeywords(
          args, kwargs, "I:KDTree", const_cast<char **>(keywords), &maxsize))
  {
    return -1;
  }

  if (UINT_IS_NEG(maxsize)) {
    PyErr_SetString(PyExc_ValueError, "negative 'size' given");
    return -1;
  }

  self->obj = kdtree_new<float3>(maxsize);
  self->maxsize = maxsize;
  self->count = 0;
  /* Initialize `uint-max` to avoid crashes on unbalanced trees. */
  self->count_balance = uint(-1);

  return 0;
}

static void PyKDTree__tp_dealloc(PyKDTree *self)
{
  kdtree_free<float3>(self->obj);
  Py_TYPE(self)->tp_free(reinterpret_cast<PyObject *>(self));
}

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_insert_doc,
    ".. method:: insert(co, index)\n"
    "\n"
    "   Insert a point into the KDTree.\n"
    "\n"
    "   :param co: Point 3d position.\n"
    "   :type co: Sequence[float]\n"
    "   :param index: The index of the point (must be non-negative).\n"
    "   :type index: int\n");
static PyObject *py_kdtree_insert(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
  PyObject *py_co;
  float co[3];
  int index;
  const char *keywords[] = {"co", "index", nullptr};

  if (!PyArg_ParseTupleAndKeywords(
          args, kwargs, "Oi:insert", const_cast<char **>(keywords), &py_co, &index))
  {
    return nullptr;
  }

  if (mathutils_array_parse(co, 3, 3, py_co, "insert: invalid 'co' arg") == -1) {
    return nullptr;
  }

  if (index < 0) {
    PyErr_SetString(PyExc_ValueError, "negative index given");
    return nullptr;
  }

  if (self->count >= self->maxsize) {
    PyErr_SetString(PyExc_RuntimeError, "Trying to insert more items than KDTree has room for");
    return nullptr;
  }

  kdtree_insert<float3>(self->obj, index, co);
  self->count++;

  Py_RETURN_NONE;
}

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_balance_doc,
    ".. method:: balance()\n"
    "\n"
    "   Balance the tree.\n"
    "\n"
    "   .. note::\n"
    "\n"
    "      This builds the entire tree, avoid calling after each insertion.\n");
static PyObject *py_kdtree_balance(PyKDTree *self)
{
  kdtree_balance<float3>(self->obj);
  self->count_balance = self->count;
  Py_RETURN_NONE;
}

struct PyKDTree_NearestData {
  PyObject *py_filter;
  bool is_error;
};

static int py_find_nearest_cb(void *user_data, int index, const float3 &co, float dist_sq)
{
  UNUSED_VARS(co, dist_sq);

  PyKDTree_NearestData *data = static_cast<PyKDTree_NearestData *>(user_data);

  PyObject *py_args = PyTuple_New(1);
  PyTuple_SET_ITEM(py_args, 0, PyLong_FromLong(index));
  PyObject *result = PyObject_CallObject(data->py_filter, py_args);
  Py_DECREF(py_args);

  if (result) {
    bool use_node;
    const int ok = PyC_ParseBool(result, &use_node);
    Py_DECREF(result);
    if (ok) {
      return int(use_node);
    }
  }

  data->is_error = true;
  return -1;
}

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_find_doc,
    ".. method:: find(co, *, filter=None)\n"
    "\n"
    "   Find nearest point to ``co``.\n"
    "\n"
    "   :param co: 3D coordinate.\n"
    "   :type co: Sequence[float]\n"
    "   :param filter: function which takes an index and returns True for indices to "
    "include in the search.\n"
    "   :type filter: Callable[[int], bool] | None\n"
    "   :return: Returns (position, index, distance),\n"
    "      or (None, None, None) when no match is found.\n"
    "   :rtype: tuple[:class:`Vector`, int, float] | tuple[None, None, None]\n");
static PyObject *py_kdtree_find(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
  PyObject *py_co, *py_filter = Py_None;
  float co[3];
  KDTreeNearest<float3> nearest;
  const char *keywords[] = {"co", "filter", nullptr};

  if (!PyArg_ParseTupleAndKeywords(
          args, kwargs, "O|$O:find", const_cast<char **>(keywords), &py_co, &py_filter))
  {
    return nullptr;
  }

  if (mathutils_array_parse(co, 3, 3, py_co, "find: invalid 'co' arg") == -1) {
    return nullptr;
  }

  if (self->count != self->count_balance) {
    PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find()");
    return nullptr;
  }

  nearest.index = -1;

  if (py_filter == Py_None) {
    kdtree_find_nearest<float3>(self->obj, co, &nearest);
  }
  else {
    PyKDTree_NearestData data = {nullptr};

    data.py_filter = py_filter;
    data.is_error = false;

    kdtree_find_nearest_cb<float3>(
        self->obj, co, &nearest, [&](int index, const float3 &co_nearest, float dist_sq) {
          return py_find_nearest_cb(&data, index, co_nearest, dist_sq);
        });

    if (data.is_error) {
      return nullptr;
    }
  }

  return kdtree_nearest_to_py_and_check(&nearest);
}

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_find_n_doc,
    ".. method:: find_n(co, n)\n"
    "\n"
    "   Find nearest ``n`` points to ``co``.\n"
    "\n"
    "   :param co: 3D coordinate.\n"
    "   :type co: Sequence[float]\n"
    "   :param n: Number of points to find.\n"
    "   :type n: int\n"
    "   :return: Returns a list of tuples (position, index, distance).\n"
    "   :rtype: list[tuple[:class:`Vector`, int, float]]\n");
static PyObject *py_kdtree_find_n(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
  PyObject *py_list;
  PyObject *py_co;
  float co[3];
  KDTreeNearest<float3> *nearest;
  uint n;
  int i, found;
  const char *keywords[] = {"co", "n", nullptr};

  if (!PyArg_ParseTupleAndKeywords(
          args, kwargs, "OI:find_n", const_cast<char **>(keywords), &py_co, &n))
  {
    return nullptr;
  }

  if (mathutils_array_parse(co, 3, 3, py_co, "find_n: invalid 'co' arg") == -1) {
    return nullptr;
  }

  if (UINT_IS_NEG(n)) {
    PyErr_SetString(PyExc_RuntimeError, "negative 'n' given");
    return nullptr;
  }

  if (self->count != self->count_balance) {
    PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_n()");
    return nullptr;
  }

  nearest = MEM_new_array_uninitialized<KDTreeNearest<float3>>(n, __func__);

  found = kdtree_find_nearest_n<float3>(self->obj, co, nearest, n);

  py_list = PyList_New(found);

  for (i = 0; i < found; i++) {
    PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
  }

  MEM_delete(nearest);

  return py_list;
}

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_find_range_doc,
    ".. method:: find_range(co, radius)\n"
    "\n"
    "   Find all points within ``radius`` of ``co``.\n"
    "\n"
    "   :param co: 3D coordinate.\n"
    "   :type co: Sequence[float]\n"
    "   :param radius: Maximum distance to search for points.\n"
    "   :type radius: float\n"
    "   :return: Returns a list of tuples (position, index, distance).\n"
    "   :rtype: list[tuple[:class:`Vector`, int, float]]\n");
static PyObject *py_kdtree_find_range(PyKDTree *self, PyObject *args, PyObject *kwargs)
{
  PyObject *py_list;
  PyObject *py_co;
  float co[3];
  KDTreeNearest<float3> *nearest = nullptr;
  float radius;
  int i, found;

  const char *keywords[] = {"co", "radius", nullptr};

  if (!PyArg_ParseTupleAndKeywords(
          args, kwargs, "Of:find_range", const_cast<char **>(keywords), &py_co, &radius))
  {
    return nullptr;
  }

  if (mathutils_array_parse(co, 3, 3, py_co, "find_range: invalid 'co' arg") == -1) {
    return nullptr;
  }

  if (radius < 0.0f) {
    PyErr_SetString(PyExc_RuntimeError, "negative radius given");
    return nullptr;
  }

  if (self->count != self->count_balance) {
    PyErr_SetString(PyExc_RuntimeError, "KDTree must be balanced before calling find_range()");
    return nullptr;
  }

  found = kdtree_range_search<float3>(self->obj, co, &nearest, radius);

  py_list = PyList_New(found);

  for (i = 0; i < found; i++) {
    PyList_SET_ITEM(py_list, i, kdtree_nearest_to_py(&nearest[i]));
  }

  if (nearest) {
    MEM_delete(nearest);
  }

  return py_list;
}

#ifdef __GNUC__
#  ifdef __clang__
#    pragma clang diagnostic push
#    pragma clang diagnostic ignored "-Wcast-function-type"
#  else
#    pragma GCC diagnostic push
#    pragma GCC diagnostic ignored "-Wcast-function-type"
#  endif
#endif

static PyMethodDef PyKDTree_methods[] = {
    {"insert",
     reinterpret_cast<PyCFunction>(py_kdtree_insert),
     METH_VARARGS | METH_KEYWORDS,
     py_kdtree_insert_doc},
    {"balance",
     reinterpret_cast<PyCFunction>(py_kdtree_balance),
     METH_NOARGS,
     py_kdtree_balance_doc},
    {"find",
     reinterpret_cast<PyCFunction>(py_kdtree_find),
     METH_VARARGS | METH_KEYWORDS,
     py_kdtree_find_doc},
    {"find_n",
     reinterpret_cast<PyCFunction>(py_kdtree_find_n),
     METH_VARARGS | METH_KEYWORDS,
     py_kdtree_find_n_doc},
    {"find_range",
     reinterpret_cast<PyCFunction>(py_kdtree_find_range),
     METH_VARARGS | METH_KEYWORDS,
     py_kdtree_find_range_doc},
    {nullptr, nullptr, 0, nullptr},
};

#ifdef __GNUC__
#  ifdef __clang__
#    pragma clang diagnostic pop
#  else
#    pragma GCC diagnostic pop
#  endif
#endif

PyDoc_STRVAR(
    /* Wrap. */
    py_KDtree_doc,
    ".. class:: KDTree(size)\n"
    "\n"
    "   KDTree(size) -> new kd-tree initialized to hold up to ``size`` items.\n"
    "\n"
    "   :param size: Maximum number of items.\n"
    "   :type size: int\n"
    "\n"
    "   .. note::\n"
    "\n"
    "      :meth:`KDTree.balance` must have been called before using any of the ``find`` "
    "methods.\n");
PyTypeObject PyKDTree_Type = {
    /*ob_base*/ PyVarObject_HEAD_INIT(nullptr, 0)
    /*tp_name*/ "KDTree",
    /*tp_basicsize*/ sizeof(PyKDTree),
    /*tp_itemsize*/ 0,
    /*tp_dealloc*/ reinterpret_cast<destructor>(PyKDTree__tp_dealloc),
    /*tp_vectorcall_offset*/ 0,
    /*tp_getattr*/ nullptr,
    /*tp_setattr*/ nullptr,
    /*tp_as_async*/ nullptr,
    /*tp_repr*/ nullptr,
    /*tp_as_number*/ nullptr,
    /*tp_as_sequence*/ nullptr,
    /*tp_as_mapping*/ nullptr,
    /*tp_hash*/ nullptr,
    /*tp_call*/ nullptr,
    /*tp_str*/ nullptr,
    /*tp_getattro*/ nullptr,
    /*tp_setattro*/ nullptr,
    /*tp_as_buffer*/ nullptr,
    /*tp_flags*/ Py_TPFLAGS_DEFAULT,
    /*tp_doc*/ py_KDtree_doc,
    /*tp_traverse*/ nullptr,
    /*tp_clear*/ nullptr,
    /*tp_richcompare*/ nullptr,
    /*tp_weaklistoffset*/ 0,
    /*tp_iter*/ nullptr,
    /*tp_iternext*/ nullptr,
    /*tp_methods*/ static_cast<PyMethodDef *>(PyKDTree_methods),
    /*tp_members*/ nullptr,
    /*tp_getset*/ nullptr,
    /*tp_base*/ nullptr,
    /*tp_dict*/ nullptr,
    /*tp_descr_get*/ nullptr,
    /*tp_descr_set*/ nullptr,
    /*tp_dictoffset*/ 0,
    /*tp_init*/ reinterpret_cast<initproc>(PyKDTree__tp_init),
    /*tp_alloc*/ static_cast<allocfunc>(PyType_GenericAlloc),
    /*tp_new*/ static_cast<newfunc>(PyType_GenericNew),
    /*tp_free*/ static_cast<freefunc>(nullptr),
    /*tp_is_gc*/ nullptr,
    /*tp_bases*/ nullptr,
    /*tp_mro*/ nullptr,
    /*tp_cache*/ nullptr,
    /*tp_subclasses*/ nullptr,
    /*tp_weaklist*/ nullptr,
    /*tp_del*/ static_cast<destructor>(nullptr),
    /*tp_version_tag*/ 0,
    /*tp_finalize*/ nullptr,
    /*tp_vectorcall*/ nullptr,
};

PyDoc_STRVAR(
    /* Wrap. */
    py_kdtree_doc,
    "Generic 3-dimensional kd-tree to perform spatial searches.");
static PyModuleDef kdtree_moduledef = {
    /*m_base*/ PyModuleDef_HEAD_INIT,
    /*m_name*/ "mathutils.kdtree",
    /*m_doc*/ py_kdtree_doc,
    /*m_size*/ 0,
    /*m_methods*/ nullptr,
    /*m_slots*/ nullptr,
    /*m_traverse*/ nullptr,
    /*m_clear*/ nullptr,
    /*m_free*/ nullptr,
};

PyMODINIT_FUNC PyInit_mathutils_kdtree()
{
  PyObject *m = PyModule_Create(&kdtree_moduledef);

  if (m == nullptr) {
    return nullptr;
  }

  /* Register the 'KDTree' class */
  if (PyType_Ready(&PyKDTree_Type)) {
    return nullptr;
  }
  PyModule_AddType(m, &PyKDTree_Type);

  return m;
}

}  // namespace blender
