/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include <gtest/gtest.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/nested_msg.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/expr.h>

#include <algorithm>
#include <array>
#include <cstring>
#include <functional>
#include <iterator>
#include <new>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

using namespace tvm;
using namespace tvm::relax;

namespace {

TensorType ScalarTensorType(PrimType dtype) {
  auto n = tvm::ffi::make_object<TensorTypeNode>();
  n->dtype = std::move(dtype);
  n->ndim = 0;
  return TensorType(n);
}

}  // namespace

TEST(NestedMsg, Basic) {
  // start with no annotation
  relax::Var x("x", std::nullopt), y("y", std::nullopt);

  // constructor from array, T and nullopt.
  NestedMsg<relax::Expr> msg({x, std::nullopt, x});

  EXPECT_TRUE(msg.IsNested());
  EXPECT_FALSE(msg.IsLeaf());
  EXPECT_TRUE(msg != nullptr);

  EXPECT_ANY_THROW(msg.LeafValue());

  auto arr = msg.NestedArray();
  EXPECT_TRUE(arr[0].LeafValue().same_as(x));
  EXPECT_TRUE(arr[1] == nullptr);
  EXPECT_TRUE(arr[1].IsNull());

  EXPECT_TRUE(arr[2].LeafValue().same_as(x));

  auto a0 = arr[0];
  EXPECT_TRUE(a0.IsLeaf());

  // assignment
  // assign null
  a0 = std::nullopt;
  EXPECT_TRUE(a0 == nullptr);

  // assign array
  a0 = {x, {x, std::nullopt, y}};
  EXPECT_TRUE(a0.IsNested());
  auto t0 = a0.NestedArray()[1];
  EXPECT_TRUE(t0.IsNested());
  EXPECT_TRUE(t0.NestedArray()[2].LeafValue().same_as(y));

  // assign leaf
  a0 = x;

  EXPECT_TRUE(a0.IsLeaf());
  EXPECT_TRUE(a0.LeafValue().same_as(x));
}

TEST(NestedMsg, IntAndAny) {
  NestedMsg<int64_t> msg({1, std::nullopt, 2});
  Any any_msg = msg;
  NestedMsg<int64_t> msg2 = any_msg.cast<NestedMsg<int64_t>>();

  EXPECT_TRUE(msg2.IsNested());
  EXPECT_EQ(msg2.NestedArray()[0].LeafValue(), 1);
  EXPECT_TRUE(msg2.NestedArray()[1].IsNull());
  EXPECT_EQ(msg2.NestedArray()[2].LeafValue(), 2);
}

TEST(NestedMsg, ForEachLeaf) {
  relax::Var x("x", std::nullopt), y("y", std::nullopt);
  NestedMsg<Expr> msg = {x, {x, y}, std::nullopt, {x, {x, y}}};

  int x_count = 0, y_count = 0;

  ForEachLeaf(msg, [&](const Expr& v) {
    if (v.same_as(x)) ++x_count;
    if (v.same_as(y)) ++y_count;
  });
  EXPECT_EQ(x_count, 4);
  EXPECT_EQ(y_count, 2);
}

TEST(NestedMsg, Equal) {
  relax::Var x("x", std::nullopt), y("y", std::nullopt);
  relax::Var z("z", std::nullopt);

  auto fequal = [](Expr lhs, Expr rhs) { return lhs.same_as(rhs); };

  using M = NestedMsg<relax::Expr>;

  EXPECT_TRUE(Equal(M(std::nullopt), M(std::nullopt), fequal));

  EXPECT_TRUE(Equal(M(x), M(x), fequal));

  EXPECT_TRUE(Equal(M({x, y}), M({x, y}), fequal));

  EXPECT_TRUE(Equal(M({x, std::nullopt}), M({x, std::nullopt}), fequal));

  EXPECT_TRUE(Equal(M({x, {std::nullopt, y}}), M({x, {std::nullopt, y}}), fequal));

  EXPECT_TRUE(Equal(M({x, {std::nullopt, y}, {x, z}}), M({x, {std::nullopt, y}, {x, z}}), fequal));

  // type mismatch
  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}, x}), M({x, {std::nullopt, y}, {x, z}}), fequal));

  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}, {x, std::nullopt}}),
                     M({x, {std::nullopt, y}, {x, z}}), fequal));

  EXPECT_FALSE(Equal(M({x, {std::nullopt, y}}), M({x, {std::nullopt, y}, {x, z}}), fequal));

  EXPECT_FALSE(Equal(M(x), M(std::nullopt), fequal));

  EXPECT_FALSE(Equal(M(std::nullopt), M(x), fequal));

  EXPECT_FALSE(Equal(M(x), M(ffi::Array<M>({x})), fequal));

  EXPECT_FALSE(Equal(M(ffi::Array<M>({x})), M(x), fequal));
}

TEST(NestedMsg, MapAndDecompose) {
  relax::Var x("x", PrimType::Int(16));
  relax::Var y("y", PrimType::Int(32));
  relax::Var z("z", PrimType::Int(64));

  BlockBuilder bb = BlockBuilder::Create(std::nullopt);
  relax::Expr t0 = bb->Normalize(Tuple({x, y}));
  relax::Expr t1 = bb->Normalize(Tuple({t0, x, z, t0}));

  auto c0 = IntImm::Int32(0);
  auto c1 = IntImm::Int32(1);
  auto c2 = IntImm::Int32(2);

  auto output = MapToNestedMsg<IntImm>(t1, [&](Expr value) {
    if (value.same_as(x)) return c0;
    if (value.same_as(y)) return c1;
    return c2;
  });

  NestedMsg<IntImm> expected = {{c0, c1}, c0, c2, {c0, c1}};

  EXPECT_TRUE(Equal(output, expected,
                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));

  auto output2 = MapToNestedMsg<IntImm>(GetType(t1), [&](Type ty) -> NestedMsg<IntImm> {
    const auto* prim_ty = ty.as<PrimTypeNode>();
    if (prim_ty == nullptr) return std::nullopt;
    int bits = prim_ty->dtype.bits;
    if (bits == 16) return c0;
    if (bits == 32) return c1;
    if (bits == 64) return c2;
    return std::nullopt;
  });

  EXPECT_TRUE(Equal(output2, expected,
                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));

  int x_count = 0, y_count = 0, z_count = 0;

  DecomposeNestedMsg(t1, expected, [&](Expr value, NestedMsg<IntImm> msg) {
    if (value.same_as(x)) {
      EXPECT_TRUE(msg.LeafValue().same_as(c0));
      ++x_count;
    } else if (value.same_as(y)) {
      EXPECT_TRUE(msg.LeafValue().same_as(c1));
      ++y_count;
    } else {
      EXPECT_TRUE(msg.LeafValue().same_as(c2));
      ++z_count;
    }
  });
  EXPECT_EQ(x_count, 3);
  EXPECT_EQ(y_count, 2);
  EXPECT_EQ(z_count, 1);
}

TEST(NestedMsg, MapToNestedMsgByType) {
  auto sf0 = ScalarTensorType(PrimType::Float(32));
  auto sf1 = TupleType({sf0, sf0});
  auto sf2 = TupleType({sf0, sf0});
  auto x = relax::Var("x", TupleType({sf1, sf2, sf0}));

  auto msg = MapToNestedMsgByType<Expr>(x, [](Expr value) { return value; });

  EXPECT_TRUE(msg.IsNested());
  auto arr = msg.NestedArray();

  EXPECT_TRUE(arr[1].IsNested());
  auto arr1 = arr[1].NestedArray();

  EXPECT_TRUE(arr1[0].IsLeaf());
  EXPECT_TRUE(
      tvm::ffi::StructuralEqual()(arr1[0].LeafValue(), TupleGetItem(TupleGetItem(x, 1), 0)));

  EXPECT_TRUE(arr[2].IsLeaf());
  EXPECT_TRUE(tvm::ffi::StructuralEqual()(arr[2].LeafValue(), TupleGetItem(x, 2)));
}

TEST(NestedMsg, NestedMsgToExpr) {
  auto sf0 = ScalarTensorType(PrimType::Float(32));
  auto sf1 = TupleType({sf0, sf0});

  auto c0 = IntImm::Int32(0);
  auto c1 = IntImm::Int32(1);
  auto c2 = IntImm::Int32(2);

  relax::Var x("x", sf0), y("y", sf0), z("z", sf0);

  NestedMsg<IntImm> msg = {c0, {c0, c1}, {c0, {c1, c2}}};
  auto expr = NestedMsgToExpr<IntImm>(msg, [&](ffi::Optional<IntImm> leaf) {
    TVM_FFI_ICHECK(leaf.has_value());
    int value = leaf.value()->value;
    switch (value) {
      case 0:
        return x;
      case 1:
        return y;
      default:
        return z;
    }
  });

  Expr expected = Tuple({x, Tuple({x, y}), Tuple({x, Tuple({y, z})})});
  EXPECT_TRUE(tvm::ffi::StructuralEqual()(expr, expected));

  // test simplified
  relax::Var t("t", sf1);
  NestedMsg<Expr> msg1 = {TupleGetItem(t, 0), TupleGetItem(t, 1)};
  auto expr1 = NestedMsgToExpr<Expr>(msg1, [](ffi::Optional<Expr> leaf) { return leaf.value(); });
  EXPECT_TRUE(tvm::ffi::StructuralEqual()(expr1, t));
}

TEST(NestedMsg, CombineNestedMsg) {
  auto c0 = IntImm::Int32(0);
  auto c1 = IntImm::Int32(1);
  auto c2 = IntImm::Int32(2);

  NestedMsg<IntImm> lhs = {c0, {c0, c1}, std::nullopt, {c0, {c1, c2}}};
  NestedMsg<IntImm> rhs = {c1, {c2, std::nullopt}, std::nullopt, {c1, {c2, c2}}};
  NestedMsg<IntImm> expected = {c1, {c2, c1}, std::nullopt, {c1, {c2, c2}}};

  auto output = CombineNestedMsg(lhs, rhs, [](IntImm x, IntImm y) {
    if (x->value > y->value) return x;
    return y;
  });

  EXPECT_TRUE(Equal(output, expected,
                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
}

TEST(NestedMsg, MapNestedMsg) {
  auto c0 = IntImm::Int32(0);
  auto c1 = IntImm::Int32(1);
  auto c2 = IntImm::Int32(2);
  auto c3 = IntImm::Int32(3);

  NestedMsg<IntImm> msg = {c0, {c0, c1}, std::nullopt, {c0, {c2, c1}}};
  NestedMsg<IntImm> expected = {c3, {c3, std::nullopt}, std::nullopt, {c3, {c2, std::nullopt}}};

  auto output = MapNestedMsg(msg, [](IntImm x) {
    if (x->value == 0) {
      return NestedMsg<IntImm>(IntImm::Int32(3));
    } else if (x->value == 1) {
      return NestedMsg<IntImm>();
    } else {
      return NestedMsg<IntImm>(x);
    }
  });

  EXPECT_TRUE(Equal(output, expected,
                    [](IntImm lhs, IntImm rhs) -> bool { return lhs->value == rhs->value; }));
}

TEST(NestedMsg, TransformTupleLeaf) {
  auto c0 = IntImm::Int32(0);
  auto c1 = IntImm::Int32(1);
  auto c2 = IntImm::Int32(2);
  using NInt = NestedMsg<IntImm>;

  NInt msg1 = {c0, {c0, c1}, c2, {c0, {c1, c2}}};
  NInt msg2 = {c1, {c2, c0}, c2, {c1, {c2, c0}}};

  PrimType s = PrimType::Int(32);
  relax::Var x("x", s), y("y", s), z("z", s);
  BlockBuilder bb = BlockBuilder::Create(std::nullopt);
  Expr expr = bb->Normalize(Tuple({x, Tuple({x, x}), x, Tuple({x, Tuple({x, x})})}));

  auto ftransleaf = [&](Expr value, std::array<NInt, 2> msgs) -> Expr {
    int lhs = msgs[0].LeafValue().as_or_throw<IntImm>()->value;
    int rhs = msgs[1].LeafValue().as_or_throw<IntImm>()->value;
    if (lhs > rhs)
      return z;
    else if (lhs == rhs)
      return value;
    else
      return y;
  };

  Expr expected = Tuple({y, Tuple({y, z}), x, Tuple({y, Tuple({y, z})})});

  EXPECT_TRUE(tvm::ffi::StructuralEqual()(
      TransformTupleLeaf(expr, std::array<NInt, 2>({msg1, msg2}), ftransleaf), expected));

  EXPECT_TRUE(
      expr.same_as(TransformTupleLeaf(expr, std::array<NInt, 2>({msg1, msg1}), ftransleaf)));
}