// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_extensions/kernels/kernels/traditionalml/include_traditionalml_kernels.h"

#include "onnx_core/runtime/runtime_context.h"
#include <cstring>
#include <string>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_kernels::kernel {

namespace {

template <typename T>
onnx_kernels::Shape ValidateAndComputeOutputShape(const Tensor &x,
                                                  const std::vector<T> &class_labels) {
  EXT_ENFORCE_INVALID(x.data_type == static_cast<int32_t>(DataType::FLOAT),
                      "kernel::ZipMap expects a float input tensor, got data_type=",
                      std::to_string(x.data_type), ".");
  EXT_ENFORCE_INVALID(!class_labels.empty(), "kernel::ZipMap expects non-empty class labels.");
  EXT_ENFORCE_INVALID(x.shape.size() == 1 || x.shape.size() == 2,
                      "kernel::ZipMap expects input rank 1 or 2.");
  for (int64_t d : x.shape) {
    EXT_ENFORCE_INVALID(d >= 0, "kernel::ZipMap input shape dimensions must be non-negative.");
  }

  const int64_t class_count = x.shape.back();
  EXT_ENFORCE_INVALID(class_count == static_cast<int64_t>(class_labels.size()),
                      "kernel::ZipMap class labels size must match input class dimension, got ",
                      std::to_string(class_labels.size()), " labels for class dimension ",
                      std::to_string(class_count), ".");

  if (x.shape.size() == 1) {
    return onnx_kernels::Shape{1, class_count};
  }
  return x.shape;
}

template <typename T>
Tensor ComputeZipMapOutput(const Tensor &x, const std::vector<T> &class_labels,
                           RawBufferAllocator *allocator) {
  const onnx_kernels::Shape output_shape = ValidateAndComputeOutputShape(x, class_labels);
  Tensor out = MakeOutputTensor(static_cast<int32_t>(DataType::FLOAT), output_shape, x.size_bytes(),
                                allocator);
  if (x.size_bytes() > 0) {
    std::memcpy(out.mutable_bytes(), x.bytes(), x.size_bytes());
  }
  return out;
}

template <typename T>
void ComputeZipMapOutput(const Tensor &x, const std::vector<T> &class_labels, Tensor &output) {
  const onnx_kernels::Shape expected_shape = ValidateAndComputeOutputShape(x, class_labels);
  EXT_ENFORCE_INVALID(output.data_type == static_cast<int32_t>(DataType::FLOAT),
                      "kernel::ZipMap preallocated output dtype must be float.");
  EXT_ENFORCE_INVALID(output.shape == expected_shape,
                      "kernel::ZipMap preallocated output shape is incorrect.");
  EXT_ENFORCE_INVALID(output.size_bytes() == x.size_bytes(),
                      "kernel::ZipMap preallocated output buffer is incorrectly sized.");
  if (x.size_bytes() > 0) {
    std::memcpy(output.mutable_bytes(), x.bytes(), x.size_bytes());
  }
}

} // namespace

Tensor ZipMap::operator()(const Tensor &x, const std::vector<int64_t> &class_labels,
                          RuntimeContext * /*rt*/) const {
  return ComputeZipMapOutput(x, class_labels, ctx_.allocator);
}

Tensor ZipMap::operator()(const Tensor &x, const ParamStrings &class_labels,
                          RuntimeContext * /*rt*/) const {
  return ComputeZipMapOutput(x, class_labels, ctx_.allocator);
}

void ZipMap::operator()(const Tensor &x, const std::vector<int64_t> &class_labels,
                        Tensor &output) const {
  ComputeZipMapOutput(x, class_labels, output);
}

void ZipMap::operator()(const Tensor &x, const ParamStrings &class_labels, Tensor &output) const {
  ComputeZipMapOutput(x, class_labels, output);
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_kernels::kernel
