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

#include "onnx_extensions/shapes/shapes/training/shape_training.h"

#include "onnx_core/shapes/shape_inference.h"
#include "onnx_core/shapes/shapes_context.h"
#include "onnx_core/symbolic/sym_tensor.h"
#include "onnx_proto/onnx.h"

#include <gtest/gtest.h>

#include <stdexcept>
#include <string>

using namespace ONNX_LIGHT_NAMESPACE;

namespace Test {

namespace {

// Builds an Adam node with ``n`` optimised tensors. Each optimised tensor
// ``X_i`` gets matching ``G_i``, ``V_i`` and ``H_i`` inputs, and the node
// declares the corresponding ``3 * n`` outputs.
NodeProto MakeAdamNode(int n) {
  NodeProto node;
  node.set_op_type("Adam");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < n; ++i) {
    node.add_input("X" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("G" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("V" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("H" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("X" + std::to_string(i + 1) + "_new");
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("V" + std::to_string(i + 1) + "_new");
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("H" + std::to_string(i + 1) + "_new");
  }
  return node;
}

void SeedScalar(core::shapes::ShapesContext &ctx, const std::string &name,
                core::symbolic::TensorType dtype) {
  ctx.Set(name, core::symbolic::SymTensor(nullptr, dtype, core::symbolic::SymShape{}));
}

void SeedTensor(core::shapes::ShapesContext &ctx, const std::string &name,
                core::symbolic::TensorType dtype, core::symbolic::SymShape shape) {
  ctx.Set(name, core::symbolic::SymTensor(nullptr, dtype, std::move(shape)));
}

} // namespace

TEST(OnnxOptimShapeAdam, PropagatesShapesForSingleOptimizedTensor) {
  NodeProto node = MakeAdamNode(1);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(2), core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape);

  ctx.ComputeShapeNode(node);

  ASSERT_TRUE(ctx.Has("X1_new"));
  ASSERT_TRUE(ctx.Has("V1_new"));
  ASSERT_TRUE(ctx.Has("H1_new"));
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape);
  EXPECT_EQ(ctx.Get("V1_new").Shape(), shape);
  EXPECT_EQ(ctx.Get("H1_new").Shape(), shape);
}

TEST(OnnxOptimShapeAdam, PropagatesShapesForMultipleOptimizedTensors) {
  NodeProto node = MakeAdamNode(2);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape_x1{core::symbolic::SymDim(4)};
  const core::symbolic::SymShape shape_x2{core::symbolic::SymDim(2), core::symbolic::SymDim(5)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "X2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "G2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "V2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "H2", core::symbolic::TensorType::kDouble, shape_x2);

  ctx.ComputeShapeNode(node);

  // Outputs are laid out as [X1_new, X2_new, V1_new, V2_new, H1_new, H2_new]
  // and each output mirrors the dtype and shape of its corresponding input.
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("X2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("X2_new").Shape(), shape_x2);
  EXPECT_EQ(ctx.Get("V1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("V1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("V2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("V2_new").Shape(), shape_x2);
  EXPECT_EQ(ctx.Get("H1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("H1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("H2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("H2_new").Shape(), shape_x2);
}

TEST(OnnxOptimShapeAdam, RejectsInputCountNotMultipleOfFour) {
  // 2 (R, T) + 5 variadic inputs = not a multiple of 4 for the variadic part.
  NodeProto node;
  node.set_op_type("Adam");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < 5; ++i) {
    node.add_input("V" + std::to_string(i));
  }
  node.add_output("Y0");
  node.add_output("Y1");
  node.add_output("Y2");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  for (int i = 0; i < 5; ++i) {
    SeedTensor(ctx, "V" + std::to_string(i), core::symbolic::TensorType::kFloat,
               core::symbolic::SymShape{core::symbolic::SymDim(1)});
  }

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeAdam, RejectsWrongOutputCount) {
  // Build a 1-optimised-tensor Adam node but declare only 2 outputs instead
  // of the expected 3 (X_new, V_new, H_new).
  NodeProto node;
  node.set_op_type("Adam");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  node.add_input("X1");
  node.add_input("G1");
  node.add_input("V1");
  node.add_input("H1");
  node.add_output("X1_new");
  node.add_output("V1_new");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape);

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeAdam, DirectCallRejectsWrongOpType) {
  NodeProto node;
  node.set_op_type("NotAdam");
  node.add_output("Y");

  core::shapes::ShapesContext ctx;
  EXPECT_THROW(onnx_shapes::shapes::training::ComputeShapeAdam(ctx, node), std::invalid_argument);
}

namespace {

// Builds an Adagrad node with ``n`` optimised tensors. Each optimised tensor
// ``X_i`` gets matching ``G_i`` and ``H_i`` inputs, and the node declares the
// corresponding ``2 * n`` outputs.
NodeProto MakeAdagradNode(int n) {
  NodeProto node;
  node.set_op_type("Adagrad");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < n; ++i) {
    node.add_input("X" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("G" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("H" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("X" + std::to_string(i + 1) + "_new");
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("H" + std::to_string(i + 1) + "_new");
  }
  return node;
}

// Builds a Momentum node with ``n`` optimised tensors. Each optimised tensor
// ``X_i`` gets matching ``G_i`` and ``V_i`` inputs, and the node declares the
// corresponding ``2 * n`` outputs.
NodeProto MakeMomentumNode(int n) {
  NodeProto node;
  node.set_op_type("Momentum");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < n; ++i) {
    node.add_input("X" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("G" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_input("V" + std::to_string(i + 1));
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("X" + std::to_string(i + 1) + "_new");
  }
  for (int i = 0; i < n; ++i) {
    node.add_output("V" + std::to_string(i + 1) + "_new");
  }
  return node;
}

} // namespace

TEST(OnnxOptimShapeAdagrad, PropagatesShapesForSingleOptimizedTensor) {
  NodeProto node = MakeAdagradNode(1);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(2), core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape);

  ctx.ComputeShapeNode(node);

  ASSERT_TRUE(ctx.Has("X1_new"));
  ASSERT_TRUE(ctx.Has("H1_new"));
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape);
  EXPECT_EQ(ctx.Get("H1_new").Shape(), shape);
}

TEST(OnnxOptimShapeAdagrad, PropagatesShapesForMultipleOptimizedTensors) {
  NodeProto node = MakeAdagradNode(2);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape_x1{core::symbolic::SymDim(4)};
  const core::symbolic::SymShape shape_x2{core::symbolic::SymDim(2), core::symbolic::SymDim(5)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "X2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "G2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "H2", core::symbolic::TensorType::kDouble, shape_x2);

  ctx.ComputeShapeNode(node);

  // Outputs are laid out as [X1_new, X2_new, H1_new, H2_new] and each output
  // mirrors the dtype and shape of its corresponding input.
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("X2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("X2_new").Shape(), shape_x2);
  EXPECT_EQ(ctx.Get("H1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("H1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("H2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("H2_new").Shape(), shape_x2);
}

TEST(OnnxOptimShapeAdagrad, RejectsInputCountNotMultipleOfThree) {
  // 2 (R, T) + 4 variadic inputs = not a multiple of 3 for the variadic part.
  NodeProto node;
  node.set_op_type("Adagrad");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < 4; ++i) {
    node.add_input("V" + std::to_string(i));
  }
  node.add_output("Y0");
  node.add_output("Y1");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  for (int i = 0; i < 4; ++i) {
    SeedTensor(ctx, "V" + std::to_string(i), core::symbolic::TensorType::kFloat,
               core::symbolic::SymShape{core::symbolic::SymDim(1)});
  }

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeAdagrad, RejectsWrongOutputCount) {
  // Build a 1-optimised-tensor Adagrad node but declare only 1 output instead
  // of the expected 2 (X_new, H_new).
  NodeProto node;
  node.set_op_type("Adagrad");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  node.add_input("X1");
  node.add_input("G1");
  node.add_input("H1");
  node.add_output("X1_new");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "H1", core::symbolic::TensorType::kFloat, shape);

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeAdagrad, DirectCallRejectsWrongOpType) {
  NodeProto node;
  node.set_op_type("NotAdagrad");
  node.add_output("Y");

  core::shapes::ShapesContext ctx;
  EXPECT_THROW(onnx_shapes::shapes::training::ComputeShapeAdagrad(ctx, node),
               std::invalid_argument);
}

TEST(OnnxOptimShapeMomentum, PropagatesShapesForSingleOptimizedTensor) {
  NodeProto node = MakeMomentumNode(1);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(2), core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape);

  ctx.ComputeShapeNode(node);

  ASSERT_TRUE(ctx.Has("X1_new"));
  ASSERT_TRUE(ctx.Has("V1_new"));
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape);
  EXPECT_EQ(ctx.Get("V1_new").Shape(), shape);
}

TEST(OnnxOptimShapeMomentum, PropagatesShapesForMultipleOptimizedTensors) {
  NodeProto node = MakeMomentumNode(2);

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape_x1{core::symbolic::SymDim(4)};
  const core::symbolic::SymShape shape_x2{core::symbolic::SymDim(2), core::symbolic::SymDim(5)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "X2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "G2", core::symbolic::TensorType::kDouble, shape_x2);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape_x1);
  SeedTensor(ctx, "V2", core::symbolic::TensorType::kDouble, shape_x2);

  ctx.ComputeShapeNode(node);

  // Outputs are laid out as [X1_new, X2_new, V1_new, V2_new] and each output
  // mirrors the dtype and shape of its corresponding input.
  EXPECT_EQ(ctx.Get("X1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("X1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("X2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("X2_new").Shape(), shape_x2);
  EXPECT_EQ(ctx.Get("V1_new").Dtype(), core::symbolic::TensorType::kFloat);
  EXPECT_EQ(ctx.Get("V1_new").Shape(), shape_x1);
  EXPECT_EQ(ctx.Get("V2_new").Dtype(), core::symbolic::TensorType::kDouble);
  EXPECT_EQ(ctx.Get("V2_new").Shape(), shape_x2);
}

TEST(OnnxOptimShapeMomentum, RejectsInputCountNotMultipleOfThree) {
  // 2 (R, T) + 4 variadic inputs = not a multiple of 3 for the variadic part.
  NodeProto node;
  node.set_op_type("Momentum");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  for (int i = 0; i < 4; ++i) {
    node.add_input("V" + std::to_string(i));
  }
  node.add_output("Y0");
  node.add_output("Y1");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  for (int i = 0; i < 4; ++i) {
    SeedTensor(ctx, "V" + std::to_string(i), core::symbolic::TensorType::kFloat,
               core::symbolic::SymShape{core::symbolic::SymDim(1)});
  }

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeMomentum, RejectsWrongOutputCount) {
  // Build a 1-optimised-tensor Momentum node but declare only 1 output instead
  // of the expected 2 (X_new, V_new).
  NodeProto node;
  node.set_op_type("Momentum");
  node.set_domain("ai.onnx.preview.training");
  node.add_input("R");
  node.add_input("T");
  node.add_input("X1");
  node.add_input("G1");
  node.add_input("V1");
  node.add_output("X1_new");

  core::shapes::ShapesContext ctx;
  SeedScalar(ctx, "R", core::symbolic::TensorType::kFloat);
  SeedScalar(ctx, "T", core::symbolic::TensorType::kInt64);
  const core::symbolic::SymShape shape{core::symbolic::SymDim(3)};
  SeedTensor(ctx, "X1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "G1", core::symbolic::TensorType::kFloat, shape);
  SeedTensor(ctx, "V1", core::symbolic::TensorType::kFloat, shape);

  EXPECT_THROW(ctx.ComputeShapeNode(node), std::invalid_argument);
}

TEST(OnnxOptimShapeMomentum, DirectCallRejectsWrongOpType) {
  NodeProto node;
  node.set_op_type("NotMomentum");
  node.add_output("Y");

  core::shapes::ShapesContext ctx;
  EXPECT_THROW(onnx_shapes::shapes::training::ComputeShapeMomentum(ctx, node),
               std::invalid_argument);
}

} // namespace Test
