/*
 * 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 "mxnet/c_api.h"  // MXGetGPUCount()
#include "mxnet-cpp/MxNetCpp.h"

using namespace mxnet::cpp;

static bool test_single_key(const Context &context, const std::string &context_str) {
  std::string key = "singlekeytest-" + context_str;

  NDArray result(Shape(4), context);
  NDArray result_cpu;

  // initialize data
  NDArray data_cpu({0.f, 233.f, -0.12f, 9.f}, Shape(4), Context::cpu());
  NDArray data = data_cpu.Copy(context);
  NDArray::WaitAll();

  KVStore::Init(key, data);
  NDArray::WaitAll();

  // retrieve result
  KVStore::Pull(key, &result);
  NDArray::WaitAll();

  result_cpu = result.Copy(Context::cpu());
  NDArray::WaitAll();

  // compare
  for (size_t j=0; j < result_cpu.Size(); j++) {
    if (result_cpu.GetData()[j] != data_cpu.GetData()[j]) {
      LG << "Error: wrong initialized data in singlekeytest-" << context_str
          << ", expect " << data_cpu.GetData()[j]
          << " got " << result_cpu.GetData()[j];
      return false;
    }
  }

  // push gradient
  NDArray grad_cpu({0.1f, -2.f, -4.4f, 0.f}, Shape(4), Context::cpu());
  NDArray grad = grad_cpu.Copy(context);
  NDArray::WaitAll();

  KVStore::Push(key, grad);
  NDArray::WaitAll();

  // retrieve result
  KVStore::Pull(key, &result);
  NDArray::WaitAll();

  result_cpu = result.Copy(Context::cpu());
  NDArray::WaitAll();

  // compare
  for (size_t j=0; j < result_cpu.Size(); j++) {
    if (result_cpu.GetData()[j] != grad_cpu.GetData()[j]) {
      LG << "Error: wrong gradient data in singlekeytest-" << context_str
          << ", expect " << grad_cpu.GetData()[j]
          << " got " << result_cpu.GetData()[j];
      return false;
    }
  }

  return true;
}

static bool test_multiple_key(const Context &context, const std::string &context_str) {
  std::vector<std::string> keys(2);
  keys[0] = "multikeytest-0-" + context_str;
  keys[1] = "multikeytest-1-" + context_str;

  std::vector<NDArray> results(2);
  results[0] = NDArray(Shape(4), context);
  results[1] = NDArray(Shape(4), context);
  std::vector<NDArray> results_cpu(2);

  // initialize data
  std::vector<NDArray> data_cpu(2);
  data_cpu[0] = NDArray({0.f, 2.f, -3.12f, 4.f}, Shape(4), Context::cpu());
  data_cpu[1] = NDArray({0.8f, -2.f, 6.6f, 77.f}, Shape(4), Context::cpu());
  std::vector<NDArray> data(2);
  data[0] = data_cpu[0].Copy(context);
  data[1] = data_cpu[1].Copy(context);
  NDArray::WaitAll();

  KVStore::Init(keys, data);
  NDArray::WaitAll();

  // retrieve result
  KVStore::Pull(keys, &results);
  NDArray::WaitAll();

  results_cpu[0] = results[0].Copy(Context::cpu());
  results_cpu[1] = results[1].Copy(Context::cpu());
  NDArray::WaitAll();

  // compare
  for (size_t i=0; i < results_cpu.size(); i++) {
    for (size_t j=0; j < results_cpu[i].Size(); j++) {
      if (results_cpu[i].GetData()[j] != data_cpu[i].GetData()[j]) {
        LG << "Error: wrong initialized data in multikeytest-" << context_str
            << ", expect " << data_cpu[i].GetData()[j]
            << " got " << results_cpu[i].GetData()[j];
        return false;
      }
    }
  }

  // push gradient, reduce for the second
  std::vector<std::string> push_keys(3);
  push_keys[0] = "multikeytest-0-" + context_str;
  push_keys[1] = "multikeytest-1-" + context_str;
  push_keys[2] = "multikeytest-1-" + context_str;

  std::vector<NDArray> grads_cpu(3);
  grads_cpu[0] = NDArray({0.2f, -0.3f, -1.1f, 0.0f}, Shape(4), Context::cpu());
  grads_cpu[1] = NDArray({2.f, 4.f, -4.f, -5.f}, Shape(4), Context::cpu());
  grads_cpu[2] = NDArray({-3.f, -0.2f, 12.f, -9.f}, Shape(4), Context::cpu());
  std::vector<NDArray> grads(3);
  grads[0] = grads_cpu[0].Copy(context);
  grads[1] = grads_cpu[1].Copy(context);
  grads[2] = grads_cpu[2].Copy(context);
  NDArray::WaitAll();

  KVStore::Push(push_keys, grads);
  NDArray::WaitAll();

  // retrieve result
  KVStore::Pull(keys, &results);
  NDArray::WaitAll();

  results_cpu[0] = results[0].Copy(Context::cpu());
  results_cpu[1] = results[1].Copy(Context::cpu());
  NDArray::WaitAll();

  // compare the first
  for (size_t j=0; j < results_cpu[0].Size(); j++) {
    if (results_cpu[0].GetData()[j] != grads_cpu[0].GetData()[j]) {
      LG << "Error: wrong gradient data in multikeytest-" << context_str
          << ", expect " << grads_cpu[0].GetData()[j]
          << " got " << results_cpu[0].GetData()[j];
      return false;
    }
  }

  // compare the second
  for (size_t j=0; j < results_cpu[1].Size(); j++) {
    if (results_cpu[1].GetData()[j] != (grads_cpu[1].GetData()[j] + grads_cpu[2].GetData()[j])) {
      LG << "Error: wrong reduced gradient data in multikeytest-" << context_str
          << ", expect " << (grads_cpu[1].GetData()[j] + grads_cpu[2].GetData()[j])
          << " got " << results_cpu[1].GetData()[j];
      return false;
    }
  }

  return true;
}

int main(int argc, char** argv) {
  KVStore::SetType("local");

  bool success1 = test_single_key(Context::cpu(), "cpu");
  bool success2 = test_multiple_key(Context::cpu(), "cpu");

  bool success3 = true;
  bool success4 = true;

  int gpu_count = 0;
  if (MXGetGPUCount(&gpu_count) != 0) {
    LG << "Error: MXGetGPUCount";

    MXNotifyShutdown();
    return 1;
  }

  if (gpu_count > 0) {
    success3 = test_single_key(Context::gpu(), "gpu");
    success4 = test_multiple_key(Context::gpu(), "gpu");
  }

  int ret = (success1 && success2 && success3 && success4) ? 0 : 1;

  MXNotifyShutdown();
  return ret;
}