#include <gtest/gtest.h>
#include "caffe2/core/common.h"
#include "caffe2/core/logging.h"
#include "caffe2/opt/bound_shape_inferencer.h"
#include "caffe2/utils/proto_utils.h"

using namespace caffe2;
namespace {

ShapeInfo makeTensorInfo(
    const std::vector<TensorBoundShape::DimType>& t,
    const std::vector<int64_t>& dims,
    TensorProto::DataType dtype = TensorProto_DataType_FLOAT) {
  ShapeInfo info;
  info.setDimType(t);
  TensorShape& shape = info.shape;
  for (const auto d : dims) {
    shape.add_dims(d);
  }
  shape.set_data_type(dtype);
  return info;
}

void verifyShapeInfo(
    const ShapeInfoMap& info,
    const std::string& name,
    const std::vector<TensorBoundShape::DimType>& t,
    const std::vector<int64_t>& dims,
    TensorProto::DataType dtype = TensorProto_DataType_FLOAT) {
  LOG(INFO) << "Checking " << name;
  const auto it = info.find(name);
  ASSERT_TRUE(it != info.end());
  const auto& shape_info = it->second;
  EXPECT_EQ(shape_info.getDimType(), t);
  const auto& shape = shape_info.shape;
  ASSERT_EQ(shape.dims_size(), dims.size());
  for (int i = 0; i < dims.size(); ++i) {
    EXPECT_EQ(shape.dims(i), dims[i]);
  }
  EXPECT_EQ(shape.data_type(), dtype);
}

} // namespace

TEST(BoundShapeInference, SparseLengthsSum) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "SparseLengthsSum", "", {"Weights", "Data", "Lengths"}, {"Out"}, {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "Weights",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {1000, 16}));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "Weights",
      {TensorBoundShape_DimType_CONSTANT, TensorBoundShape_DimType_CONSTANT},
      {1000, 16});
  verifyShapeInfo(
      out_shape,
      "Data",
      {TensorBoundShape_DimType_FEATURE_MAX_DEFAULT},
      {spec.max_seq_size},
      TensorProto_DataType_INT64);
  verifyShapeInfo(
      out_shape,
      "Lengths",
      {TensorBoundShape_DimType_BATCH},
      {spec.max_batch_size},
      TensorProto_DataType_INT32);
  verifyShapeInfo(
      out_shape,
      "Out",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
}

TEST(BoundShapeInference, SparseLengthsSumFused8BitRowwise) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "SparseLengthsSumFused8BitRowwise",
      "",
      {"Weights", "Data", "Lengths"},
      {"Out"},
      {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "Weights",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {1000, 58},
          TensorProto_DataType_INT8));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "Weights",
      {TensorBoundShape_DimType_CONSTANT, TensorBoundShape_DimType_CONSTANT},
      {1000, 58},
      TensorProto_DataType_INT8);
  verifyShapeInfo(
      out_shape,
      "Data",
      {TensorBoundShape_DimType_FEATURE_MAX_DEFAULT},
      {spec.max_seq_size},
      TensorProto_DataType_INT64);
  verifyShapeInfo(
      out_shape,
      "Lengths",
      {TensorBoundShape_DimType_BATCH},
      {spec.max_batch_size},
      TensorProto_DataType_INT32);
  verifyShapeInfo(
      out_shape,
      "Out",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 50});
}

TEST(BoundShapeInference, SparseLengthsSumFused4BitRowwise) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "SparseLengthsSumFused4BitRowwise",
      "",
      {"Weights", "Data", "Lengths"},
      {"Out"},
      {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "Weights",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {1000, 54},
          TensorProto_DataType_INT8));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "Weights",
      {TensorBoundShape_DimType_CONSTANT, TensorBoundShape_DimType_CONSTANT},
      {1000, 54},
      TensorProto_DataType_INT8);
  verifyShapeInfo(
      out_shape,
      "Data",
      {TensorBoundShape_DimType_FEATURE_MAX_DEFAULT},
      {spec.max_seq_size},
      TensorProto_DataType_INT64);
  verifyShapeInfo(
      out_shape,
      "Lengths",
      {TensorBoundShape_DimType_BATCH},
      {spec.max_batch_size},
      TensorProto_DataType_INT32);
  verifyShapeInfo(
      out_shape,
      "Out",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 100});
}

TEST(BoundShapeInference, LengthsRangeFill) {
  NetDef net;
  net.add_op()->CopyFrom(
      CreateOperatorDef("LengthsRangeFill", "", {"X"}, {"Y"}, {}));
  net.add_op()->CopyFrom(CreateOperatorDef("Copy", "", {"Y"}, {"Z"}, {}));
  ShapeInfoMap shape_map;
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "X",
      {TensorBoundShape_DimType_BATCH},
      {spec.max_batch_size},
      TensorProto_DataType_INT32);
  verifyShapeInfo(
      out_shape,
      "Y",
      {TensorBoundShape_DimType_FEATURE_MAX_DEFAULT},
      {spec.max_seq_size},
      TensorProto_DataType_INT32);
  verifyShapeInfo(
      out_shape,
      "Z",
      {TensorBoundShape_DimType_FEATURE_MAX_DEFAULT},
      {spec.max_seq_size},
      TensorProto_DataType_INT32);
}

TEST(BoundShapeInference, Reshape) {
  NetDef net;
  std::vector<int> new_shape{-1, 8};
  std::vector<int> new_shape2{2, 8};
  net.add_op()->CopyFrom(
      CreateOperatorDef("FC", "", {"X0", "W0", "B0"}, {"X1"}, {}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Reshape",
      "",
      {"X1"},
      {"Y1", "old_shape"},
      {MakeArgument<std::vector<int>>("shape", new_shape)}));

  // Cannot infer shape for this one because input/output shape doesn't match
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Reshape",
      "",
      {"X1"},
      {"Y2", "old_shape2"},
      {MakeArgument<std::vector<int>>("shape", new_shape2)}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "W0",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {16, 1024}));
  shape_map.emplace(
      "B0", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {16}));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "X0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 1024});
  verifyShapeInfo(
      out_shape,
      "X1",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
  verifyShapeInfo(
      out_shape,
      "Y1",
      {TensorBoundShape_DimType_BATCH,
       TensorBoundShape_DimType_CONSTANT}, // TODO
      {spec.max_batch_size * 16 / 8, 8});
  EXPECT_TRUE(out_shape.find("Y2") == out_shape.end());
}

TEST(BoundShapeInference, ConcatMissingInput) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Concat",
      "",
      {"I0", "I1"},
      {"Cout", "split_info"},
      {MakeArgument<int>("axis", 1), MakeArgument<int>("add_axis", 1)}));
  BoundShapeSpec spec(20, 1000);
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "I0",
      makeTensorInfo(
          {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
          {spec.max_batch_size, 60}));
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "I0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 60});
  verifyShapeInfo(
      out_shape,
      "Cout",
      {TensorBoundShape_DimType_BATCH,
       TensorBoundShape_DimType_CONSTANT,
       TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 2, 60});
}

TEST(BoundShapeInference, ConcatInferInputBackwards) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Concat",
      "",
      {"I0", "I1"},
      {"Cout", "split_info"},
      {MakeArgument<int>("axis", 1)}));
  net.add_op()->CopyFrom(
      CreateOperatorDef("FCTransposed", "", {"Cout", "W0", "B0"}, {"Y"}, {}));
  BoundShapeSpec spec(20, 1000);
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "I0",
      makeTensorInfo(
          {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
          {spec.max_batch_size, 60}));
  shape_map.emplace(
      "W0",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {101, 16}));
  shape_map.emplace(
      "B0", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {16}));
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "I0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 60});
  verifyShapeInfo(
      out_shape,
      "Cout",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 101});
  verifyShapeInfo(
      out_shape,
      "Y",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
  verifyShapeInfo(
      out_shape,
      "I1",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 101 - 60});
}

TEST(BoundShapeInference, Bucketize) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Bucketize",
      "",
      {"In"},
      {"Out"},
      {MakeArgument<std::vector<float>>("boundaries", {1.0, 2.0})}));
  BoundShapeSpec spec(20, 1000);
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "In",
      makeTensorInfo(
          {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
          {spec.max_batch_size, 60}));
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "Out",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 60},
      TensorProto_DataType_INT32);
}

TEST(BoundShapeInference, Split) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Split", "", {"X"}, {"Y0", "Y1"}, {MakeArgument<int>("axis", 1)}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Split",
      "",
      {"X"},
      {"Y2", "Y3", "Y4"},
      {MakeArgument<int>("axis", 1),
       MakeArgument<std::vector<int>>("split", {4, 30, 14})}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Split",
      "",
      {"X1"},
      {"Y5", "Y6"},
      {MakeArgument<int>("axis", 1), MakeArgument<int>("add_axis", 1)}));
  BoundShapeSpec spec(20, 1000);
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "X",
      makeTensorInfo(
          {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
          {spec.max_batch_size, 48}));
  shape_map.emplace(
      "X1",
      makeTensorInfo(
          {TensorBoundShape_DimType_BATCH,
           TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {spec.max_batch_size, 2, 48}));
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "X",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 48});
  verifyShapeInfo(
      out_shape,
      "X1",
      {TensorBoundShape_DimType_BATCH,
       TensorBoundShape_DimType_CONSTANT,
       TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 2, 48});
  verifyShapeInfo(
      out_shape,
      "Y0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 48 / 2});
  verifyShapeInfo(
      out_shape,
      "Y1",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 48 / 2});
  verifyShapeInfo(
      out_shape,
      "Y2",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 4});
  verifyShapeInfo(
      out_shape,
      "Y3",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 30});
  verifyShapeInfo(
      out_shape,
      "Y4",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 14});
  verifyShapeInfo(
      out_shape,
      "Y5",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 48});
  verifyShapeInfo(
      out_shape,
      "Y6",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 48});
}

TEST(BoundShapeInference, FC) {
  NetDef net;
  net.add_op()->CopyFrom(
      CreateOperatorDef("FC", "", {"X0", "W0", "B0"}, {"Out0"}, {}));
  net.add_op()->CopyFrom(
      CreateOperatorDef("FCTransposed", "", {"X1", "W1", "B1"}, {"Out1"}, {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "W0",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {16, 1024}));
  shape_map.emplace(
      "B0", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {16}));
  shape_map.emplace(
      "W1",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {16, 1024}));
  shape_map.emplace(
      "B1", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {1024}));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "X0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 1024});
  verifyShapeInfo(
      out_shape,
      "Out0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
  verifyShapeInfo(
      out_shape,
      "X1",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
  verifyShapeInfo(
      out_shape,
      "Out1",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 1024});
}

TEST(BoundShapeInference, FC3D) {
  NetDef net;
  net.add_op()->CopyFrom(
      CreateOperatorDef("FC", "", {"X0", "W0", "B0"}, {"Out0"}, {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "W0",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {16, 1, 1024}));
  shape_map.emplace(
      "B0", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {16}));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  verifyShapeInfo(
      out_shape,
      "X0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 1024});
  verifyShapeInfo(
      out_shape,
      "Out0",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 16});
}

TEST(BoundShapeInference, Combo0) {
  NetDef net;
  net.add_op()->CopyFrom(CreateOperatorDef(
      "SparseLengthsSum", "", {"Weights0", "Data0", "Lengths0"}, {"EB0"}, {}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "SparseLengthsSum", "", {"Weights1", "Data1", "Lengths1"}, {"EB1"}, {}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "Concat",
      "",
      {"EB0", "EB1"},
      {"Cout", "split_info"},
      {MakeArgument<int>("axis", 1), MakeArgument<int>("add_axis", 1)}));
  net.add_op()->CopyFrom(CreateOperatorDef(
      "BatchMatMul",
      "",
      {"Cout", "Cout"},
      {"Bout"},
      {MakeArgument<int>("trans_b", 1)}));
  net.add_op()->CopyFrom(
      CreateOperatorDef("Flatten", "", {"Bout"}, {"Fout"}, {}));
  net.add_op()->CopyFrom(
      CreateOperatorDef("BatchGather", "", {"Fout", "Indices"}, {"Gout"}, {}));
  ShapeInfoMap shape_map;
  shape_map.emplace(
      "Weights0",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {1000, 16}));
  shape_map.emplace(
      "Weights1",
      makeTensorInfo(
          {TensorBoundShape_DimType_CONSTANT,
           TensorBoundShape_DimType_CONSTANT},
          {20000, 16}));
  shape_map.emplace(
      "Indices", makeTensorInfo({TensorBoundShape_DimType_CONSTANT}, {2}));
  BoundShapeSpec spec(20, 1000);
  BoundShapeInferencer eng(spec);
  eng.InferBoundShapeAndType(net, shape_map, nullptr);
  const auto& out_shape = eng.shape_info();
  LOG(INFO) << eng.PrintShapeInfo();
  verifyShapeInfo(
      out_shape,
      "Gout",
      {TensorBoundShape_DimType_BATCH, TensorBoundShape_DimType_CONSTANT},
      {spec.max_batch_size, 2});
}
