/*
  Copyright (C) 2024  Sutou Kouhei <kou@clear-code.com>

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  This library 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
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "grn_distance.hpp"

#include <groonga/bulk.hpp>

#include "grn_ctx.hpp"
#include "grn_db.h"

namespace grn {
  namespace distance {
#ifdef GRN_WITH_SIMSIMD
    bool use_simsimd = true;
    namespace simsimd {
      simsimd_capability_t capabilities = simsimd_cap_serial_k;
    }
#else
    bool use_simsimd = false;
#endif

#ifdef GRN_WITH_XSIMD
    bool use_xsimd = true;
#else
    bool use_xsimd = false;
#endif
  } // namespace distance
} // namespace grn

namespace {
  bool
  validate_vectors(grn_ctx *ctx,
                   grn_obj *vector1,
                   grn_obj *vector2,
                   const char *tag)
  {
    if (!grn_obj_is_number_family_vector(ctx, vector1)) {
      grn::TextBulk inspected(ctx);
      grn_inspect(ctx, *inspected, vector1);
      ERR(GRN_INVALID_ARGUMENT,
          "%s the 1st argument must be number family vector: %.*s",
          tag,
          static_cast<int>(GRN_TEXT_LEN(*inspected)),
          GRN_TEXT_VALUE(*inspected));
      return false;
    }

    if (!grn_obj_is_number_family_vector(ctx, vector2)) {
      grn::TextBulk inspected(ctx);
      grn_inspect(ctx, *inspected, vector2);
      ERR(GRN_INVALID_ARGUMENT,
          "%s the 2nd argument must be number family vector: %.*s",
          tag,
          static_cast<int>(GRN_TEXT_LEN(*inspected)),
          GRN_TEXT_VALUE(*inspected));
      return false;
    }

    if (vector1->header.domain != vector2->header.domain) {
      ERR(GRN_INVALID_ARGUMENT,
          "%s both vectors must be the same type: %s: %s",
          tag,
          grn_type_id_to_string_builtin(ctx, vector1->header.domain),
          grn_type_id_to_string_builtin(ctx, vector2->header.domain));
      return false;
    }

    auto n_elements1 = grn_vector_size(ctx, vector1);
    auto n_elements2 = grn_vector_size(ctx, vector2);
    if (n_elements1 != n_elements2) {
      ERR(GRN_INVALID_ARGUMENT,
          "%s both vectors must be the same size: %d: %d",
          tag,
          n_elements1,
          n_elements2);
      return false;
    }

    return true;
  }

  template <typename ElementType>
  float
  compute_distance_cosine(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
  {
    return grn::distance::compute_distance_cosine<ElementType>(vector1,
                                                               vector2);
  }

  template <typename ElementType>
  float
  compute_distance_inner_product(grn_ctx *ctx,
                                 grn_obj *vector1,
                                 grn_obj *vector2)
  {
    return grn::distance::compute_distance_inner_product<ElementType>(vector1,
                                                                      vector2);
  }

  template <typename ElementType>
  float
  compute_distance_l1_norm(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
  {
    return grn::distance::compute_difference_l1_norm<ElementType>(vector1,
                                                                  vector2);
  }

  template <typename ElementType>
  float
  compute_distance_l2_norm_squared(grn_ctx *ctx,
                                   grn_obj *vector1,
                                   grn_obj *vector2)
  {
    return grn::distance::compute_difference_l2_norm_squared<ElementType>(
      vector1,
      vector2);
  }

  enum class Method {
    Cosine,
    InnerProduct,
    L1Norm,
    L2NormSquared,
  };

  template <typename ElementType, typename DistanceFunc>
  void
  distance_apply(grn_ctx *ctx,
                 grn_applier_data *data,
                 grn_obj *table,
                 grn_obj *output_column,
                 grn_obj *input_column,
                 grn_id input_column_range,
                 grn_obj *literal,
                 const char *function_name,
                 DistanceFunc distance_func)
  {
    auto n_records = grn_table_size(ctx, table);
    if (n_records == 0) {
      return;
    }
    auto n_elements = GRN_BULK_VSIZE(literal) / sizeof(ElementType);
    auto task_executor = grn_ctx_get_task_executor(ctx);
    auto n_workers = task_executor->get_n_workers();
    if (n_workers == 0) {
      n_workers = 1;
    }
    n_workers = std::min(n_workers, n_records);
    int limit = (n_records == n_workers) ? 1 : ((n_records / n_workers) + 1);
    for (size_t i = 0; i < n_workers; i++) {
      auto execute = [&, i]() {
        grn_obj input;
        GRN_VALUE_FIX_SIZE_INIT(&input,
                                GRN_OBJ_VECTOR | GRN_OBJ_DO_SHALLOW_COPY,
                                input_column_range);
        grn_obj output;
        GRN_FLOAT_INIT(&output, 0);
        int offset = limit * i;
        grn_ctx *task_ctx = ctx;
        grn_ctx *child_ctx = nullptr;
        if (task_executor->is_parallel()) {
          task_ctx = child_ctx = grn_ctx_pull_child(ctx);
        }
        grn::ChildCtxReleaser releaser(ctx, child_ctx);
        auto cursor = grn_table_cursor_open(task_ctx,
                                            table,
                                            nullptr,
                                            0,
                                            nullptr,
                                            0,
                                            offset,
                                            limit,
                                            GRN_CURSOR_BY_ID);
        if (cursor) {
          grn_id id;
          while ((id = grn_table_cursor_next(task_ctx, cursor)) != GRN_ID_NIL) {
            uint32_t size;
            auto value = grn_obj_get_value_(task_ctx, input_column, id, &size);
            GRN_TEXT_SET(task_ctx, &input, value, size);
            if ((GRN_BULK_VSIZE(&input) / sizeof(ElementType)) != n_elements) {
              continue;
            }
            auto distance_raw = distance_func(task_ctx, &input, literal);
            if (task_ctx->rc != GRN_SUCCESS) {
              continue;
            }
            GRN_FLOAT_SET(task_ctx, &output, distance_raw);
            grn_obj_set_value(task_ctx,
                              output_column,
                              id,
                              &output,
                              GRN_OBJ_SET);
            if (task_ctx->rc != GRN_SUCCESS) {
              continue;
            }
          }
          grn_table_cursor_close(task_ctx, cursor);
        }
        GRN_OBJ_FIN(task_ctx, &input);
        GRN_OBJ_FIN(task_ctx, &output);
        return true;
      };
      task_executor->execute(i, execute, function_name);
    }
    task_executor->wait_all();
  }

  grn_rc
  distance_applier(grn_ctx *ctx,
                   grn_applier_data *data,
                   Method method,
                   const char *function_name)
  {
    auto table = grn_applier_data_get_table(ctx, data);
    auto output_column = grn_applier_data_get_output_column(ctx, data);
    size_t n_args;
    auto args = grn_applier_data_get_args(ctx, data, &n_args);

    if (n_args != 2) {
      ERR(GRN_INVALID_ARGUMENT,
          "%s: wrong number of arguments (%" GRN_FMT_SIZE " for 2)",
          function_name,
          n_args);
      return ctx->rc;
    }

    auto input_column = args[0];
    if (!(grn_obj_is_vector_column(ctx, input_column) ||
          grn_obj_is_vector_accessor(ctx, input_column))) {
      grn::TextBulk inspected(ctx);
      grn_inspect(ctx, *inspected, input_column);
      ERR(GRN_INVALID_ARGUMENT,
          "%s: 1st argument must be a vector column or accessor: %.*s",
          function_name,
          static_cast<int>(GRN_TEXT_LEN(*inspected)),
          GRN_TEXT_VALUE(*inspected));
      return ctx->rc;
    }
    auto input_column_range = grn_obj_get_range(ctx, input_column);
    switch (input_column_range) {
    case GRN_DB_FLOAT32:
    case GRN_DB_FLOAT:
      break;
    default:
      ERR(GRN_INVALID_ARGUMENT,
          "%s: 1st argument must be Float or Float32: %s",
          function_name,
          grn_type_id_to_string_builtin(ctx, input_column_range));
      return ctx->rc;
    }

    auto literal = args[1];
    grn_obj casted_literal_buffer;
    grn_obj *casted_literal = NULL;
    if (grn_obj_is_uvector(ctx, literal) &&
        literal->header.domain == input_column_range) {
      /* We can use this as-is. */
    } else {
      GRN_VALUE_FIX_SIZE_INIT(&casted_literal_buffer,
                              GRN_OBJ_VECTOR,
                              input_column_range);
      if (grn_obj_cast(ctx, literal, &casted_literal_buffer, false) !=
          GRN_SUCCESS) {
        GRN_OBJ_FIN(ctx, &casted_literal_buffer);
        grn::TextBulk inspected(ctx);
        grn_inspect(ctx, *inspected, literal);
        ERR(GRN_INVALID_ARGUMENT,
            "%s: 2nd argument must be a Float or Float32 vector literal: %.*s",
            function_name,
            static_cast<int>(GRN_TEXT_LEN(*inspected)),
            GRN_TEXT_VALUE(*inspected));
        return ctx->rc;
      }
      casted_literal = &casted_literal_buffer;
      literal = casted_literal;
    }

    switch (input_column_range) {
    case GRN_DB_FLOAT32:
      {
        auto apply = [&](auto distance_func) {
          return distance_apply<float>(ctx,
                                       data,
                                       table,
                                       output_column,
                                       input_column,
                                       input_column_range,
                                       literal,
                                       function_name,
                                       distance_func);
        };
        switch (method) {
        case Method::Cosine:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_cosine<float>(ctx, vector1, vector2);
          });
          break;
        case Method::InnerProduct:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_inner_product<float>(ctx, vector1, vector2);
          });
          break;
        case Method::L1Norm:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_l1_norm<float>(ctx, vector1, vector2);
          });
          break;
        case Method::L2NormSquared:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_l2_norm_squared<float>(ctx,
                                                           vector1,
                                                           vector2);
          });
          break;
        }
        break;
      }
    case GRN_DB_FLOAT:
      {
        auto apply = [&](auto distance_func) {
          distance_apply<double>(ctx,
                                 data,
                                 table,
                                 output_column,
                                 input_column,
                                 input_column_range,
                                 literal,
                                 function_name,
                                 distance_func);
        };
        switch (method) {
        case Method::Cosine:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_cosine<double>(ctx, vector1, vector2);
          });
          break;
        case Method::InnerProduct:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_inner_product<double>(ctx,
                                                          vector1,
                                                          vector2);
          });
          break;
        case Method::L1Norm:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_l1_norm<double>(ctx, vector1, vector2);
          });
          break;
        case Method::L2NormSquared:
          apply([](grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2) {
            return compute_distance_l2_norm_squared<double>(ctx,
                                                            vector1,
                                                            vector2);
          });
          break;
        }
        break;
      }
    default:
      break;
    }

    if (literal == casted_literal) {
      GRN_OBJ_FIN(ctx, casted_literal);
    }

    return ctx->rc;
  }
} // namespace

extern "C" void
grn_distance_init_external_libraries(void)
{
#ifdef GRN_WITH_SIMSIMD
  grn::distance::simsimd::capabilities = simsimd_capabilities();
#endif
}

extern "C" void
grn_distance_init_from_env(void)
{
  {
    char grn_distance_simd_env[GRN_ENV_BUFFER_SIZE];
    grn_getenv("GRN_DISTANCE_SIMD", grn_distance_simd_env, GRN_ENV_BUFFER_SIZE);
    if (strcmp(grn_distance_simd_env, "simsimd") == 0) {
      grn::distance::use_simsimd = true;
      grn::distance::use_xsimd = false;
    } else if (strcmp(grn_distance_simd_env, "xsimd") == 0) {
      grn::distance::use_simsimd = false;
      grn::distance::use_xsimd = true;
    } else if (strcmp(grn_distance_simd_env, "no") == 0) {
      grn::distance::use_simsimd = false;
      grn::distance::use_xsimd = false;
    }
  }
}

extern "C" float
grn_distance_cosine(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
{
  const char *tag = "[distance][cosine]";
  float distance = 0.0;

  if (!validate_vectors(ctx, vector1, vector2, tag)) {
    return distance;
  }

  switch (vector1->header.domain) {
  case GRN_DB_FLOAT32:
    distance = compute_distance_cosine<float>(ctx, vector1, vector2);
    break;
  case GRN_DB_FLOAT:
    distance = compute_distance_cosine<double>(ctx, vector1, vector2);
    break;
  default:
    // TODO: We should add support for all integer types
    ERR(GRN_FUNCTION_NOT_IMPLEMENTED,
        "%s unsupported type: %s",
        tag,
        grn_type_id_to_string_builtin(ctx, vector1->header.domain));
    break;
  }

  return distance;
}

extern "C" grn_rc
grn_distance_cosine_applier(grn_ctx *ctx, grn_applier_data *data)
{
  return distance_applier(ctx, data, Method::Cosine, "distance_cosine()");
}

extern "C" float
grn_distance_inner_product(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
{
  const char *tag = "[distance][inner-product]";
  float distance = 0.0;

  if (!validate_vectors(ctx, vector1, vector2, tag)) {
    return distance;
  }

  switch (vector1->header.domain) {
  case GRN_DB_FLOAT32:
    distance = compute_distance_inner_product<float>(ctx, vector1, vector2);
    break;
  case GRN_DB_FLOAT:
    distance = compute_distance_inner_product<double>(ctx, vector1, vector2);
    break;
  default:
    // TODO: We should add support for all integer types
    ERR(GRN_FUNCTION_NOT_IMPLEMENTED,
        "%s unsupported type: %s",
        tag,
        grn_type_id_to_string_builtin(ctx, vector1->header.domain));
    break;
  }

  return distance;
}

extern "C" grn_rc
grn_distance_inner_product_applier(grn_ctx *ctx, grn_applier_data *data)
{
  return distance_applier(ctx,
                          data,
                          Method::InnerProduct,
                          "distance_inner_product()");
}

extern "C" float
grn_distance_l1_norm(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
{
  const char *tag = "[distance][l1-norm]";
  float distance = 0.0;

  if (!validate_vectors(ctx, vector1, vector2, tag)) {
    return distance;
  }

  switch (vector1->header.domain) {
  case GRN_DB_FLOAT32:
    distance = compute_distance_l1_norm<float>(ctx, vector1, vector2);
    break;
  case GRN_DB_FLOAT:
    distance = compute_distance_l1_norm<double>(ctx, vector1, vector2);
    break;
  default:
    // TODO: We should add support for all integer types
    ERR(GRN_FUNCTION_NOT_IMPLEMENTED,
        "%s unsupported type: %s",
        tag,
        grn_type_id_to_string_builtin(ctx, vector1->header.domain));
    break;
  }

  return distance;
}

extern "C" grn_rc
grn_distance_l1_norm_applier(grn_ctx *ctx, grn_applier_data *data)
{
  return distance_applier(ctx, data, Method::L1Norm, "distance_l1_norm()");
}

extern "C" float
grn_distance_l2_norm_squared(grn_ctx *ctx, grn_obj *vector1, grn_obj *vector2)
{
  const char *tag = "[distance][l2-norm-squared]";
  float distance = 0.0;

  if (!validate_vectors(ctx, vector1, vector2, tag)) {
    return distance;
  }

  switch (vector1->header.domain) {
  case GRN_DB_FLOAT32:
    distance = compute_distance_l2_norm_squared<float>(ctx, vector1, vector2);
    break;
  case GRN_DB_FLOAT:
    distance = compute_distance_l2_norm_squared<double>(ctx, vector1, vector2);
    break;
  default:
    // TODO: We should add support for all integer types
    ERR(GRN_FUNCTION_NOT_IMPLEMENTED,
        "%s unsupported type: %s",
        tag,
        grn_type_id_to_string_builtin(ctx, vector1->header.domain));
    break;
  }

  return distance;
}

extern "C" grn_rc
grn_distance_l2_norm_squared_applier(grn_ctx *ctx, grn_applier_data *data)
{
  return distance_applier(ctx,
                          data,
                          Method::L2NormSquared,
                          "distance_l2_norm_squared()");
}