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

#include "onnx_core/backend_test/expect.h"
#include "onnx_core/runtime/kernels/controlflow/include_controlflow_kernels.h"
#include "onnx_core/runtime/kernels/random.h"
#include "onnx_extensions/backend_test/cases/controlflow/include_controlflow_cases.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <string>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test {

namespace {

// Returns the bytes of an INT64 scalar with value ``v`` in little-endian
// order (matching ``Tensor::data`` layout).
std::vector<uint8_t> Int64Bytes(int64_t v) {
  std::vector<uint8_t> out(sizeof(int64_t));
  std::memcpy(out.data(), &v, sizeof(int64_t));
  return out;
}

// Adds a graph input named ``name`` with a tensor type ``dtype`` and a
// scalar shape (empty dims).
void AddGraphInputScalar(GraphProto &g, const std::string &name, TensorProto::DataType dtype) {
  ValueInfoProto *vi = g.add_input();
  vi->set_name(name);
  TypeProto::Tensor *tensor_type = vi->ref_type().mutable_tensor_type();
  tensor_type->set_elem_type(static_cast<int>(dtype));
  tensor_type->mutable_shape();
}

// Adds a graph input named ``name`` with a tensor type ``dtype`` and the
// given (non-scalar) ``shape``.
void AddGraphInputTensor(GraphProto &g, const std::string &name, TensorProto::DataType dtype,
                         const std::vector<int64_t> &shape) {
  ValueInfoProto *vi = g.add_input();
  vi->set_name(name);
  TypeProto::Tensor *tensor_type = vi->ref_type().mutable_tensor_type();
  tensor_type->set_elem_type(static_cast<int>(dtype));
  TensorShapeProto *s = tensor_type->mutable_shape();
  for (int64_t d : shape) {
    s->add_dim()->set_dim_value(d);
  }
}

void AddGraphOutputTensor(GraphProto &g, const std::string &name, TensorProto::DataType dtype) {
  ValueInfoProto *vi = g.add_output();
  vi->set_name(name);
  TypeProto::Tensor *tensor_type = vi->ref_type().mutable_tensor_type();
  tensor_type->set_elem_type(static_cast<int>(dtype));
  tensor_type->mutable_shape();
}

// Builds a Loop body subgraph for our simple test case:
//
//   inputs : (i [INT64 scalar], cond_in [BOOL scalar])
//   nodes  : identity-on-cond + a Constant node emitting [42] (scan_output)
//   outputs: (cond_out=cond_in, scan_out [INT64 shape=[1]])
//
// The body has no loop-carried dependencies (N == 0) and produces one scan
// output (K == 1). The reference kernel does not execute the body; it is
// included in the model purely so the registered ``TestCase`` is a
// well-formed ONNX model with a valid ``Loop`` node.
GraphProto BuildSimpleLoopBody() {
  GraphProto g;
  g.set_name("loop_body");

  AddGraphInputScalar(g, "i", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);

  // cond_out = Identity(cond_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }
  // scan_out = Constant(value=int64 [42])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("scan_out");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::INT64);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(Int64Bytes(42)));
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  // scan_out: the Constant node above produces an INT64 tensor of shape [1].
  ValueInfoProto *scan_vi = g.add_output();
  scan_vi->set_name("scan_out");
  TypeProto::Tensor *scan_tt = scan_vi->ref_type().mutable_tensor_type();
  scan_tt->set_elem_type(static_cast<int>(TensorProto::DataType::INT64));
  scan_tt->ref_shape().add_dim()->set_dim_value(1);
  return g;
}

// Builds a ``Loop`` node with the given input/output names and a body
// subgraph attribute.
NodeProto MakeLoopNode(const std::string &m, const std::string &cond,
                       const std::string &scan_output_name) {
  NodeProto node;
  node.set_op_type("Loop");
  node.add_input(m);
  node.add_input(cond);
  // No v_initial inputs (N == 0): the Loop node has 2 inputs.
  node.add_output(scan_output_name);

  AttributeProto *body_attr = node.add_attribute();
  body_attr->set_name("body");
  body_attr->set_type(AttributeProto::AttributeType::GRAPH);
  *body_attr->add_g() = BuildSimpleLoopBody();
  return node;
}

} // namespace

// ---------------------------------------------------------------------------
// Loop — loop11 mirror.
//
// Mirrors ONNX's ``test_loop11`` node case: trip-count 5 with a single
// FLOAT[1] loop-carried dependency that accumulates a prefix sum of
// ``x = [1, 2, 3, 4, 5]`` starting from ``y = [-2]`` and a single
// FLOAT[1] scan output emitting the running sum at each iteration.
//
//   res_y    = [13]
//   res_scan = [[-1], [1], [4], [8], [13]]
//
// The body subgraph is well-formed and now also numerically correct:
// ``y_out = y_in + x[iter_count]`` with ``x = [1, 2, 3, 4, 5]`` as a body
// Constant, and ``scan_out = y_out``. The ``Loop`` reference does
// not execute the body, but external backends (e.g. onnxruntime) do — so
// the body must produce the same expected outputs registered below.
// ---------------------------------------------------------------------------
namespace {

// Returns the little-endian bytes of a FLOAT value (matching the layout of
// ``Tensor::data`` for ``DataType::FLOAT``).
std::vector<uint8_t> FloatBytes(float v) {
  std::vector<uint8_t> out(sizeof(float));
  std::memcpy(out.data(), &v, sizeof(float));
  return out;
}

// Returns the little-endian bytes of ``values`` packed as FLOAT.
std::vector<uint8_t> FloatBytes(const std::vector<float> &values) {
  std::vector<uint8_t> out(values.size() * sizeof(float));
  std::memcpy(out.data(), values.data(), out.size());
  return out;
}

// Body subgraph for the loop11 case: 3 inputs (iter_count, cond_in, y_in),
// 3 outputs (cond_out, y_out, scan_out). The body computes
// ``y_out = y_in + x[iter_count]`` where ``x = [1, 2, 3, 4, 5]`` is a
// per-body Constant, and emits ``y_out`` as the scan output. With initial
// ``y = [-2]`` this produces the expected prefix-sum sequence
// ``[[-1], [1], [4], [8], [13]]`` and final ``y = [13]``.
GraphProto BuildLoop11Body() {
  GraphProto g;
  g.set_name("loop11_body");

  AddGraphInputScalar(g, "iter_count", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);
  AddGraphInputTensor(g, "y_in", TensorProto::DataType::FLOAT, {1});

  // cond_out = Identity(cond_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }
  // x = Constant(value=float [1, 2, 3, 4, 5])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("x");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(5);
    t->set_raw_data(utils::ByteSpan(FloatBytes({1.0f, 2.0f, 3.0f, 4.0f, 5.0f})));
  }
  // iter_1d = Unsqueeze(iter_count, axes=[0]) — turn the INT64 scalar into
  // an INT64 tensor of shape [1] suitable as Gather indices producing a
  // FLOAT[1] output.
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Unsqueeze");
    n->add_input("iter_count");
    n->add_output("iter_1d");
    AttributeProto *a = n->add_attribute();
    a->set_name("axes");
    a->set_type(AttributeProto::AttributeType::INTS);
    a->add_ints(static_cast<int64_t>(0));
  }
  // x_i = Gather(x, iter_1d, axis=0) — FLOAT[1].
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Gather");
    n->add_input("x");
    n->add_input("iter_1d");
    n->add_output("x_i");
    AttributeProto *a = n->add_attribute();
    a->set_name("axis");
    a->set_type(AttributeProto::AttributeType::INT);
    a->set_i(0);
  }
  // y_out = Add(y_in, x_i)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Add");
    n->add_input("y_in");
    n->add_input("x_i");
    n->add_output("y_out");
  }
  // scan_out = Identity(y_out)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("y_out");
    n->add_output("scan_out");
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  // y_out: FLOAT[1]
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("y_out");
    TypeProto::Tensor *tt = vi->ref_type().mutable_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->ref_shape().add_dim()->set_dim_value(1);
  }
  // scan_out: FLOAT[1]
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("scan_out");
    TypeProto::Tensor *tt = vi->ref_type().mutable_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->ref_shape().add_dim()->set_dim_value(1);
  }
  return g;
}

} // namespace

static void RegisterLoop11Case(std::vector<TestCase> &registry) {
  // ONNX's ``test_loop11`` model imports opset 11 — match that here so the
  // case can be located by substring against ONNX's test name.
  const OpsetId opset = DefaultOpset(11);
  const KernelContext ctx{opset};
  const Loop loop_kernel{ctx};

  // Build the Loop node: 3 inputs (M, cond, y) and 2 outputs (res_y, res_scan).
  NodeProto node;
  node.set_op_type("Loop");
  node.add_input("trip_count");
  node.add_input("cond");
  node.add_input("y");
  node.add_output("res_y");
  node.add_output("res_scan");

  AttributeProto *body_attr = node.add_attribute();
  body_attr->set_name("body");
  body_attr->set_type(AttributeProto::AttributeType::GRAPH);
  *body_attr->add_g() = BuildLoop11Body();

  Expect(registry, std::move(node), "test_cc_loop11_carried_state", {opset},
         [loop_kernel]() -> IoData {
           // Inputs to the Loop node.
           Tensor trip_count("trip_count", DataType::INT64, {}, Int64Bytes(5));
           Tensor cond("cond", DataType::BOOL, {}, std::vector<uint8_t>{1});
           Tensor y("y", DataType::FLOAT, {1}, FloatBytes(-2.0f));

           // Per-iteration scan-output values, matching ONNX's ``test_loop11``
           // expected ``res_scan = [[-1], [1], [4], [8], [13]]``.
           const std::vector<Tensor> per_iter_scan = {
               Tensor("", DataType::FLOAT, {1}, FloatBytes(-1.0f)),
               Tensor("", DataType::FLOAT, {1}, FloatBytes(1.0f)),
               Tensor("", DataType::FLOAT, {1}, FloatBytes(4.0f)),
               Tensor("", DataType::FLOAT, {1}, FloatBytes(8.0f)),
               Tensor("", DataType::FLOAT, {1}, FloatBytes(13.0f)),
           };

           // Final loop-carried state value matches ``res_y = [13]``.
           const Tensor res_y("", DataType::FLOAT, {1}, FloatBytes(13.0f));

           std::vector<Tensor> out = loop_kernel(trip_count, cond, /*v_initial=*/{y},
                                                 /*final_state=*/{res_y}, {per_iter_scan});
           return IoData{{std::move(trip_count), std::move(cond), std::move(y)}, std::move(out)};
         });
}

namespace {

// ---------------------------------------------------------------------------
// Loop — minimal cases exercising the reference kernel's stacking semantics.
//
// Two registered cases use a Loop node with no loop-carried dependencies
// (N=0) and a single scan output (K=1) producing the constant int64 ``[42]``
// per iteration. They differ only in the value of ``M``: 3 (the scan output
// is ``[3, 1]`` containing ``[[42], [42], [42]]``) and 0 (the scan output
// is empty along its leading axis, ``[0, 1]``).
// ---------------------------------------------------------------------------
void RegisterTripCountVariants(std::vector<TestCase> &registry) {
  const OpsetId opset = DefaultOpset(13);
  const KernelContext ctx{opset};
  const Loop loop_kernel{ctx};

  const Tensor per_iter_value("", DataType::INT64, {1}, Int64Bytes(42));
  const Tensor cond_undef; // omitted scalar BOOL cond.

  auto register_case = [&](const std::string &test_name, int64_t m_value) {
    NodeProto node = MakeLoopNode("M", "", "scan_outputs");
    Expect(registry, std::move(node), test_name, {opset}, [=]() -> IoData {
      const Tensor m("", DataType::INT64, {}, Int64Bytes(m_value));
      // Always supply at least one per-iteration template tensor so the
      // stacked output keeps its dtype/trailing-shape even when the loop
      // runs zero iterations.
      const std::size_t row_len = std::max<std::size_t>(1, static_cast<std::size_t>(m_value));
      std::vector<Tensor> per_iter(row_len, per_iter_value);
      std::vector<Tensor> out =
          loop_kernel(m, cond_undef, /*v_initial=*/{}, /*final_state=*/{}, {per_iter});
      return IoData{{std::move(m)}, std::move(out)};
    });
  };

  // M = 3 → stacked scan output has shape [3, 1] = [[42], [42], [42]].
  register_case("test_cc_loop_basic_trip_count", 3);
  // M = 0 → stacked scan output has shape [0, 1] (empty along axis 0).
  register_case("test_cc_loop_zero_trip_count", 0);
}

} // namespace

// ---------------------------------------------------------------------------
// Loop — multi loop-carried state with a scan output.
//
// Trip count 3, two FLOAT[1] loop-carried dependencies ``y`` (incremented
// by 1 each iter) and ``w`` (doubled each iter), and a single FLOAT[1] scan
// output emitting ``y`` per iteration. Starting from ``y = [0]`` and
// ``w = [1]`` the expected outputs are:
//
//   res_y    = [3]
//   res_w    = [8]
//   res_scan = [[1], [2], [3]]
//
// This exercises ``N == 2`` loop-carried dependencies (the loop11 mirror
// uses ``N == 1``).
// ---------------------------------------------------------------------------
namespace {

// Body for the multi-carried case:
//   inputs:  iter [INT64 scalar], cond_in [BOOL scalar],
//            y_in [FLOAT[1]], w_in [FLOAT[1]]
//   outputs: cond_out [BOOL scalar], y_out [FLOAT[1]],
//            w_out [FLOAT[1]], scan_out [FLOAT[1]]
GraphProto BuildMultiCarriedBody() {
  GraphProto g;
  g.set_name("loop_multi_carried_body");

  AddGraphInputScalar(g, "iter", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);
  AddGraphInputTensor(g, "y_in", TensorProto::DataType::FLOAT, {1});
  AddGraphInputTensor(g, "w_in", TensorProto::DataType::FLOAT, {1});

  // cond_out = Identity(cond_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }
  // one = Constant(value=float [1.0])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("one");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(FloatBytes(1.0f)));
  }
  // two = Constant(value=float [2.0])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("two");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(FloatBytes(2.0f)));
  }
  // y_out = Add(y_in, one)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Add");
    n->add_input("y_in");
    n->add_input("one");
    n->add_output("y_out");
  }
  // w_out = Mul(w_in, two)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Mul");
    n->add_input("w_in");
    n->add_input("two");
    n->add_output("w_out");
  }
  // scan_out = Identity(y_out)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("y_out");
    n->add_output("scan_out");
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  for (const char *name : {"y_out", "w_out", "scan_out"}) {
    ValueInfoProto *vi = g.add_output();
    vi->set_name(name);
    TypeProto::Tensor *tt = vi->ref_type().mutable_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->ref_shape().add_dim()->set_dim_value(1);
  }
  return g;
}

void RegisterMultiCarriedCase(std::vector<TestCase> &registry) {
  const OpsetId opset = DefaultOpset(13);
  const KernelContext ctx{opset};
  const Loop loop_kernel{ctx};

  NodeProto node;
  node.set_op_type("Loop");
  node.add_input("trip_count");
  node.add_input("cond");
  node.add_input("y");
  node.add_input("w");
  node.add_output("res_y");
  node.add_output("res_w");
  node.add_output("res_scan");

  AttributeProto *body_attr = node.add_attribute();
  body_attr->set_name("body");
  body_attr->set_type(AttributeProto::AttributeType::GRAPH);
  *body_attr->add_g() = BuildMultiCarriedBody();

  Expect(registry, std::move(node), "test_cc_loop_multi_carried_state", {opset},
         [loop_kernel]() -> IoData {
           Tensor trip_count("trip_count", DataType::INT64, {}, Int64Bytes(3));
           Tensor cond("cond", DataType::BOOL, {}, std::vector<uint8_t>{1});
           Tensor y("y", DataType::FLOAT, {1}, FloatBytes(0.0f));
           Tensor w("w", DataType::FLOAT, {1}, FloatBytes(1.0f));

           // Drive the body-runner overload so the registered expected outputs are
           // produced by the same code path that ``RunLoopNode`` exercises at
           // runtime.
           Loop::BodyRunner runner = [](int64_t /*iter*/, bool cond_in,
                                        const std::vector<Tensor> &state) -> std::vector<Tensor> {
             const float y_in = state[0].AsFloat()[0];
             const float w_in = state[1].AsFloat()[0];
             const float y_out = y_in + 1.0f;
             const float w_out = w_in * 2.0f;
             return {Tensor("", DataType::BOOL, {},
                            std::vector<uint8_t>{static_cast<uint8_t>(cond_in)}),
                     Tensor("", DataType::FLOAT, {1}, FloatBytes(y_out)),
                     Tensor("", DataType::FLOAT, {1}, FloatBytes(w_out)),
                     Tensor("", DataType::FLOAT, {1}, FloatBytes(y_out))};
           };

           std::vector<Tensor> out = loop_kernel(trip_count, cond, /*v_initial=*/{y, w},
                                                 /*num_scan_outputs=*/1, runner);
           return IoData{{std::move(trip_count), std::move(cond), std::move(y), std::move(w)},
                         std::move(out)};
         });
}

} // namespace

// ---------------------------------------------------------------------------
// Loop — opset 13 with a sequence-typed loop-carried state.
//
// Mirrors ONNX's ``test_loop13_seq`` in spirit: trip-count 5 with an
// empty input sequence as loop-carried state and a body that appends a
// per-iteration tensor via ``SequenceInsert``. Upstream uses
// ``x[:i+1]`` (variable per-iteration shapes ``[1]..[5]``) which the
// project's ``DataSet`` (uniform-shape ``Tensor``s) cannot represent, so
// the C++ case appends ``Gather(x, [i])`` instead — a fixed ``FLOAT[1]``
// element per iteration. The resulting sequence is materialized as a
// stacked ``FLOAT[5, 1]`` tensor via ``ConcatFromSequence(new_axis=1)``
// for byte comparison by the test harness, yielding
// ``[[1], [2], [3], [4], [5]]``.
//
// The graph topology is therefore: ``SequenceEmpty`` → ``Loop`` →
// ``ConcatFromSequence``. The body subgraph computes
// ``seq_out = SequenceInsert(seq_in, Gather(x, Unsqueeze(iter, [0])))``
// where ``x = [1, 2, 3, 4, 5]`` is a body-local ``Constant``.
// ---------------------------------------------------------------------------
namespace {

// Body subgraph for the loop13_seq case. Inputs are
// ``(iter [INT64 scalar], cond_in [BOOL scalar], seq_in [Sequence<FLOAT, [1]>])``;
// outputs are ``(cond_out [BOOL scalar], seq_out [Sequence<FLOAT, [1]>])``.
GraphProto BuildLoop13SeqBody() {
  GraphProto g;
  g.set_name("loop13_seq_body");

  AddGraphInputScalar(g, "iter", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);
  // seq_in: Sequence<Tensor<FLOAT, [1]>>.
  {
    ValueInfoProto *vi = g.add_input();
    vi->set_name("seq_in");
    TypeProto *tp = vi->ref_type().add_sequence_type()->add_elem_type();
    TypeProto::Tensor *tt = tp->add_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->add_shape()->add_dim()->set_dim_value(1);
  }

  // cond_out = Identity(cond_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }
  // x = Constant(value=float [1, 2, 3, 4, 5])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("x");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(5);
    t->set_raw_data(utils::ByteSpan(FloatBytes({1.0f, 2.0f, 3.0f, 4.0f, 5.0f})));
  }
  // axes_const = Constant(value=int64 [0])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("axes_const");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::INT64);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(Int64Bytes(0)));
  }
  // iter_1d = Unsqueeze(iter, axes_const) — INT64 scalar → INT64[1].
  // Opset 13 takes axes as an input, not an attribute.
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Unsqueeze");
    n->add_input("iter");
    n->add_input("axes_const");
    n->add_output("iter_1d");
  }
  // x_i = Gather(x, iter_1d, axis=0) — FLOAT[1].
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Gather");
    n->add_input("x");
    n->add_input("iter_1d");
    n->add_output("x_i");
    AttributeProto *a = n->add_attribute();
    a->set_name("axis");
    a->set_type(AttributeProto::AttributeType::INT);
    a->set_i(0);
  }
  // seq_out = SequenceInsert(seq_in, x_i)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("SequenceInsert");
    n->add_input("seq_in");
    n->add_input("x_i");
    n->add_output("seq_out");
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  // seq_out: Sequence<Tensor<FLOAT, [1]>>.
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("seq_out");
    TypeProto *tp = vi->ref_type().add_sequence_type()->add_elem_type();
    TypeProto::Tensor *tt = tp->add_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->add_shape()->add_dim()->set_dim_value(1);
  }
  return g;
}

void RegisterLoop13SeqCase(std::vector<TestCase> &registry) {
  const std::string name = "test_cc_loop13_seq";
  const OpsetId opset = DefaultOpset(13);

  // Inputs: ``trip_count = 5`` and ``cond = true``. The empty input
  // sequence is materialised inside the graph via ``SequenceEmpty`` so
  // the DataSet (Tensor-only) does not need to carry a sequence value.
  const Tensor trip_count("trip_count", DataType::INT64, {}, Int64Bytes(5));
  const Tensor cond("cond", DataType::BOOL, {}, std::vector<uint8_t>{1});

  // Expected stacked output: ``[[1], [2], [3], [4], [5]]`` of shape [5, 1].
  Tensor stacked("res_stacked", DataType::FLOAT, {5, 1},
                 FloatBytes({1.0f, 2.0f, 3.0f, 4.0f, 5.0f}));

  TestCase tc(name, name);
  ModelProto &model = tc.emplace_model();
  InitModel(model, /*ir_version=*/13, {opset});
  GraphProto *graph = model.add_graph();
  graph->set_name(name);

  // Node 1: seq_empty = SequenceEmpty(dtype=FLOAT)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("SequenceEmpty");
    n->add_output("seq_empty");
    AttributeProto *a = n->add_attribute();
    a->set_name("dtype");
    a->set_type(AttributeProto::AttributeType::INT);
    a->set_i(static_cast<int64_t>(TensorProto::DataType::FLOAT));
  }
  // Node 2: seq_res = Loop(trip_count, cond, seq_empty)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("Loop");
    n->add_input("trip_count");
    n->add_input("cond");
    n->add_input("seq_empty");
    n->add_output("seq_res");
    AttributeProto *a = n->add_attribute();
    a->set_name("body");
    a->set_type(AttributeProto::AttributeType::GRAPH);
    *a->add_g() = BuildLoop13SeqBody();
  }
  // Node 3: res_stacked = ConcatFromSequence(seq_res, axis=0, new_axis=1)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("ConcatFromSequence");
    n->add_input("seq_res");
    n->add_output("res_stacked");
    AttributeProto *axis_attr = n->add_attribute();
    axis_attr->set_name("axis");
    axis_attr->set_type(AttributeProto::AttributeType::INT);
    axis_attr->set_i(0);
    AttributeProto *na_attr = n->add_attribute();
    na_attr->set_name("new_axis");
    na_attr->set_type(AttributeProto::AttributeType::INT);
    na_attr->set_i(1);
  }

  // Graph inputs/outputs.
  FillValueInfo(trip_count, *graph->add_input());
  FillValueInfo(cond, *graph->add_input());
  FillValueInfo(stacked, *graph->add_output());

  DataSet ds;
  ds.inputs.push_back(trip_count);
  ds.inputs.push_back(cond);
  ds.outputs.push_back(stacked);
  tc.data_sets().emplace_back(std::move(ds));

  registry.emplace_back(std::move(tc));
}

// ---------------------------------------------------------------------------
// ``test_cc_loop16_seq_none`` — Mirrors ONNX's ``test_loop16_seq_none``.
// A Loop whose body receives an Optional<Sequence<FLOAT, [1]>> as
// loop-carried state.  On the first iteration the Optional is present
// (contains a sequence ``[[0]]``), so the ``If`` "else" branch extracts it
// via ``OptionalGetElement``. On subsequent iterations the state is a plain
// sequence (the Optional has an element). Each iteration appends
// ``Gather(x, [iter])`` where ``x = [1, 2, 3, 4, 5]``.  After 5 iterations
// the output sequence is ``[[0], [1], [2], [3], [4], [5]]``, stacked to
// ``FLOAT[6, 1]`` via ``ConcatFromSequence(new_axis=1)``.
//
// This exercises OptionalHasElement, Not, If (with subgraphs),
// OptionalGetElement, SequenceConstruct, SequenceInsert inside a Loop body.
// ---------------------------------------------------------------------------

// Body subgraph for loop16_seq_none.
// Inputs: (iter [INT64 scalar], cond_in [BOOL scalar],
//          opt_seq_in [Optional<Sequence<FLOAT, [1]>>])
// Outputs: (cond_out [BOOL scalar], seq_out [Sequence<FLOAT, [1]>])
GraphProto BuildLoop16SeqNoneBody() {
  GraphProto g;
  g.set_name("loop16_seq_none_body");

  AddGraphInputScalar(g, "iter_count", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);
  // opt_seq_in: Optional<Sequence<Tensor<FLOAT, [1]>>>
  {
    ValueInfoProto *vi = g.add_input();
    vi->set_name("opt_seq_in");
    TypeProto *opt_tp = vi->ref_type().add_optional_type()->add_elem_type();
    TypeProto *seq_tp = opt_tp->add_sequence_type()->add_elem_type();
    TypeProto::Tensor *tt = seq_tp->add_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->add_shape()->add_dim()->set_dim_value(1);
  }

  // cond_out = Identity(cond_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }

  // optional_has_elem = OptionalHasElement(opt_seq_in)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("OptionalHasElement");
    n->add_input("opt_seq_in");
    n->add_output("optional_has_elem");
  }

  // optional_is_none = Not(optional_has_elem)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Not");
    n->add_input("optional_has_elem");
    n->add_output("optional_is_none");
  }

  // If(optional_is_none):
  //   then: SequenceConstruct(Constant(0.0)) → init_seq_in
  //   else: OptionalGetElement(opt_seq_in) → seq_in
  {
    // --- then_body ---
    GraphProto then_body;
    then_body.set_name("then_body");
    {
      NodeProto *cn = then_body.add_node();
      cn->set_op_type("Constant");
      cn->add_output("constant_in");
      AttributeProto *a = cn->add_attribute();
      a->set_name("value");
      a->set_type(AttributeProto::AttributeType::TENSOR);
      TensorProto *t = a->add_t();
      t->set_data_type(TensorProto::DataType::FLOAT);
      t->add_dims(1);
      t->set_raw_data(utils::ByteSpan(FloatBytes(0.0f)));
    }
    {
      NodeProto *sc = then_body.add_node();
      sc->set_op_type("SequenceConstruct");
      sc->add_input("constant_in");
      sc->add_output("init_seq_in");
    }
    {
      ValueInfoProto *vi = then_body.add_output();
      vi->set_name("init_seq_in");
      TypeProto *tp = vi->ref_type().add_sequence_type()->add_elem_type();
      TypeProto::Tensor *tt = tp->add_tensor_type();
      tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
      tt->add_shape()->add_dim()->set_dim_value(1);
    }

    // --- else_body ---
    GraphProto else_body;
    else_body.set_name("else_body");
    {
      NodeProto *ge = else_body.add_node();
      ge->set_op_type("OptionalGetElement");
      ge->add_input("opt_seq_in");
      ge->add_output("seq_in");
    }
    {
      ValueInfoProto *vi = else_body.add_output();
      vi->set_name("seq_in");
      TypeProto *tp = vi->ref_type().add_sequence_type()->add_elem_type();
      TypeProto::Tensor *tt = tp->add_tensor_type();
      tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
      tt->add_shape()->add_dim()->set_dim_value(1);
    }

    NodeProto *if_node = g.add_node();
    if_node->set_op_type("If");
    if_node->add_input("optional_is_none");
    if_node->add_output("sequence");
    {
      AttributeProto *a = if_node->add_attribute();
      a->set_name("then_branch");
      a->set_type(AttributeProto::AttributeType::GRAPH);
      *a->add_g() = then_body;
    }
    {
      AttributeProto *a = if_node->add_attribute();
      a->set_name("else_branch");
      a->set_type(AttributeProto::AttributeType::GRAPH);
      *a->add_g() = else_body;
    }
  }

  // x = Constant(value=float [1, 2, 3, 4, 5])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("x");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(5);
    t->set_raw_data(utils::ByteSpan(FloatBytes({1.0f, 2.0f, 3.0f, 4.0f, 5.0f})));
  }

  // one = Constant(value=int64 1)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("one");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::INT64);
    t->set_raw_data(utils::ByteSpan(Int64Bytes(1)));
  }

  // axes_const = Constant(value=int64 [0])
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Constant");
    n->add_output("axes_const");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::INT64);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(Int64Bytes(0)));
  }

  // end = Add(iter_count, one)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Add");
    n->add_input("iter_count");
    n->add_input("one");
    n->add_output("end");
  }

  // iter_1d = Unsqueeze(iter_count, axes_const)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Unsqueeze");
    n->add_input("iter_count");
    n->add_input("axes_const");
    n->add_output("iter_1d");
  }

  // x_i = Gather(x, iter_1d, axis=0) → FLOAT[1]
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Gather");
    n->add_input("x");
    n->add_input("iter_1d");
    n->add_output("x_i");
    AttributeProto *a = n->add_attribute();
    a->set_name("axis");
    a->set_type(AttributeProto::AttributeType::INT);
    a->set_i(0);
  }

  // seq_out = SequenceInsert(sequence, x_i)
  {
    NodeProto *n = g.add_node();
    n->set_op_type("SequenceInsert");
    n->add_input("sequence");
    n->add_input("x_i");
    n->add_output("seq_out");
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  // seq_out: Sequence<Tensor<FLOAT, [1]>>
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("seq_out");
    TypeProto *tp = vi->ref_type().add_sequence_type()->add_elem_type();
    TypeProto::Tensor *tt = tp->add_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    tt->add_shape()->add_dim()->set_dim_value(1);
  }
  return g;
}

void RegisterLoop16SeqNoneCase(std::vector<TestCase> &registry) {
  const std::string name = "test_cc_loop16_seq_none";
  const OpsetId opset = DefaultOpset(16);

  // Inputs: trip_count=5, cond=true.
  // The initial optional sequence [[0]] is constructed via
  // SequenceConstruct + Optional inside the graph.
  const Tensor trip_count("trip_count", DataType::INT64, {}, Int64Bytes(5));
  const Tensor cond("cond", DataType::BOOL, {}, std::vector<uint8_t>{1});

  // Expected stacked output: [[0], [1], [2], [3], [4], [5]] → FLOAT[6, 1].
  Tensor stacked("res_stacked", DataType::FLOAT, {6, 1},
                 FloatBytes({0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f}));

  TestCase tc(name, name);
  ModelProto &model = tc.emplace_model();
  InitModel(model, /*ir_version=*/9, {opset});
  GraphProto *graph = model.add_graph();
  graph->set_name(name);

  // Node 1: init_elem = Constant(value=float [0])
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("Constant");
    n->add_output("init_elem");
    AttributeProto *a = n->add_attribute();
    a->set_name("value");
    a->set_type(AttributeProto::AttributeType::TENSOR);
    TensorProto *t = a->add_t();
    t->set_data_type(TensorProto::DataType::FLOAT);
    t->add_dims(1);
    t->set_raw_data(utils::ByteSpan(FloatBytes(0.0f)));
  }

  // Node 2: init_seq = SequenceConstruct(init_elem)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("SequenceConstruct");
    n->add_input("init_elem");
    n->add_output("init_seq");
  }

  // Node 3: opt_seq = Optional(init_seq)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("Optional");
    n->add_input("init_seq");
    n->add_output("opt_seq");
  }

  // Node 4: seq_res = Loop(trip_count, cond, opt_seq)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("Loop");
    n->add_input("trip_count");
    n->add_input("cond");
    n->add_input("opt_seq");
    n->add_output("seq_res");
    AttributeProto *a = n->add_attribute();
    a->set_name("body");
    a->set_type(AttributeProto::AttributeType::GRAPH);
    *a->add_g() = BuildLoop16SeqNoneBody();
  }

  // Node 5: res_stacked = ConcatFromSequence(seq_res, axis=0, new_axis=1)
  {
    NodeProto *n = graph->add_node();
    n->set_op_type("ConcatFromSequence");
    n->add_input("seq_res");
    n->add_output("res_stacked");
    AttributeProto *axis_attr = n->add_attribute();
    axis_attr->set_name("axis");
    axis_attr->set_type(AttributeProto::AttributeType::INT);
    axis_attr->set_i(0);
    AttributeProto *na_attr = n->add_attribute();
    na_attr->set_name("new_axis");
    na_attr->set_type(AttributeProto::AttributeType::INT);
    na_attr->set_i(1);
  }

  // Graph inputs/outputs.
  FillValueInfo(trip_count, *graph->add_input());
  FillValueInfo(cond, *graph->add_input());
  FillValueInfo(stacked, *graph->add_output());

  DataSet ds;
  ds.inputs.push_back(trip_count);
  ds.inputs.push_back(cond);
  ds.outputs.push_back(stacked);
  tc.data_sets().emplace_back(std::move(ds));

  registry.emplace_back(std::move(tc));
}

// Builds a Loop body that carries a large FLOAT tensor unchanged (identity),
// emitting it as the per-iteration scan output. Used by the benchmark case.
GraphProto BuildLoopBenchmarkBody(const std::vector<int64_t> &shape) {
  GraphProto g;
  g.set_name("loop_benchmark_body");

  AddGraphInputScalar(g, "iter_count", TensorProto::DataType::INT64);
  AddGraphInputScalar(g, "cond_in", TensorProto::DataType::BOOL);
  AddGraphInputTensor(g, "y_in", TensorProto::DataType::FLOAT, shape);

  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("cond_in");
    n->add_output("cond_out");
  }
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("y_in");
    n->add_output("y_out");
  }
  {
    NodeProto *n = g.add_node();
    n->set_op_type("Identity");
    n->add_input("y_in");
    n->add_output("scan_out");
  }

  AddGraphOutputTensor(g, "cond_out", TensorProto::DataType::BOOL);
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("y_out");
    TypeProto::Tensor *tt = vi->ref_type().mutable_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    TensorShapeProto &s = tt->ref_shape();
    for (int64_t d : shape) {
      s.add_dim()->set_dim_value(d);
    }
  }
  {
    ValueInfoProto *vi = g.add_output();
    vi->set_name("scan_out");
    TypeProto::Tensor *tt = vi->ref_type().mutable_tensor_type();
    tt->set_elem_type(static_cast<int>(TensorProto::DataType::FLOAT));
    TensorShapeProto &s = tt->ref_shape();
    for (int64_t d : shape) {
      s.add_dim()->set_dim_value(d);
    }
  }
  return g;
}

// Registers a Loop benchmark case: carries a large FLOAT tensor unchanged
// through a handful of iterations, producing a stacked scan output.
void RegisterLoopBenchmarkCase(std::vector<TestCase> &registry) {
  const OpsetId opset = DefaultOpset(13);
  const KernelContext ctx{opset};
  const Loop loop_kernel{ctx};

  const std::vector<int64_t> shape = {256, 256};
  const int64_t trip = 4;

  NodeProto node;
  node.set_op_type("Loop");
  node.add_input("trip_count");
  node.add_input("cond");
  node.add_input("y");
  node.add_output("res_y");
  node.add_output("res_scan");

  AttributeProto *body_attr = node.add_attribute();
  body_attr->set_name("body");
  body_attr->set_type(AttributeProto::AttributeType::GRAPH);
  *body_attr->add_g() = BuildLoopBenchmarkBody(shape);

  Expect(registry, std::move(node), "test_cc_loop_benchmark", {opset},
         [loop_kernel, shape, trip]() -> IoData {
           Tensor trip_count("trip_count", DataType::INT64, {}, Int64Bytes(trip));
           Tensor cond("cond", DataType::BOOL, {}, std::vector<uint8_t>{1});
           Tensor y = RandnTensor(DataType::FLOAT, shape, 4401);

           std::vector<Tensor> per_iter_scan;
           per_iter_scan.reserve(static_cast<size_t>(trip));
           for (int64_t i = 0; i < trip; ++i) {
             per_iter_scan.push_back(Tensor("", DataType::FLOAT, shape, y.data));
           }
           const Tensor res_y("", DataType::FLOAT, shape, y.data);

           std::vector<Tensor> out = loop_kernel(trip_count, cond, /*v_initial=*/{y},
                                                 /*final_state=*/{res_y}, {per_iter_scan});
           return IoData{{std::move(trip_count), std::move(cond), std::move(y)}, std::move(out)};
         });
}

} // namespace

void RegisterLoopCases(std::vector<TestCase> &registry, TestMode mode) {
  if (mode == TestMode::BENCHMARK) {
    RegisterLoopBenchmarkCase(registry);
    return;
  }
  RegisterTripCountVariants(registry);
  RegisterLoop11Case(registry);
  RegisterMultiCarriedCase(registry);
  RegisterLoop13SeqCase(registry);
  RegisterLoop16SeqNoneCase(registry);
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_backend_test
