#include "tensorflow/core/public/tensor.h"

#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include <gtest/gtest.h>

namespace tensorflow {

TEST(TensorTest, Default) {
  Tensor t;
  EXPECT_EQ(t.dtype(), DT_FLOAT);
  EXPECT_EQ(t.dims(), 1);
  EXPECT_EQ(t.NumElements(), 0);
}

TEST(TensorTest, DataType_Traits) {
  EXPECT_TRUE(std::is_trivial<float>::value);
  EXPECT_TRUE(std::is_trivial<double>::value);
  EXPECT_TRUE(std::is_trivial<int32>::value);
  EXPECT_TRUE(std::is_trivial<uint8>::value);
  EXPECT_TRUE(std::is_trivial<int16>::value);
  EXPECT_TRUE(std::is_trivial<int8>::value);
  EXPECT_TRUE(std::is_trivial<int64>::value);
  EXPECT_TRUE(std::is_trivial<bool>::value);
  EXPECT_FALSE(std::is_trivial<string>::value);

  EXPECT_EQ(sizeof(bool), 1);

  // Unfortunately. std::complex::complex() initializes (0, 0).
  EXPECT_FALSE(std::is_trivial<complex64>::value);
  EXPECT_FALSE(std::is_trivial<std::complex<double>>::value);
  EXPECT_TRUE(std::is_trivial<float[2]>::value);
  struct MyComplex {
    float re, im;
  };
  EXPECT_TRUE(std::is_trivial<MyComplex>::value);
}

template <typename T>
void TestCopies(const Tensor& t) {
  {
    LOG(INFO) << "CopyFrom()";
    Tensor t2(t.dtype());
    EXPECT_TRUE(t2.CopyFrom(t, t.shape()));
    test::ExpectTensorEqual<T>(t, t2);
  }
  {
    LOG(INFO) << "operator=()";
    Tensor t2(t.dtype());
    t2 = t;
    test::ExpectTensorEqual<T>(t, t2);
  }
  {
    LOG(INFO) << "deep copy";
    Tensor t2(t.dtype(), t.shape());
    t2.flat<T>() = t.flat<T>();
    test::ExpectTensorEqual<T>(t, t2);
  }
  {
    LOG(INFO) << "AsProtoField()";
    TensorProto proto;
    t.AsProtoField(&proto);
    Tensor t2(t.dtype());
    EXPECT_TRUE(t2.FromProto(proto));
    test::ExpectTensorEqual<T>(t, t2);
  }
  {
    LOG(INFO) << "AsProtoTensorContent()";
    TensorProto proto;
    t.AsProtoTensorContent(&proto);
    Tensor t2(t.dtype());
    EXPECT_TRUE(t2.FromProto(proto));
    test::ExpectTensorEqual<T>(t, t2);
    // Make another copy via tensor_content field.
    *proto.mutable_tensor_content() = proto.tensor_content();
    Tensor t3(t.dtype());
    EXPECT_TRUE(t3.FromProto(proto));
    test::ExpectTensorEqual<T>(t, t2);
  }
  {
    LOG(INFO) << "AsTensor";
    gtl::ArraySlice<T> values(t.flat<T>().data(), t.NumElements());
    Tensor t2 = test::AsTensor(values, t.shape());
    test::ExpectTensorEqual<T>(t, t2);
  }
}

TEST(Tensor_Float, Simple) {
  Tensor t(DT_FLOAT, TensorShape({10, 20}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({10, 20})));
  for (int64 a = 0; a < t.shape().dim_size(0); a++) {
    for (int64 b = 0; b < t.shape().dim_size(1); b++) {
      t.matrix<float>()(a, b) = static_cast<float>(a * b);
    }
  }
  TestCopies<float>(t);
}

TEST(Tensor_QInt8, Simple) {
  Tensor t(DT_QINT8, TensorShape({2, 2}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({2, 2})));
  for (int64 a = 0; a < t.shape().dim_size(0); a++) {
    for (int64 b = 0; b < t.shape().dim_size(1); b++) {
      t.matrix<qint8>()(a, b) = qint8(a * b);
    }
  }
  TestCopies<qint8>(t);
}

TEST(Tensor_QUInt8, Simple) {
  Tensor t(DT_QUINT8, TensorShape({2, 2}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({2, 2})));
  for (int64 a = 0; a < t.shape().dim_size(0); a++) {
    for (int64 b = 0; b < t.shape().dim_size(1); b++) {
      t.matrix<Eigen::QUInt8>()(a, b) = Eigen::QUInt8(a * b);
    }
  }
  TestCopies<Eigen::QUInt8>(t);
}

TEST(Tensor_QInt32, Simple) {
  Tensor t(DT_QINT32, TensorShape({2, 2}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({2, 2})));
  for (int64 a = 0; a < t.shape().dim_size(0); a++) {
    for (int64 b = 0; b < t.shape().dim_size(1); b++) {
      t.matrix<qint32>()(a, b) = qint32(static_cast<int32>(a * b));
    }
  }
  TestCopies<qint32>(t);
}

TEST(Tensor_Float, Reshape) {
  Tensor t(DT_FLOAT, TensorShape({2, 3, 4, 5}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({2, 3, 4, 5})));

  {
    auto tensor = t.tensor<float, 4>();
    EXPECT_EQ(2, tensor.dimension(0));
    EXPECT_EQ(3, tensor.dimension(1));
    EXPECT_EQ(4, tensor.dimension(2));
    EXPECT_EQ(5, tensor.dimension(3));

    // Set first and last elements.
    tensor(0, 0, 0, 0) = 0.01f;
    tensor(1, 2, 3, 4) = 0.02f;
  }
  {
    auto shaped = t.shaped<float, 1>({120});
    EXPECT_EQ(120, shaped.dimension(0));
    EXPECT_EQ(shaped(0), 0.01f);
    EXPECT_EQ(shaped(119), 0.02f);
  }
  {
    auto shaped = t.shaped<float, 2>({6, 20});
    EXPECT_EQ(6, shaped.dimension(0));
    EXPECT_EQ(20, shaped.dimension(1));
    EXPECT_EQ(shaped(0, 0), 0.01f);
    EXPECT_EQ(shaped(5, 19), 0.02f);
  }
  {
    auto shaped = t.shaped<float, 3>({6, 4, 5});
    EXPECT_EQ(6, shaped.dimension(0));
    EXPECT_EQ(4, shaped.dimension(1));
    EXPECT_EQ(5, shaped.dimension(2));
    EXPECT_EQ(shaped(0, 0, 0), 0.01f);
    EXPECT_EQ(shaped(5, 3, 4), 0.02f);
  }
  {
    auto shaped = t.shaped<float, 4>({2, 3, 4, 5});
    EXPECT_EQ(2, shaped.dimension(0));
    EXPECT_EQ(3, shaped.dimension(1));
    EXPECT_EQ(4, shaped.dimension(2));
    EXPECT_EQ(5, shaped.dimension(3));

    EXPECT_EQ(shaped(0, 0, 0, 0), 0.01f);
    EXPECT_EQ(shaped(1, 2, 3, 4), 0.02f);
  }
  {
    auto flat = t.flat<float>();
    EXPECT_EQ(flat(0), 0.01f);
    EXPECT_EQ(120, flat.dimension(0));
    EXPECT_EQ(flat(0), 0.01f);
    EXPECT_EQ(flat(119), 0.02f);
  }
  {
    auto flat_inner_dims = t.flat_inner_dims<float>();
    EXPECT_EQ(24, flat_inner_dims.dimension(0));
    EXPECT_EQ(5, flat_inner_dims.dimension(1));
    EXPECT_EQ(flat_inner_dims(0, 0), 0.01f);
    EXPECT_EQ(flat_inner_dims(23, 4), 0.02f);
  }
}

TEST(Tensor_Scalar, Basics) {
  {
    Tensor t(DT_FLOAT, TensorShape({}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.scalar<float>();
    EXPECT_EQ(1, Tt.size());
    EXPECT_EQ(0, Tt.rank());
    t.scalar<float>()() = 123.45f;
    EXPECT_FLOAT_EQ(123.45f, Tt());
  }
  {
    Tensor t(DT_FLOAT, TensorShape({1}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.vec<float>();
    EXPECT_EQ(1, Tt.size());
    t.vec<float>()(0) = 123.45f;
    EXPECT_FLOAT_EQ(123.45f, Tt(0));
  }
  {
    Tensor t(DT_FLOAT, TensorShape({1, 1, 1}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.scalar<float>();
    EXPECT_EQ(1, Tt.size());
    EXPECT_EQ(0, Tt.rank());
    t.flat<float>()(0) = 123.45f;
    EXPECT_FLOAT_EQ(123.45f, Tt());
  }
  {
    Tensor t(DT_STRING, TensorShape({}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.scalar<string>();
    EXPECT_EQ(1, Tt.size());
    EXPECT_EQ(0, Tt.rank());
    t.scalar<string>()() = "foo";
    EXPECT_EQ("foo", Tt());
  }
  {
    Tensor t(DT_STRING, TensorShape({1}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.vec<string>();
    EXPECT_EQ(1, Tt.size());
    t.flat<string>()(0) = "foo";
    EXPECT_EQ("foo", Tt(0));
  }
  {
    Tensor t(DT_STRING, TensorShape({1, 1, 1}));
    EXPECT_EQ(1, t.NumElements());
    auto Tt = t.scalar<string>();
    EXPECT_EQ(1, Tt.size());
    EXPECT_EQ(0, Tt.rank());
    t.flat<string>()(0) = "bar";
    EXPECT_EQ("bar", Tt());
  }
  {
    Tensor t(DT_FLOAT, TensorShape({0, 1}));
    EXPECT_EQ(0, t.NumElements());
    auto Tt = t.flat<float>();
    EXPECT_EQ(0, Tt.size());
    auto Tm = t.matrix<float>();
    EXPECT_EQ(0, Tm.size());
    EXPECT_EQ(0, Tm.dimensions()[0]);
    EXPECT_EQ(1, Tm.dimensions()[1]);
  }
}

TEST(Tensor_Float, Reshape_And_Slice_Assignment) {
  // A test to experiment with a way to assign to a subset of a tensor
  Tensor t(DT_FLOAT, TensorShape({10, 4, 3, 2}));
  EXPECT_TRUE(t.shape().IsSameSize(TensorShape({10, 4, 3, 2})));

  // Get the N dimensional tensor (N==4 here)
  auto e_t = t.tensor<float, 4>();
  // Reshape to view it as a two-dimensional tensor
  auto e_2d = t.shaped<float, 2>({10, 4 * 3 * 2});
  for (int i = 0; i < 10; i++) {
    // Assign a 1 x 4*3*2 matrix (really vector) to a slice of size
    // 1 x 4*3*2 in e_t.
    Eigen::Tensor<float, 2, Eigen::RowMajor> m(1, 4 * 3 * 2);
    m.setConstant(i * 2.0);

    Eigen::DSizes<Eigen::DenseIndex, 2> indices(i, 0);
    Eigen::DSizes<Eigen::DenseIndex, 2> sizes(1, 4 * 3 * 2);
    e_2d.slice(indices, sizes) = m;
  }
  for (int i = 0; i < 10; i++) {
    for (int j = 0; j < 4; j++) {
      for (int k = 0; k < 3; k++) {
        for (int l = 0; l < 2; l++) {
          EXPECT_EQ(e_t(i, j, k, l), i * 2.0f);
          LOG(INFO) << i << "," << j << "," << k << "," << l
                    << " &e_t(i, j, k, l): " << &e_t(i, j, k, l) << " = "
                    << e_t(i, j, k, l);
        }
      }
    }
  }
}

TEST(Tensor_String, Simple) {
  Tensor t = test::AsTensor<string>(
      {"hello", "world", "machine", "learning", "new", "york"},
      TensorShape({3, 2}));
  auto s = t.shape();
  ASSERT_EQ(s.dims(), 2);
  ASSERT_EQ(s.dim_size(0), 3);
  ASSERT_EQ(s.dim_size(1), 2);
  auto m = t.matrix<string>();
  EXPECT_EQ(t.TotalBytes(), 3 * 2 * sizeof(string) + 5 + 5 + 7 + 8 + 3 + 4);

  EXPECT_EQ(m(0, 0), "hello");
  EXPECT_EQ(m(0, 1), "world");
  EXPECT_EQ(m(1, 0), "machine");
  EXPECT_EQ(m(1, 1), "learning");
  EXPECT_EQ(m(2, 0), "new");
  EXPECT_EQ(m(2, 1), "york");

  TestCopies<string>(t);
}

TEST(Tensor_Float, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<float>({0, 1, 2, 3, 4, 5}, {2, 3});
  Tensor t2(t1.dtype(), t1.shape());
  t2.flat<float>() = t1.flat<float>() * 2.0f;
  Tensor t3 = test::AsTensor<float>({0, 2, 4, 6, 8, 10}, t1.shape());
  test::ExpectTensorEqual<float>(t2, t3);
}

TEST(Tensor_Int32, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<int32>({0, 1, 2, 3, 4, 5}, {2, 3});
  Tensor t2(t1.dtype(), t1.shape());
  t2.flat<int32>() = t1.flat<int32>() * 2;
  Tensor t3 = test::AsTensor<int32>({0, 2, 4, 6, 8, 10}, t1.shape());
  test::ExpectTensorEqual<int32>(t2, t3);
}

TEST(Tensor_QInt8, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<qint8>({0, 1, 2, 3, 4, 5}, {2, 3});
  Tensor t2(t1.dtype(), t1.shape());
  t2.flat<qint8>() = t1.flat<qint8>() + qint8(-2);
  Tensor t3 = test::AsTensor<qint8>({-2, -1, 0, 1, 2, 3}, {2, 3});
  test::ExpectTensorEqual<qint8>(t2, t3);
}

TEST(Tensor_QUInt8, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<quint8>({0, 1, 2, 3, 4, 5}, {2, 3});
  Tensor t2(t1.dtype(), t1.shape());
  t2.flat<quint8>() = t1.flat<quint8>() + quint8(2);
  Tensor t3 = test::AsTensor<quint8>({2, 3, 4, 5, 6, 7}, {2, 3});
  test::ExpectTensorEqual<quint8>(t2, t3);
}

TEST(Tensor_Int64, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<int64>(
      {0LL << 48, 1LL << 48, 2LL << 48, 3LL << 48, 4LL << 48, 5LL << 48},
      {2, 3});
  Tensor t2(t1.dtype(), t1.shape());
  t2.flat<int64>() = t1.flat<int64>() * static_cast<int64>(2);
  Tensor t3 = test::AsTensor<int64>(
      {0LL << 48, 2LL << 48, 4LL << 48, 6LL << 48, 8LL << 48, 10LL << 48},
      {2, 3});
  test::ExpectTensorEqual<int64>(t2, t3);
}

TEST(Tensor_String, SimpleWithHelper) {
  Tensor t1 = test::AsTensor<string>({"0", "1", "2", "3", "4", "5"}, {2, 3});
  Tensor t2(DT_STRING, {2, 3});
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      t2.matrix<string>()(i, j) = strings::StrCat(i * 3 + j);
    }
  }

  // Test with helper.
  test::ExpectTensorEqual<string>(t1, t2);
}

TEST(Tensor_Bool, SimpleWithHelper) {
  Tensor t1 =
      test::AsTensor<bool>({false, true, false, true, false, true}, {2, 3});

  Tensor t2(DT_BOOL, {2, 3});
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      t2.matrix<bool>()(i, j) = (((i + j) % 2) != 0);
    }
  }

  // Test with helper.
  test::ExpectTensorEqual<bool>(t1, t2);
}

TEST(Tensor_Complex, Simple) {
  Tensor t(DT_COMPLEX64, {4, 5, 3, 7});
  t.flat<complex64>().setRandom();
  TestCopies<complex64>(t);
}

TEST(Tensor_Complex, SimpleWithHelper) {
  {
    Tensor t1 = test::AsTensor<complex64>({0,
                                           {1, 1},
                                           complex64(2),
                                           complex64(3, 3),
                                           complex64(0, 4),
                                           complex64(2, 5)},
                                          {2, 3});
    Tensor t2(t1.dtype(), t1.shape());
    t2.flat<complex64>() = t1.flat<complex64>() * complex64(0, 2);
    Tensor t3 = test::AsTensor<complex64>(
        {0, {-2, 2}, {0, 4}, {-6, 6}, {-8, 0}, {-10, 4}},
        // shape
        {2, 3});
    test::ExpectTensorEqual<complex64>(t2, t3);
  }

  // Does some numeric operations for complex numbers.
  {
    const float PI = std::acos(-1);
    const complex64 rotate_45 = std::polar(1.0f, PI / 4);

    // x contains all the 8-th root of unity.
    Tensor x(DT_COMPLEX64, TensorShape({8}));
    for (int i = 0; i < 8; ++i) {
      x.vec<complex64>()(i) = std::pow(rotate_45, i);
    }

    // Shift the roots by 45 degree.
    Tensor y(DT_COMPLEX64, TensorShape({8}));
    y.vec<complex64>() = x.vec<complex64>() * rotate_45;
    Tensor y_expected(DT_COMPLEX64, TensorShape({8}));
    for (int i = 0; i < 8; ++i) {
      y_expected.vec<complex64>()(i) = std::pow(rotate_45, i + 1);
    }
    test::ExpectTensorNear<complex64>(y, y_expected, 1e-5);

    // Raise roots to the power of 8.
    Tensor z(DT_COMPLEX64, TensorShape({8}));
    z.vec<complex64>() = x.vec<complex64>().pow(8);
    Tensor z_expected(DT_COMPLEX64, TensorShape({8}));
    for (int i = 0; i < 8; ++i) {
      z_expected.vec<complex64>()(i) = 1;
    }
    test::ExpectTensorNear<complex64>(z, z_expected, 1e-5);
  }
}

// On the alignment.
//
// As of 2015/8, tensorflow::Tensor allocates its buffer with 32-byte
// alignment. Tensor::tensor/flat/vec/matrix methods requires the the
// buffer satisfies Eigen::Aligned (e.g., 16-bytes aligned usually,
// and 32-bytes for AVX). Tensor::Slice requires the caller to ensure
// its result is aligned if the caller intends to use those methods.
// In this test case, we simply make sure each slice is 32-byte
// aligned: sizeof(float) * 4 * 2 = 32.
TEST(Tensor, Slice_Basic) {
  Tensor saved;
  {  // General
    Tensor x(DT_FLOAT, TensorShape({10, 4, 34}));
    // Fills in known values.
    for (int i = 0; i < 10; ++i) {
      x.Slice(i, i + 1).flat<float>().setConstant(i * 1.f);
    }
    // A simple slice along dim0.
    Tensor y = x.Slice(4, 8);
    EXPECT_TRUE(y.shape().IsSameSize(TensorShape({4, 4, 34})));
    auto tx = x.tensor<float, 3>();
    auto ty = y.tensor<float, 3>();
    for (int i = 0; i < 4; ++i) {
      for (int j = 0; j < 4; ++j) {
        for (int k = 0; k < 34; ++k) {
          EXPECT_EQ(ty(i, j, k), 4.0 + i);
          EXPECT_EQ(&tx(4 + i, j, k), &ty(i, j, k));
        }
      }
    }
    // A simple slice equivalent to identity.
    TestCopies<float>(y);
    y = x.Slice(0, 10);
    test::ExpectTensorEqual<float>(x, y);
    EXPECT_EQ(x.flat<float>().data(), y.flat<float>().data());

    // A slice of a slice.
    auto z = x.Slice(4, 8).Slice(2, 3);
    auto tz = z.tensor<float, 3>();
    EXPECT_EQ(1, z.dim_size(0));
    for (int j = 0; j < 4; ++j) {
      for (int k = 0; k < 34; ++k) {
        EXPECT_EQ(tz(0, j, k), 6.0);
      }
    }

    // x and y will be out of scope. But 'saved' should be alive.
    saved = z;
  }
  {
    EXPECT_EQ(1, saved.dim_size(0));
    auto tsaved = saved.tensor<float, 3>();
    for (int j = 0; j < 4; ++j) {
      for (int k = 0; k < 34; ++k) {
        EXPECT_EQ(tsaved(0, j, k), 6.0);
      }
    }
  }
  {  // Empty
    Tensor x(DT_FLOAT, TensorShape({10, 0, 34}));
    x.flat<float>().setRandom();
    Tensor y = x.Slice(4, 8);
    EXPECT_TRUE(y.shape().IsSameSize(TensorShape({4, 0, 34})));
  }

  {
    // Test unaligned access via a Slice.
    Tensor x(DT_FLOAT, TensorShape({30}));
    x.flat<float>().setConstant(0.0);

    // Take an unaligned slice.
    Tensor y = x.Slice(1, 13);
    y.unaligned_flat<float>().setConstant(1.0);
    for (int64 i = 0; i < y.NumElements(); ++i) {
      EXPECT_EQ(1.0, y.unaligned_flat<float>()(i));
    }
  }
}

static void BM_CreateAndDestroy(int iters) {
  TensorShape shape({10, 20});
  while (--iters) {
    Tensor t(DT_FLOAT, shape);
  }
}
BENCHMARK(BM_CreateAndDestroy);

static void BM_Assign(int iters) {
  Tensor a(DT_FLOAT, TensorShape({10, 20}));
  Tensor b(DT_FLOAT, TensorShape({10, 20}));
  bool a_to_b = true;
  while (--iters) {
    if (a_to_b) {
      b = a;
    } else {
      a = b;
    }
    a_to_b = !a_to_b;
  }
}
BENCHMARK(BM_Assign);

// Ensure tensor_data() works on empty tensors
TEST(Tensor, EmptyTensorData) {
  Tensor empty;
  EXPECT_EQ(empty.tensor_data().size(), 0);
}

}  // namespace tensorflow