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

#include <algorithm>
#include <numeric>

#include "NOD_rna_define.hh"

#include "UI_interface_layout.hh"
#include "UI_resources.hh"

#include "BLI_array_utils.hh"

#include "NOD_socket_search_link.hh"

#include "RNA_enum_types.hh"

#include "node_geometry_util.hh"

namespace blender::nodes::node_geo_attribute_statistic_cc {

static void node_declare(NodeDeclarationBuilder &b)
{
  const bNode *node = b.node_or_null();

  b.add_input<decl::Geometry>("Geometry"_ustr).description("Geometry to get the statistics from");
  b.add_input<decl::Bool>("Selection"_ustr)
      .default_value(true)
      .evaluated_geometry_field()
      .hide_value();

  if (node != nullptr) {
    const eCustomDataType data_type = eCustomDataType(node->custom1);
    b.add_input(data_type, "Attribute"_ustr).hide_value().evaluated_geometry_field();

    b.add_output(data_type, "Mean"_ustr);
    b.add_output(data_type, CTX_N_(BLT_I18NCONTEXT_ID_NODETREE, "Median"_ustr))
        .translation_context(BLT_I18NCONTEXT_ID_NODETREE);
    b.add_output(data_type, "Sum"_ustr);
    b.add_output(data_type, "Min"_ustr);
    b.add_output(data_type, "Max"_ustr);
    b.add_output(data_type, "Range"_ustr);
    b.add_output(data_type, "Standard Deviation"_ustr);
    b.add_output(data_type, "Variance"_ustr);
  }
}

static void node_layout(ui::Layout &layout, bContext * /*C*/, PointerRNA *ptr)
{
  layout.prop(ptr, "data_type", UI_ITEM_NONE, "", ICON_NONE);
  layout.prop(ptr, "domain", UI_ITEM_NONE, "", ICON_NONE);
}

static void node_init(bNodeTree * /*tree*/, bNode *node)
{
  node->custom1 = CD_PROP_FLOAT;
  node->custom2 = int16_t(AttrDomain::Point);
}

static std::optional<eCustomDataType> node_type_from_other_socket(const bNodeSocket &socket)
{
  switch (socket.type) {
    case SOCK_FLOAT:
    case SOCK_BOOLEAN:
    case SOCK_INT:
      return CD_PROP_FLOAT;
    case SOCK_VECTOR:
    case SOCK_RGBA:
    case SOCK_ROTATION:
      return CD_PROP_FLOAT3;
    default:
      return {};
  }
}

static void node_gather_link_searches(GatherLinkSearchOpParams &params)
{
  const bke::bNodeType &node_type = params.node_type();
  const NodeDeclaration &declaration = *params.node_type().static_declaration;
  search_link_ops_for_declarations(params, declaration.inputs);

  const std::optional<eCustomDataType> type = node_type_from_other_socket(params.other_socket());
  if (!type) {
    return;
  }

  if (params.in_out() == SOCK_IN) {
    params.add_item(IFACE_("Attribute"), [node_type, type](LinkSearchOpParams &params) {
      bNode &node = params.add_node(node_type);
      node.custom1 = *type;
      params.update_and_connect_available_socket(node, "Attribute"_ustr);
    });
  }
  else {
    for (const UString name : {"Mean"_ustr,
                               "Median"_ustr,
                               "Sum"_ustr,
                               "Min"_ustr,
                               "Max"_ustr,
                               "Range"_ustr,
                               "Standard Deviation"_ustr,
                               "Variance"_ustr})
    {
      params.add_item(IFACE_(name.ref()), [node_type, name, type](LinkSearchOpParams &params) {
        bNode &node = params.add_node(node_type);
        node.custom1 = *type;
        params.update_and_connect_available_socket(node, name);
      });
    }
  }
}

static float compute_variance(const Span<float> data, const float mean)
{
  if (data.size() <= 1) {
    return 0.0f;
  }

  float sum_of_squared_differences = std::accumulate(
      data.begin(), data.end(), 0.0f, [mean](float accumulator, float value) {
        float difference = mean - value;
        return accumulator + difference * difference;
      });

  return sum_of_squared_differences / data.size();
}

static float median_of_sorted_span(const Span<float> data)
{
  if (data.is_empty()) {
    return 0.0f;
  }

  const float median = data[data.size() / 2];

  /* For spans of even length, the median is the average of the middle two elements. */
  if (data.size() % 2 == 0) {
    return (median + data[data.size() / 2 - 1]) * 0.5f;
  }
  return median;
}

static void node_geo_exec(GeoNodeExecParams params)
{
  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry"_ustr);
  const bNode &node = params.node();
  const eCustomDataType data_type = eCustomDataType(node.custom1);
  const AttrDomain domain = AttrDomain(node.custom2);
  Vector<const GeometryComponent *> components = geometry_set.get_components();

  const Field<bool> selection_field = params.extract_input<Field<bool>>("Selection"_ustr);

  switch (data_type) {
    case CD_PROP_FLOAT: {
      const Field<float> input_field = params.extract_input<Field<float>>("Attribute"_ustr);
      Vector<float> data;
      for (const GeometryComponent *component : components) {
        const int domain_size = component->attribute_domain_size(domain);
        if (domain_size == 0) {
          continue;
        }
        const bke::GeometryFieldContext field_context{*component, domain};
        fn::FieldEvaluator data_evaluator{field_context, domain_size};
        data_evaluator.add(input_field);
        data_evaluator.set_selection(selection_field);
        data_evaluator.evaluate();
        const VArray<float> component_data = data_evaluator.get_evaluated<float>(0);
        const IndexMask selection = data_evaluator.get_evaluated_selection_as_mask();

        const int next_data_index = data.size();
        data.resize(next_data_index + selection.size());
        MutableSpan<float> selected_data = data.as_mutable_span().slice(next_data_index,
                                                                        selection.size());
        array_utils::gather(component_data, selection, selected_data);
      }

      float mean = 0.0f;
      float median = 0.0f;
      float sum = 0.0f;
      float min = 0.0f;
      float max = 0.0f;
      float range = 0.0f;
      float standard_deviation = 0.0f;
      float variance = 0.0f;
      const bool sort_required = params.output_is_required("Min"_ustr) ||
                                 params.output_is_required("Max"_ustr) ||
                                 params.output_is_required("Range"_ustr) ||
                                 params.output_is_required("Median"_ustr);
      const bool sum_required = params.output_is_required("Sum"_ustr) ||
                                params.output_is_required("Mean"_ustr);
      const bool variance_required = params.output_is_required("Standard Deviation"_ustr) ||
                                     params.output_is_required("Variance"_ustr);

      if (data.size() != 0) {
        if (sort_required) {
          std::ranges::sort(data);
          median = median_of_sorted_span(data);

          min = data.first();
          max = data.last();
          range = max - min;
        }
        if (sum_required || variance_required) {
          sum = array_utils::compute_sum<float>(data);
          mean = sum / data.size();

          if (variance_required) {
            variance = compute_variance(data, mean);
            standard_deviation = std::sqrt(variance);
          }
        }
      }

      if (sum_required) {
        params.set_output("Sum"_ustr, sum);
        params.set_output("Mean"_ustr, mean);
      }
      if (sort_required) {
        params.set_output("Min"_ustr, min);
        params.set_output("Max"_ustr, max);
        params.set_output("Range"_ustr, range);
        params.set_output("Median"_ustr, median);
      }
      if (variance_required) {
        params.set_output("Standard Deviation"_ustr, standard_deviation);
        params.set_output("Variance"_ustr, variance);
      }
      break;
    }
    case CD_PROP_FLOAT3: {
      const Field<float3> input_field = params.extract_input<Field<float3>>("Attribute"_ustr);
      Vector<float3> data;
      for (const GeometryComponent *component : components) {
        const std::optional<AttributeAccessor> attributes = component->attributes();
        if (!attributes.has_value()) {
          continue;
        }
        if (attributes->domain_supported(domain)) {
          const bke::GeometryFieldContext field_context{*component, domain};
          fn::FieldEvaluator data_evaluator{field_context, attributes->domain_size(domain)};
          data_evaluator.add(input_field);
          data_evaluator.set_selection(selection_field);
          data_evaluator.evaluate();
          const VArray<float3> component_data = data_evaluator.get_evaluated<float3>(0);
          const IndexMask selection = data_evaluator.get_evaluated_selection_as_mask();

          const int next_data_index = data.size();
          data.resize(data.size() + selection.size());
          MutableSpan<float3> selected_data = data.as_mutable_span().slice(next_data_index,
                                                                           selection.size());
          array_utils::gather(component_data, selection, selected_data);
        }
      }

      float3 median{0};
      float3 min{0};
      float3 max{0};
      float3 range{0};
      float3 sum{0};
      float3 mean{0};
      float3 variance{0};
      float3 standard_deviation{0};
      const bool sort_required = params.output_is_required("Min"_ustr) ||
                                 params.output_is_required("Max"_ustr) ||
                                 params.output_is_required("Range"_ustr) ||
                                 params.output_is_required("Median"_ustr);
      const bool sum_required = params.output_is_required("Sum"_ustr) ||
                                params.output_is_required("Mean"_ustr);
      const bool variance_required = params.output_is_required("Standard Deviation"_ustr) ||
                                     params.output_is_required("Variance"_ustr);

      Array<float> data_x;
      Array<float> data_y;
      Array<float> data_z;
      if (sort_required || variance_required) {
        data_x.reinitialize(data.size());
        data_y.reinitialize(data.size());
        data_z.reinitialize(data.size());
        for (const int i : data.index_range()) {
          data_x[i] = data[i].x;
          data_y[i] = data[i].y;
          data_z[i] = data[i].z;
        }
      }

      if (data.size() != 0) {
        if (sort_required) {
          std::ranges::sort(data_x);
          std::ranges::sort(data_y);
          std::ranges::sort(data_z);

          const float x_median = median_of_sorted_span(data_x);
          const float y_median = median_of_sorted_span(data_y);
          const float z_median = median_of_sorted_span(data_z);
          median = float3(x_median, y_median, z_median);

          min = float3(data_x.first(), data_y.first(), data_z.first());
          max = float3(data_x.last(), data_y.last(), data_z.last());
          range = max - min;
        }
        if (sum_required || variance_required) {
          sum = array_utils::compute_sum(data.as_span());
          mean = sum / data.size();

          if (variance_required) {
            const float x_variance = compute_variance(data_x, mean.x);
            const float y_variance = compute_variance(data_y, mean.y);
            const float z_variance = compute_variance(data_z, mean.z);
            variance = float3(x_variance, y_variance, z_variance);
            standard_deviation = float3(
                std::sqrt(variance.x), std::sqrt(variance.y), std::sqrt(variance.z));
          }
        }
      }

      if (sum_required) {
        params.set_output("Sum"_ustr, sum);
        params.set_output("Mean"_ustr, mean);
      }
      if (sort_required) {
        params.set_output("Min"_ustr, min);
        params.set_output("Max"_ustr, max);
        params.set_output("Range"_ustr, range);
        params.set_output("Median"_ustr, median);
      }
      if (variance_required) {
        params.set_output("Standard Deviation"_ustr, standard_deviation);
        params.set_output("Variance"_ustr, variance);
      }
      break;
    }
    default:
      break;
  }
}

static void node_rna(StructRNA *srna)
{
  RNA_def_node_enum(
      srna,
      "data_type",
      "Data Type",
      "The data type the attribute is converted to before calculating the results",
      rna_enum_attribute_type_items,
      NOD_inline_enum_accessors(custom1),
      CD_PROP_FLOAT,
      [](bContext * /*C*/, PointerRNA * /*ptr*/, PropertyRNA * /*prop*/, bool *r_free) {
        *r_free = true;
        return enum_items_filter(rna_enum_attribute_type_items, [](const EnumPropertyItem &item) {
          return ELEM(item.value, CD_PROP_FLOAT, CD_PROP_FLOAT3);
        });
      });

  RNA_def_node_enum(srna,
                    "domain",
                    "Domain",
                    "Which domain to read the data from",
                    rna_enum_attribute_domain_items,
                    NOD_inline_enum_accessors(custom2),
                    int(AttrDomain::Point),
                    nullptr,
                    true);
}

static void node_register()
{
  static bke::bNodeType ntype;
  geo_node_type_base(&ntype, "GeometryNodeAttributeStatistic"_ustr, GEO_NODE_ATTRIBUTE_STATISTIC);
  ntype.ui_name = "Attribute Statistic";
  ntype.ui_description =
      "Calculate statistics about a data set from a field evaluated on a geometry";
  ntype.enum_name_legacy = "ATTRIBUTE_STATISTIC";
  ntype.nclass = NODE_CLASS_ATTRIBUTE;
  ntype.initfunc = node_init;
  ntype.declare = node_declare;
  ntype.geometry_node_execute = node_geo_exec;
  ntype.draw_buttons = node_layout;
  ntype.gather_link_search_ops = node_gather_link_searches;
  bke::node_register_type(ntype);

  node_rna(ntype.rna_ext.srna);
}
NOD_REGISTER_NODE(node_register)

}  // namespace blender::nodes::node_geo_attribute_statistic_cc
