// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#include "main.h"

#include <Eigen/CXX11/Tensor>

using Eigen::Tensor;
using Eigen::TensorMap;

static void test_assign() {
  std::string data1[6];
  TensorMap<Tensor<std::string, 2>> mat1(data1, 2, 3);
  std::string data2[6];
  const TensorMap<Tensor<const std::string, 2>> mat2(data2, 2, 3);

  for (int i = 0; i < 6; ++i) {
    std::ostringstream s1;
    s1 << "abc" << i * 3;
    data1[i] = s1.str();
    std::ostringstream s2;
    s2 << "def" << i * 5;
    data2[i] = s2.str();
  }

  Tensor<std::string, 2> rslt1;
  rslt1 = mat1;
  Tensor<std::string, 2> rslt2;
  rslt2 = mat2;

  Tensor<std::string, 2> rslt3 = mat1;
  Tensor<std::string, 2> rslt4 = mat2;

  Tensor<std::string, 2> rslt5(mat1);
  Tensor<std::string, 2> rslt6(mat2);

  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      VERIFY_IS_EQUAL(rslt1(i, j), data1[i + 2 * j]);
      VERIFY_IS_EQUAL(rslt2(i, j), data2[i + 2 * j]);
      VERIFY_IS_EQUAL(rslt3(i, j), data1[i + 2 * j]);
      VERIFY_IS_EQUAL(rslt4(i, j), data2[i + 2 * j]);
      VERIFY_IS_EQUAL(rslt5(i, j), data1[i + 2 * j]);
      VERIFY_IS_EQUAL(rslt6(i, j), data2[i + 2 * j]);
    }
  }
}

static void test_concat() {
  Tensor<std::string, 2> t1(2, 3);
  Tensor<std::string, 2> t2(2, 3);

  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      std::ostringstream s1;
      s1 << "abc" << i + j * 2;
      t1(i, j) = s1.str();
      std::ostringstream s2;
      s2 << "def" << i * 5 + j * 32;
      t2(i, j) = s2.str();
    }
  }

  Tensor<std::string, 2> result = t1.concatenate(t2, 1);
  VERIFY_IS_EQUAL(result.dimension(0), 2);
  VERIFY_IS_EQUAL(result.dimension(1), 6);

  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      VERIFY_IS_EQUAL(result(i, j), t1(i, j));
      VERIFY_IS_EQUAL(result(i, j + 3), t2(i, j));
    }
  }
}

static void test_slices() {
  Tensor<std::string, 2> data(2, 6);
  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      std::ostringstream s1;
      s1 << "abc" << i + j * 2;
      data(i, j) = s1.str();
    }
  }

  const Eigen::DSizes<ptrdiff_t, 2> half_size(2, 3);
  const Eigen::DSizes<ptrdiff_t, 2> first_half(0, 0);
  const Eigen::DSizes<ptrdiff_t, 2> second_half(0, 3);

  Tensor<std::string, 2> t1 = data.slice(first_half, half_size);
  Tensor<std::string, 2> t2 = data.slice(second_half, half_size);

  for (int i = 0; i < 2; ++i) {
    for (int j = 0; j < 3; ++j) {
      VERIFY_IS_EQUAL(data(i, j), t1(i, j));
      VERIFY_IS_EQUAL(data(i, j + 3), t2(i, j));
    }
  }
}

static void test_additions() {
  Tensor<std::string, 1> data1(3);
  Tensor<std::string, 1> data2(3);
  for (int i = 0; i < 3; ++i) {
    data1(i) = "abc";
    std::ostringstream s1;
    s1 << i;
    data2(i) = s1.str();
  }

  Tensor<std::string, 1> sum = data1 + data2;
  for (int i = 0; i < 3; ++i) {
    std::ostringstream concat;
    concat << "abc" << i;
    std::string expected = concat.str();
    VERIFY_IS_EQUAL(sum(i), expected);
  }
}

static void test_initialization() {
  Tensor<std::string, 2> a(2, 3);
  a.setConstant(std::string("foo"));
  for (int i = 0; i < 2 * 3; ++i) {
    VERIFY_IS_EQUAL(a(i), std::string("foo"));
  }
}

EIGEN_DECLARE_TEST(cxx11_tensor_of_strings) {
  // Beware: none of this is likely to ever work on a GPU.
  CALL_SUBTEST(test_assign());
  CALL_SUBTEST(test_concat());
  CALL_SUBTEST(test_slices());
  CALL_SUBTEST(test_additions());
  CALL_SUBTEST(test_initialization());
}