/*
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#ifdef WITH_OPENVDB
#  include <openvdb/openvdb.h>
#  include <openvdb/tools/LevelSetUtil.h>
#  include <openvdb/tools/ParticlesToLevelSet.h>
#endif

#include "node_geometry_util.hh"

#include "DNA_mesh_types.h"

#include "BKE_geometry_set.hh"
#include "BKE_lib_id.h"
#include "BKE_material.h"
#include "BKE_mesh.h"
#include "BKE_volume.h"
#include "BKE_volume_to_mesh.hh"

namespace blender::nodes::node_geo_field_to_volume_cc {
static void node_declare(NodeDeclarationBuilder &b)
{
  b.add_input<decl::Float>("Density").implicit_field();

  b.add_input<decl::Vector>("Bounds Min").default_value(float3(-1.0f));
  b.add_input<decl::Vector>("Bounds Max").default_value(float3(1.0f));

  b.add_input<decl::Int>("Resolution X").default_value(32);
  b.add_input<decl::Int>("Resolution Y").default_value(32);
  b.add_input<decl::Int>("Resolution Z").default_value(32);

  b.add_output<decl::Geometry>("Volume");
}

// Is there a library map-range function?
static float map(float x, float in_min, float in_max, float out_min, float out_max)
{
  return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min;
}

class Grid3DFieldContext : public FieldContext {
 private:
  int3 resolution_;
  float3 bounds_min_;
  float3 bounds_max_;

 public:
  Grid3DFieldContext(const int3 resolution, const float3 bounds_min, const float3 bounds_max)
      : resolution_(resolution), bounds_min_(bounds_min), bounds_max_(bounds_max)
  {
  }

  GVArray get_varray_for_input(const FieldInput &field_input,
                               const IndexMask mask,
                               ResourceScope &UNUSED(scope)) const
  {
    const bke::AttributeFieldInput *attribute_field_input =
        dynamic_cast<const bke::AttributeFieldInput *>(&field_input);
    if (attribute_field_input == nullptr) {
      return {};
    }
    if (attribute_field_input->attribute_name() != "position") {
      return {};
    }

    const int64_t voxel_num = static_cast<int64_t>(resolution_.x) *
                              static_cast<int64_t>(resolution_.y) *
                              static_cast<int64_t>(resolution_.z);
    Array<float3> positions(voxel_num);
    int64_t counter = 0;
    for (const int64_t x_i : IndexRange(resolution_.x)) {
      const float x = map(x_i, 0.0f, resolution_.x - 1, bounds_min_.x, bounds_max_.x);
      for (const int64_t y_i : IndexRange(resolution_.y)) {
        const float y = map(y_i, 0.0f, resolution_.y - 1, bounds_min_.y, bounds_max_.y);
        for (const int64_t z_i : IndexRange(resolution_.z)) {
          const float z = map(z_i, 0.0f, resolution_.z - 1, bounds_min_.z, bounds_max_.z);
          positions[counter] = float3(x, y, z);
          counter++;
        }
      }
    }
    return VArray<float3>::ForContainer(std::move(positions));
  }
};

static void node_geo_exec(GeoNodeExecParams params)
{
  const float3 bounds_min = params.extract_input<float3>("Bounds Min");
  const float3 bounds_max = params.extract_input<float3>("Bounds Max");
  const int resX = params.extract_input<int>("Resolution X");
  const int resY = params.extract_input<int>("Resolution Y");
  const int resZ = params.extract_input<int>("Resolution Z");

  if (resX < 1 || resY < 1 || resZ < 1) {
    params.error_message_add(NodeWarningType::Info, TIP_("Resolution must be greater than 0"));
    params.set_default_remaining_outputs();
    return;
  }

  Field<float> input_field = params.extract_input<Field<float>>("Density");

  // Evaluate Field on the point grid
  Grid3DFieldContext context({resX, resY, resZ}, bounds_min, bounds_max);

  const int64_t domain_size = static_cast<int64_t>(resX) * static_cast<int64_t>(resY) *
                              static_cast<int64_t>(resZ);
  FieldEvaluator evaluator(context, domain_size);
  int eval_idx = evaluator.add(input_field);
  {
    evaluator.evaluate();
  }

  // Store resulting distances in openvdb grid
  openvdb::FloatGrid::Ptr grid = openvdb::FloatGrid::create(FLT_MAX);
  {
    VArray<float> results = evaluator.get_evaluated(eval_idx).typed<float>();
    openvdb::FloatGrid::Accessor accessor = grid->getAccessor();
    openvdb::Coord ijk;
    int &i = ijk[0], &j = ijk[1], &k = ijk[2];
    int idx = 0;
    for (i = 0; i < resX; i++) {
      for (j = 0; j < resY; j++) {
        for (k = 0; k < resZ; k++) {
          accessor.setValue(ijk, results[idx]);
          idx++;
        }
      }
    }
  }

  float3 scale_fac = (bounds_max - bounds_min) / float3(resX, resY, resZ);
  grid->transform().postScale(openvdb::math::Vec3<float>(scale_fac.x, scale_fac.y, scale_fac.z));
  grid->transform().postTranslate(
      openvdb::math::Vec3<float>(bounds_min.x, bounds_min.y, bounds_min.z));

  Volume *volume = (Volume *)BKE_id_new_nomain(ID_VO, nullptr);
  BKE_volume_init_grids(volume);

  BKE_volume_grid_add_vdb(*volume, "density", std::move(grid));

  GeometrySet r_geometry_set;
  r_geometry_set.replace_volume(volume);
  params.set_output("Volume", r_geometry_set);
}
}  // namespace blender::nodes::node_geo_field_to_volume_cc

void register_node_type_field_to_volume()
{
  namespace file_ns = blender::nodes::node_geo_field_to_volume_cc;

  static bNodeType ntype;

  geo_node_type_base(&ntype, GEO_NODE_FIELD_TO_VOLUME, "Field to Volume", NODE_CLASS_GEOMETRY);

  ntype.declare = file_ns::node_declare;
  ntype.geometry_node_execute = file_ns::node_geo_exec;
  nodeRegisterType(&ntype);
}