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

#include "onnx_core/backend_test/expect.h"
#include "onnx_core/runtime/kernels/random.h"
#include "onnx_extensions/backend_test/cases/nn/include_nn_cases.h"
#include "onnx_extensions/kernels/kernels/nn/include_nn_kernels.h"
#include "onnx_proto/onnx_helper.h"

#include <cstdint>
#include <string>
#include <utility>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

// ---------------------------------------------------------------------------
// GRU — one-layer GRU with the default
// Sigmoid/Tanh activations, mirroring the upstream ``test_gru_*``
// reference cases (from onnx/backend/test/case/node/gru.py):
//
//   * ``test_cc_gru_defaults`` — tiny 1x3x2 input with X/W/R inputs only,
//     hidden_size=5; only Y_h is produced.
//   * ``test_cc_gru_with_initial_bias`` — exercises the optional ``B``
//     (concatenated ``[Wb, Rb]``) input with a custom Wb bias and zero
//     Rb, hidden_size=3.
//   * ``test_cc_gru_seq_length`` — multi-step (seq_length=2) GRU with the
//     optional ``B`` input, hidden_size=5; only Y_h is produced.
//   * ``test_cc_gru_batchwise`` — ``layout=1`` variant where ``X`` is laid
//     out as ``[batch_size, seq_length, input_size]``, hidden_size=6;
//     both Y and Y_h are produced.
// ---------------------------------------------------------------------------
void RegisterGRUCases(std::vector<TestCase> &registry, TestMode mode) {
  const OpsetId opset = DefaultOpset(22);
  const KernelContext ctx{opset};
  const onnx_kernels::kernel::GRU gru_kernel{ctx};

  constexpr int64_t kNumGates = 3;

  if (mode == TestMode::BENCHMARK) {
    const int64_t seq_length = 64;
    const int64_t batch_size = 32;
    const int64_t input_size = 128;
    const int64_t hidden_size = 128;
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_output("");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", hidden_size);
    Expect(registry, std::move(node), "test_cc_gru_benchmark", {opset},
           {seq_length * batch_size * input_size, kNumGates * hidden_size * input_size,
            kNumGates * hidden_size * hidden_size},
           {batch_size * hidden_size},
           [gru_kernel, seq_length, batch_size, input_size, hidden_size]() -> IoData {
             const std::vector<int64_t> x_shape = {seq_length, batch_size, input_size};
             const std::vector<int64_t> w_shape = {1, kNumGates * hidden_size, input_size};
             const std::vector<int64_t> r_shape = {1, kNumGates * hidden_size, hidden_size};
             Tensor x = RandnTensor(DataType::FLOAT, x_shape, 2001);
             Tensor w = RandnTensor(DataType::FLOAT, w_shape, 2002);
             Tensor r = RandnTensor(DataType::FLOAT, r_shape, 2003);
             auto [y_unused, y_h] = gru_kernel(x, w, r);
             (void)y_unused;
             return IoData{{std::move(x), std::move(w), std::move(r)}, {std::move(y_h)}};
           });
    return;
  }

  // ``gru_defaults``: seq_length=1, batch_size=3, input_size=2,
  // hidden_size=5 with weight_scale=0.1. Only Y_h is produced (Y is
  // skipped via an empty output name).
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_output("");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 5);
    Expect(registry, std::move(node), "test_cc_gru_defaults", {opset}, [=]() -> IoData {
      const int64_t seq_length = 1;
      const int64_t batch_size = 3;
      const int64_t input_size = 2;
      const int64_t hidden_size = 5;
      const float weight_scale = 0.1f;

      const std::vector<float> x_data{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
      std::vector<float> w_data(static_cast<size_t>(kNumGates * hidden_size * input_size),
                                weight_scale);
      std::vector<float> r_data(static_cast<size_t>(kNumGates * hidden_size * hidden_size),
                                weight_scale);
      Tensor x = Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_data);
      Tensor w = Tensor::FromFloat("", {1, kNumGates * hidden_size, input_size}, w_data);
      Tensor r = Tensor::FromFloat("", {1, kNumGates * hidden_size, hidden_size}, r_data);

      auto [y_unused, y_h] = gru_kernel(x, w, r);
      (void)y_unused; // Y is skipped (empty output name).

      return IoData{{std::move(x), std::move(w), std::move(r)}, {std::move(y_h)}};
    });
  }

  // ``gru_with_initial_bias``: seq_length=1, batch_size=3, input_size=3,
  // hidden_size=3 with a custom-valued ``B`` (Wb is all ``custom_bias``,
  // Rb is all zeros); only Y_h is produced.
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_input("B");
    node.add_output("");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 3);
    Expect(registry, std::move(node), "test_cc_gru_with_initial_bias", {opset}, [=]() -> IoData {
      const int64_t seq_length = 1;
      const int64_t batch_size = 3;
      const int64_t input_size = 3;
      const int64_t hidden_size = 3;
      const float weight_scale = 0.1f;
      const float custom_bias = 0.1f;

      const std::vector<float> x_data{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
      std::vector<float> w_data(static_cast<size_t>(kNumGates * hidden_size * input_size),
                                weight_scale);
      std::vector<float> r_data(static_cast<size_t>(kNumGates * hidden_size * hidden_size),
                                weight_scale);
      // B = concat(Wb, Rb) along the last axis: Wb = custom_bias * ones,
      // Rb = zeros. Shape: [1, 6 * hidden_size].
      std::vector<float> b_data(static_cast<size_t>(2 * kNumGates * hidden_size), 0.0f);
      for (int64_t i = 0; i < kNumGates * hidden_size; ++i) {
        b_data[static_cast<size_t>(i)] = custom_bias;
      }
      Tensor x = Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_data);
      Tensor w = Tensor::FromFloat("", {1, kNumGates * hidden_size, input_size}, w_data);
      Tensor r = Tensor::FromFloat("", {1, kNumGates * hidden_size, hidden_size}, r_data);
      Tensor b = Tensor::FromFloat("", {1, 2 * kNumGates * hidden_size}, b_data);

      auto [y_unused, y_h] = gru_kernel(x, w, r, b);
      (void)y_unused;

      return IoData{{std::move(x), std::move(w), std::move(r), std::move(b)}, {std::move(y_h)}};
    });
  }

  // ``gru_seq_length``: seq_length=2, batch_size=3, input_size=3,
  // hidden_size=5 with deterministic (non-random) ``W``/``R``/``B``
  // values, mirroring upstream ``test_gru_seq_length`` (which uses
  // ``np.random.randn`` — here we substitute a deterministic ramp so the
  // expected output is reproducible). Only Y_h is produced.
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_input("B");
    node.add_output("");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 5);
    Expect(registry, std::move(node), "test_cc_gru_seq_length", {opset}, [=]() -> IoData {
      const int64_t seq_length = 2;
      const int64_t batch_size = 3;
      const int64_t input_size = 3;
      const int64_t hidden_size = 5;

      std::vector<float> x_data(static_cast<size_t>(seq_length * batch_size * input_size));
      for (size_t i = 0; i < x_data.size(); ++i) {
        x_data[i] = static_cast<float>(i + 1);
      }
      std::vector<float> w_data(static_cast<size_t>(kNumGates * hidden_size * input_size));
      for (size_t i = 0; i < w_data.size(); ++i) {
        w_data[i] = static_cast<float>(i % 7) * 0.05f - 0.1f;
      }
      std::vector<float> r_data(static_cast<size_t>(kNumGates * hidden_size * hidden_size));
      for (size_t i = 0; i < r_data.size(); ++i) {
        r_data[i] = static_cast<float>(i % 9) * 0.04f - 0.15f;
      }
      std::vector<float> b_data(static_cast<size_t>(2 * kNumGates * hidden_size));
      for (size_t i = 0; i < b_data.size(); ++i) {
        b_data[i] = static_cast<float>(i) * 0.02f - 0.05f;
      }
      Tensor x = Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_data);
      Tensor w = Tensor::FromFloat("", {1, kNumGates * hidden_size, input_size}, w_data);
      Tensor r = Tensor::FromFloat("", {1, kNumGates * hidden_size, hidden_size}, r_data);
      Tensor b = Tensor::FromFloat("", {1, 2 * kNumGates * hidden_size}, b_data);

      auto [y_unused, y_h] = gru_kernel(x, w, r, b);
      (void)y_unused;

      return IoData{{std::move(x), std::move(w), std::move(r), std::move(b)}, {std::move(y_h)}};
    });
  }

  // ``gru_batchwise``: ``layout=1`` variant with batch_size=3,
  // seq_length=1, input_size=2, hidden_size=6. The kernel itself only
  // implements ``layout=0`` so we run it on the axis-swapped ``X`` and
  // re-permute the outputs to the batchwise layout:
  //   * ``Y``   : [seq, 1, batch, hidden]  -> [batch, seq, 1, hidden]
  //   * ``Y_h`` : [1, batch, hidden]       -> [batch, 1, hidden]
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_output("Y");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 6);
    AddAttribute<int64_t>(node, "layout", 1);
    Expect(registry, std::move(node), "test_cc_gru_batchwise", {opset}, [=]() -> IoData {
      const int64_t batch_size = 3;
      const int64_t seq_length = 1;
      const int64_t input_size = 2;
      const int64_t hidden_size = 6;
      const float weight_scale = 0.2f;

      // Input is [batch_size, seq_length, input_size] for layout=1.
      const std::vector<float> x_batchwise_data{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
      Tensor x_batchwise =
          Tensor::FromFloat("", {batch_size, seq_length, input_size}, x_batchwise_data);

      // Same contents but laid out for layout=0
      // ([seq_length, batch_size, input_size]); for seq_length=1 the axis
      // swap is a no-op on the underlying buffer.
      std::vector<float> x_layout0_data(x_batchwise_data.size());
      for (int64_t n = 0; n < batch_size; ++n) {
        for (int64_t t = 0; t < seq_length; ++t) {
          for (int64_t k = 0; k < input_size; ++k) {
            const size_t src = static_cast<size_t>((n * seq_length + t) * input_size + k);
            const size_t dst = static_cast<size_t>((t * batch_size + n) * input_size + k);
            x_layout0_data[dst] = x_batchwise_data[src];
          }
        }
      }
      Tensor x_layout0 =
          Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_layout0_data);

      std::vector<float> w_data(static_cast<size_t>(kNumGates * hidden_size * input_size),
                                weight_scale);
      std::vector<float> r_data(static_cast<size_t>(kNumGates * hidden_size * hidden_size),
                                weight_scale);
      Tensor w = Tensor::FromFloat("", {1, kNumGates * hidden_size, input_size}, w_data);
      Tensor r = Tensor::FromFloat("", {1, kNumGates * hidden_size, hidden_size}, r_data);

      auto [y_layout0, y_h_layout0] = gru_kernel(x_layout0, w, r);

      // Permute Y: [seq, 1, batch, hidden] -> [batch, seq, 1, hidden].
      std::vector<float> y_batchwise_data(
          static_cast<size_t>(batch_size * seq_length * hidden_size));
      const float *py0 = y_layout0.AsFloat();
      for (int64_t t = 0; t < seq_length; ++t) {
        for (int64_t n = 0; n < batch_size; ++n) {
          for (int64_t h = 0; h < hidden_size; ++h) {
            const size_t src = static_cast<size_t>((t * batch_size + n) * hidden_size + h);
            const size_t dst = static_cast<size_t>((n * seq_length + t) * hidden_size + h);
            y_batchwise_data[dst] = py0[src];
          }
        }
      }
      Tensor y_batchwise =
          Tensor::FromFloat("", {batch_size, seq_length, 1, hidden_size}, y_batchwise_data);

      // Permute Y_h: [1, batch, hidden] -> [batch, 1, hidden]. With
      // num_directions=1 the leading 1-axis simply moves to the middle, so
      // the underlying buffer contents are unchanged.
      const float *pyh0 = y_h_layout0.AsFloat();
      std::vector<float> y_h_batchwise_data(pyh0,
                                            pyh0 + static_cast<size_t>(batch_size * hidden_size));
      Tensor y_h_batchwise =
          Tensor::FromFloat("", {batch_size, 1, hidden_size}, y_h_batchwise_data);

      return IoData{{std::move(x_batchwise), std::move(w), std::move(r)},
                    {std::move(y_batchwise), std::move(y_h_batchwise)}};
    });
  }

  // ``gru_reverse``: single ``reverse`` direction, seq_length=3,
  // batch_size=1, input_size=2, hidden_size=5, mirroring upstream
  // ``test_gru_reverse``. Only ``Y_h`` is produced.
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_output("");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 5);
    AddAttribute(node, "direction", std::string("reverse"));
    Expect(registry, std::move(node), "test_cc_gru_reverse", {opset}, [=]() -> IoData {
      const int64_t seq_length = 3;
      const int64_t batch_size = 1;
      const int64_t input_size = 2;
      const int64_t hidden_size = 5;
      const std::vector<float> x_data{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
      const float weight_scale = 0.1f;
      std::vector<float> w_data(static_cast<size_t>(kNumGates * hidden_size * input_size),
                                weight_scale);
      std::vector<float> r_data(static_cast<size_t>(kNumGates * hidden_size * hidden_size),
                                weight_scale);
      Tensor x = Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_data);
      Tensor w = Tensor::FromFloat("", {1, kNumGates * hidden_size, input_size}, w_data);
      Tensor r = Tensor::FromFloat("", {1, kNumGates * hidden_size, hidden_size}, r_data);

      auto [y_unused, y_h] = gru_kernel(x, w, r, Tensor{}, Tensor{}, 0, 0, "reverse");
      (void)y_unused; // Y is skipped (empty output name).

      return IoData{{std::move(x), std::move(w), std::move(r)}, {std::move(y_h)}};
    });
  }

  // ``gru_bidirectional``: ``bidirectional`` direction (num_directions=2),
  // seq_length=3, batch_size=1, input_size=2, hidden_size=5, mirroring
  // upstream ``test_gru_bidirectional``. Both ``Y`` and ``Y_h`` are produced.
  {
    NodeProto node;
    node.set_op_type("GRU");
    node.add_input("X");
    node.add_input("W");
    node.add_input("R");
    node.add_output("Y");
    node.add_output("Y_h");
    AddAttribute<int64_t>(node, "hidden_size", 5);
    AddAttribute(node, "direction", std::string("bidirectional"));
    Expect(registry, std::move(node), "test_cc_gru_bidirectional", {opset}, [=]() -> IoData {
      const int64_t num_directions = 2;
      const int64_t seq_length = 3;
      const int64_t batch_size = 1;
      const int64_t input_size = 2;
      const int64_t hidden_size = 5;
      const std::vector<float> x_data{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
      std::vector<float> w_data(
          static_cast<size_t>(num_directions * kNumGates * hidden_size * input_size));
      std::vector<float> r_data(
          static_cast<size_t>(num_directions * kNumGates * hidden_size * hidden_size));
      const float scales[2] = {0.1f, 0.2f};
      const int64_t w_block = kNumGates * hidden_size * input_size;
      const int64_t r_block = kNumGates * hidden_size * hidden_size;
      for (int64_t d = 0; d < num_directions; ++d) {
        for (int64_t i = 0; i < w_block; ++i) {
          w_data[static_cast<size_t>(d * w_block + i)] = scales[d];
        }
        for (int64_t i = 0; i < r_block; ++i) {
          r_data[static_cast<size_t>(d * r_block + i)] = scales[d];
        }
      }
      Tensor x = Tensor::FromFloat("", {seq_length, batch_size, input_size}, x_data);
      Tensor w =
          Tensor::FromFloat("", {num_directions, kNumGates * hidden_size, input_size}, w_data);
      Tensor r =
          Tensor::FromFloat("", {num_directions, kNumGates * hidden_size, hidden_size}, r_data);

      auto [y, y_h] = gru_kernel(x, w, r, Tensor{}, Tensor{}, 0, 0, "bidirectional");

      return IoData{{std::move(x), std::move(w), std::move(r)}, {std::move(y), std::move(y_h)}};
    });
  }
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
