/* SPDX-License-Identifier: GPL-2.0-or-later */

#include "BLI_array.hh"
#include "BLI_span.hh"
#include "BLI_virtual_array.hh"

#include "BKE_attribute_math.hh"
#include "BKE_spline.hh"

using blender::Array;
using blender::float3;
using blender::IndexRange;
using blender::MutableSpan;
using blender::Span;
using blender::VArray;
using blender::fn::GVArray;

void NURBSpline::copy_settings(Spline &dst) const
{
  NURBSpline &nurbs = static_cast<NURBSpline &>(dst);
  nurbs.knots_mode = knots_mode;
  nurbs.resolution_ = resolution_;
  nurbs.order_ = order_;
}

void NURBSpline::copy_data(Spline &dst) const
{
  NURBSpline &nurbs = static_cast<NURBSpline &>(dst);
  nurbs.positions_ = positions_;
  nurbs.weights_ = weights_;
  nurbs.knots_ = knots_;
  nurbs.knots_dirty_ = knots_dirty_;
  nurbs.radii_ = radii_;
  nurbs.tilts_ = tilts_;
}

int NURBSpline::size() const
{
  const int size = positions_.size();
  BLI_assert(size == radii_.size());
  BLI_assert(size == tilts_.size());
  BLI_assert(size == weights_.size());
  return size;
}

int NURBSpline::resolution() const
{
  return resolution_;
}

void NURBSpline::set_resolution(const int value)
{
  BLI_assert(value > 0);
  resolution_ = value;
  this->mark_cache_invalid();
}

uint8_t NURBSpline::order() const
{
  return order_;
}

void NURBSpline::set_order(const uint8_t value)
{
  BLI_assert(value >= 2 && value <= 6);
  order_ = value;
  this->mark_cache_invalid();
}

void NURBSpline::resize(const int size)
{
  positions_.resize(size);
  radii_.resize(size);
  tilts_.resize(size);
  weights_.resize(size);
  this->mark_cache_invalid();
  attributes.reallocate(size);
}

MutableSpan<float3> NURBSpline::positions()
{
  return positions_;
}
Span<float3> NURBSpline::positions() const
{
  return positions_;
}
MutableSpan<float> NURBSpline::radii()
{
  return radii_;
}
Span<float> NURBSpline::radii() const
{
  return radii_;
}
MutableSpan<float> NURBSpline::tilts()
{
  return tilts_;
}
Span<float> NURBSpline::tilts() const
{
  return tilts_;
}
MutableSpan<float> NURBSpline::weights()
{
  return weights_;
}
Span<float> NURBSpline::weights() const
{
  return weights_;
}

void NURBSpline::reverse_impl()
{
  this->weights().reverse();
}

void NURBSpline::mark_cache_invalid()
{
  basis_cache_dirty_ = true;
  position_cache_dirty_ = true;
  tangent_cache_dirty_ = true;
  normal_cache_dirty_ = true;
  length_cache_dirty_ = true;
}

int NURBSpline::evaluated_points_size() const
{
  if (!this->check_valid_size_and_order()) {
    return 0;
  }
  return resolution_ * this->segments_size();
}

void NURBSpline::correct_end_tangents() const
{
}

bool NURBSpline::check_valid_size_and_order() const
{
  if (this->size() < order_) {
    return false;
  }

  if (ELEM(this->knots_mode, NURBS_KNOT_MODE_BEZIER, NURBS_KNOT_MODE_ENDPOINT_BEZIER)) {
    if (this->knots_mode == NURBS_KNOT_MODE_BEZIER && this->size() <= order_) {
      return false;
    }
    return (!is_cyclic_ || this->size() % (order_ - 1) == 0);
  }

  return true;
}

int NURBSpline::knots_size() const
{
  const int size = this->size() + order_;
  return is_cyclic_ ? size + order_ - 1 : size;
}

void NURBSpline::calculate_knots() const
{
  const KnotsMode mode = this->knots_mode;
  const int order = order_;
  const bool is_bezier = ELEM(mode, NURBS_KNOT_MODE_BEZIER, NURBS_KNOT_MODE_ENDPOINT_BEZIER);
  const bool is_end_point = ELEM(mode, NURBS_KNOT_MODE_ENDPOINT, NURBS_KNOT_MODE_ENDPOINT_BEZIER);
  /* Inner knots are always repeated once except on Bezier case. */
  const int repeat_inner = is_bezier ? order - 1 : 1;
  /* How many times to repeat 0.0 at the beginning of knot. */
  const int head = is_end_point ? (order - (is_cyclic_ ? 1 : 0)) :
                                  (is_bezier ? min_ii(2, repeat_inner) : 1);
  /* Number of knots replicating widths of the starting knots.
   * Covers both Cyclic and EndPoint cases. */
  const int tail = is_cyclic_ ? 2 * order - 1 : (is_end_point ? order : 0);

  knots_.resize(this->knots_size());
  MutableSpan<float> knots = knots_;

  int r = head;
  float current = 0.0f;

  const int offset = is_end_point && is_cyclic_ ? 1 : 0;
  if (offset) {
    knots[0] = current;
    current += 1.0f;
  }

  for (const int i : IndexRange(offset, knots.size() - offset - tail)) {
    knots[i] = current;
    r--;
    if (r == 0) {
      current += 1.0;
      r = repeat_inner;
    }
  }

  const int tail_index = knots.size() - tail;
  for (const int i : IndexRange(tail)) {
    knots[tail_index + i] = current + (knots[i] - knots[0]);
  }
}

Span<float> NURBSpline::knots() const
{
  if (!knots_dirty_) {
    BLI_assert(knots_.size() == this->knots_size());
    return knots_;
  }

  std::lock_guard lock{knots_mutex_};
  if (!knots_dirty_) {
    BLI_assert(knots_.size() == this->knots_size());
    return knots_;
  }

  this->calculate_knots();

  knots_dirty_ = false;

  return knots_;
}

static void calculate_basis_for_point(const float parameter,
                                      const int size,
                                      const int degree,
                                      Span<float> knots,
                                      MutableSpan<float> r_weights,
                                      int &r_start_index)
{
  const int order = degree + 1;

  int start = 0;
  int end = 0;
  for (const int i : IndexRange(size + degree)) {
    const bool knots_equal = knots[i] == knots[i + 1];
    if (knots_equal || parameter < knots[i] || parameter > knots[i + 1]) {
      continue;
    }

    start = std::max(i - degree, 0);
    end = i;
    break;
  }

  Array<float, 12> buffer(order * 2, 0.0f);

  buffer[end - start] = 1.0f;

  for (const int i_order : IndexRange(2, degree)) {
    if (end + i_order >= knots.size()) {
      end = size + degree - i_order;
    }
    for (const int i : IndexRange(end - start + 1)) {
      const int knot_index = start + i;

      float new_basis = 0.0f;
      if (buffer[i] != 0.0f) {
        new_basis += ((parameter - knots[knot_index]) * buffer[i]) /
                     (knots[knot_index + i_order - 1] - knots[knot_index]);
      }

      if (buffer[i + 1] != 0.0f) {
        new_basis += ((knots[knot_index + i_order] - parameter) * buffer[i + 1]) /
                     (knots[knot_index + i_order] - knots[knot_index + 1]);
      }

      buffer[i] = new_basis;
    }
  }

  buffer.as_mutable_span().drop_front(end - start + 1).fill(0.0f);
  r_weights.copy_from(buffer.as_span().take_front(order));
  r_start_index = start;
}

Span<NURBSpline::BasisCache> NURBSpline::calculate_basis_cache() const
{
  if (!basis_cache_dirty_) {
    return basis_cache_;
  }

  std::lock_guard lock{basis_cache_mutex_};
  if (!basis_cache_dirty_) {
    return basis_cache_;
  }

  const int size = this->size();
  const int eval_size = this->evaluated_points_size();
  if (eval_size == 0) {
    return {};
  }

  basis_cache_.resize(eval_size);

  const int order = this->order();
  const int degree = order - 1;
  Span<float> control_weights = this->weights();
  Span<float> knots = this->knots();

  MutableSpan<BasisCache> basis_cache(basis_cache_);

  const float start = knots[degree];
  const float end = is_cyclic_ ? knots[size + degree] : knots[size];
  const float step = (end - start) / this->evaluated_edges_size();
  for (const int i : IndexRange(eval_size)) {
    /* Clamp parameter due to floating point inaccuracy. */
    const float parameter = std::clamp(start + step * i, knots[0], knots[size + degree]);

    BasisCache &basis = basis_cache[i];
    basis.weights.resize(order);
    calculate_basis_for_point(parameter,
                              size + (is_cyclic_ ? degree : 0),
                              degree,
                              knots,
                              basis.weights,
                              basis.start_index);

    for (const int j : basis.weights.index_range()) {
      const int point_index = (basis.start_index + j) % size;
      basis.weights[j] *= control_weights[point_index];
    }
  }

  basis_cache_dirty_ = false;
  return basis_cache_;
}

template<typename T>
void interpolate_to_evaluated_impl(Span<NURBSpline::BasisCache> weights,
                                   const blender::VArray<T> &src,
                                   MutableSpan<T> dst)
{
  const int size = src.size();
  BLI_assert(dst.size() == weights.size());
  blender::attribute_math::DefaultMixer<T> mixer(dst);

  for (const int i : dst.index_range()) {
    Span<float> point_weights = weights[i].weights;
    const int start_index = weights[i].start_index;
    for (const int j : point_weights.index_range()) {
      const int point_index = (start_index + j) % size;
      mixer.mix_in(i, src[point_index], point_weights[j]);
    }
  }

  mixer.finalize();
}

GVArray NURBSpline::interpolate_to_evaluated(const GVArray &src) const
{
  BLI_assert(src.size() == this->size());

  if (src.is_single()) {
    return src;
  }

  Span<BasisCache> basis_cache = this->calculate_basis_cache();

  GVArray new_varray;
  blender::attribute_math::convert_to_static_type(src.type(), [&](auto dummy) {
    using T = decltype(dummy);
    if constexpr (!std::is_void_v<blender::attribute_math::DefaultMixer<T>>) {
      Array<T> values(this->evaluated_points_size());
      interpolate_to_evaluated_impl<T>(basis_cache, src.typed<T>(), values);
      new_varray = VArray<T>::ForContainer(std::move(values));
    }
  });

  return new_varray;
}

Span<float3> NURBSpline::evaluated_positions() const
{
  if (!position_cache_dirty_) {
    return evaluated_position_cache_;
  }

  std::lock_guard lock{position_cache_mutex_};
  if (!position_cache_dirty_) {
    return evaluated_position_cache_;
  }

  const int eval_size = this->evaluated_points_size();
  evaluated_position_cache_.resize(eval_size);

  /* TODO: Avoid copying the evaluated data from the temporary array. */
  VArray<float3> evaluated = Spline::interpolate_to_evaluated(positions_.as_span());
  evaluated.materialize(evaluated_position_cache_);

  position_cache_dirty_ = false;
  return evaluated_position_cache_;
}