// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed 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 <algorithm>
#include <vector>

#include "paddle/fluid/framework/fleet/heter_ps/feature_value.h"
#include "paddle/fluid/framework/fleet/heter_ps/graph_gpu_ps_table.h"
#include "paddle/fluid/framework/fleet/heter_ps/graph_gpu_wrapper.h"
#include "paddle/fluid/framework/fleet/heter_ps/heter_comm.h"
#include "paddle/fluid/framework/fleet/heter_ps/heter_resource.h"
#include "paddle/fluid/framework/fleet/heter_ps/optimizer.cuh.h"
#include "paddle/fluid/platform/cuda_device_guard.h"

using namespace paddle::framework;
namespace platform = paddle::platform;
// paddle::framework::GpuPsCommGraph GraphTable::make_gpu_ps_graph
// paddle::framework::GpuPsCommGraph GraphTable::make_gpu_ps_graph(
//     std::vector<int64_t> ids)

std::string edges[] = {
    std::string("0\t1"),
    std::string("0\t9"),
    std::string("1\t2"),
    std::string("1\t0"),
    std::string("2\t1"),
    std::string("2\t3"),
    std::string("3\t2"),
    std::string("3\t4"),
    std::string("4\t3"),
    std::string("4\t5"),
    std::string("5\t4"),
    std::string("5\t6"),
    std::string("6\t5"),
    std::string("6\t7"),
    std::string("7\t6"),
    std::string("7\t8"),
};
char edge_file_name[] = "edges1.txt";

std::string nodes[] = {
    std::string("user\t37\ta 0.34\tb 13 14\tc hello\td abc"),
    std::string("user\t96\ta 0.31\tb 15 10\tc 96hello\td abcd"),
    std::string("user\t59\ta 0.11\tb 11 14"),
    std::string("user\t97\ta 0.11\tb 12 11"),
    std::string("item\t45\ta 0.21"),
    std::string("item\t145\ta 0.21"),
    std::string("item\t112\ta 0.21"),
    std::string("item\t48\ta 0.21"),
    std::string("item\t247\ta 0.21"),
    std::string("item\t111\ta 0.21"),
    std::string("item\t46\ta 0.21"),
    std::string("item\t146\ta 0.21"),
    std::string("item\t122\ta 0.21"),
    std::string("item\t49\ta 0.21"),
    std::string("item\t248\ta 0.21"),
    std::string("item\t113\ta 0.21")};
char node_file_name[] = "nodes.txt";
std::vector<std::string> user_feature_name = {"a", "b", "c", "d"};
std::vector<std::string> item_feature_name = {"a"};
std::vector<std::string> user_feature_dtype = {
    "float32", "int32", "string", "string"};
std::vector<std::string> item_feature_dtype = {"float32"};
std::vector<int> user_feature_shape = {1, 2, 1, 1};
std::vector<int> item_feature_shape = {1};
void prepare_file(char file_name[], bool load_edge) {
  std::ofstream ofile;
  ofile.open(file_name);
  if (load_edge) {
    for (auto x : edges) {
      ofile << x << std::endl;
    }
  } else {
    for (auto x : nodes) {
      ofile << x << std::endl;
    }
  }
  ofile.close();
}
TEST(TEST_FLEET, test_cpu_cache) {
  int gpu_num = 0;
  int st = 0, u = 0;
  std::vector<int> device_id_mapping;
  for (int i = 0; i < 2; i++) device_id_mapping.push_back(i);
  gpu_num = device_id_mapping.size();
  ::paddle::distributed::GraphParameter table_proto;
  table_proto.add_edge_types("u2u");
  table_proto.add_node_types("user");
  table_proto.add_node_types("item");
  ::paddle::distributed::GraphFeature *g_f = table_proto.add_graph_feature();

  for (int i = 0; i < user_feature_name.size(); i++) {
    g_f->add_name(user_feature_name[i]);
    g_f->add_dtype(user_feature_dtype[i]);
    g_f->add_shape(user_feature_shape[i]);
  }
  ::paddle::distributed::GraphFeature *g_f1 = table_proto.add_graph_feature();
  for (int i = 0; i < item_feature_name.size(); i++) {
    g_f1->add_name(item_feature_name[i]);
    g_f1->add_dtype(item_feature_dtype[i]);
    g_f1->add_shape(item_feature_shape[i]);
  }
  prepare_file(node_file_name, false);
  prepare_file(edge_file_name, true);
  table_proto.set_shard_num(24);
  table_proto.set_search_level(2);
  std::shared_ptr<HeterPsResource> resource =
      std::make_shared<HeterPsResource>(device_id_mapping);
  resource->enable_p2p();
  int use_nv = 1;
  GpuPsGraphTable g(resource, use_nv);
  g.init_cpu_table(table_proto);
  g.cpu_graph_table->Load(node_file_name, "nuser");
  g.cpu_graph_table->Load(node_file_name, "nitem");
  std::remove(node_file_name);
  std::vector<paddle::framework::GpuPsCommGraph> vec;
  std::vector<int64_t> node_ids;
  node_ids.push_back(37);
  node_ids.push_back(96);
  std::vector<std::vector<std::string>> node_feat(2,
                                                  std::vector<std::string>(2));
  std::vector<std::string> feature_names;
  feature_names.push_back(std::string("c"));
  feature_names.push_back(std::string("d"));
  g.cpu_graph_table->get_node_feat(0, node_ids, feature_names, node_feat);
  VLOG(0) << "get_node_feat: " << node_feat[0][0];
  VLOG(0) << "get_node_feat: " << node_feat[0][1];
  VLOG(0) << "get_node_feat: " << node_feat[1][0];
  VLOG(0) << "get_node_feat: " << node_feat[1][1];
  int n = 10;
  std::vector<int64_t> ids0, ids1;
  for (int i = 0; i < n; i++) {
    g.cpu_graph_table->add_comm_edge(0, i, (i + 1) % n);
    g.cpu_graph_table->add_comm_edge(0, i, (i - 1 + n) % n);
    if (i % 2 == 0) ids0.push_back(i);
  }
  g.cpu_graph_table->build_sampler(0);
  ids1.push_back(5);
  ids1.push_back(7);
  vec.push_back(g.cpu_graph_table->make_gpu_ps_graph(0, ids0));
  vec.push_back(g.cpu_graph_table->make_gpu_ps_graph(0, ids1));
  vec[0].display_on_cpu();
  vec[1].display_on_cpu();
  // g.build_graph_from_cpu(vec);
  g.build_graph_on_single_gpu(vec[0], 0);
  g.build_graph_on_single_gpu(vec[1], 1);
  int64_t cpu_key[3] = {0, 1, 2};
  /*
  std::vector<std::shared_ptr<char>> buffers(3);
  std::vector<int> actual_sizes(3,0);
  g.cpu_graph_table->random_sample_neighbors(cpu_key,2,buffers,actual_sizes,false);
  for(int i = 0;i < 3;i++){
    VLOG(0)<<"sample from cpu key->"<<cpu_key[i]<<" actual sample size =
  "<<actual_sizes[i]/sizeof(int64_t);
  }
  */
  void *key;
  int device_len = 2;
  for (int i = 0; i < 2; i++) {
    // platform::CUDADeviceGuard guard(i);
    LOG(0) << "query on card " << i;
    //{1,9} or {9,1} is expected for key 0
    //{0,2} or {2,0} is expected for key 1
    //{1,3} or {3,1} is expected for key 2
    int step = 2;
    int cur = 0;
    while (true) {
      auto node_query_res = g.query_node_list(i, cur, step);
      node_query_res.display();
      if (node_query_res.get_len() == 0) {
        VLOG(0) << "no more ids,break";
        break;
      }
      cur += node_query_res.get_len();
      NeighborSampleQuery query;
      query.initialize(
          i, node_query_res.get_val(), 1, node_query_res.get_len());
      query.display();
      auto c = g.graph_neighbor_sample_v3(query, false);
      c.display();
    }
  }
  g.cpu_graph_table->set_search_level(2);
  // g.cpu_graph_table->Load_to_ssd(edge_file_name,"e>u2u");
  g.cpu_graph_table->Load(edge_file_name, "e>u2u");
  g.cpu_graph_table->make_partitions(0, 64, 2);
  int index = 0;
  while (g.cpu_graph_table->load_next_partition(0) != -1) {
    auto all_ids = g.cpu_graph_table->get_all_id(0, 0, device_len);
    for (auto x : all_ids) {
      for (auto y : x) {
        VLOG(0) << "part " << index << " " << y;
      }
    }
    for (int i = 0; i < all_ids.size(); i++) {
      GpuPsCommGraph sub_graph =
          g.cpu_graph_table->make_gpu_ps_graph(0, all_ids[i]);
      g.build_graph_on_single_gpu(sub_graph, i);
      VLOG(2) << "sub graph on gpu " << i << " is built";
    }
    VLOG(0) << "start to iterate gpu graph node";
    g.cpu_graph_table->make_complementary_graph(0, 64);
    for (int i = 0; i < 2; i++) {
      // platform::CUDADeviceGuard guard(i);
      LOG(0) << "query on card " << i;
      int step = 2;
      int cur = 0;
      while (true) {
        auto node_query_res = g.query_node_list(i, cur, step);
        node_query_res.display();
        if (node_query_res.get_len() == 0) {
          VLOG(0) << "no more ids,break";
          break;
        }
        cur += node_query_res.get_len();
        NeighborSampleQuery query, q1;
        query.initialize(
            i, node_query_res.get_val(), 4, node_query_res.get_len());
        query.display();
        auto c = g.graph_neighbor_sample_v3(query, true);
        c.display();
        platform::CUDADeviceGuard guard(i);
        int64_t *key;
        VLOG(0) << "sample key 1 globally";
        g.cpu_graph_table->set_search_level(2);
        cudaMalloc((void **)&key, sizeof(int64_t));
        int64_t t_key = 1;
        cudaMemcpy(key, &t_key, sizeof(int64_t), cudaMemcpyHostToDevice);
        q1.initialize(i, (int64_t)key, 2, 1);
        auto d = g.graph_neighbor_sample_v3(q1, true);
        d.display();
        cudaFree(key);
        g.cpu_graph_table->set_search_level(1);
      }
    }
    index++;
  }
  auto iter = paddle::framework::GraphGpuWrapper::GetInstance();
  std::vector<int> device;
  device.push_back(0);
  device.push_back(1);
  iter->set_device(device);
}
