text: // 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 <gtest/gtest.h>
  
  #include <iostream>
  #include <sstream>
  
  #include "knowhere/adapter/SptagAdapter.h"
  #include "knowhere/adapter/Structure.h"
  #include "knowhere/common/Exception.h"
  #include "knowhere/index/vector_index/IndexSPTAG.h"
  #include "knowhere/index/vector_index/helpers/Definitions.h"
  
  #include "unittest/utils.h"
  
  using ::testing::Combine;
  using ::testing::TestWithParam;
  using ::testing::Values;
  
  class SPTAGTest : public DataGen, public TestWithParam<std::string> {
   protected:
      void
      SetUp() override {
          IndexType = GetParam();
          Generate(128, 100, 5);
          index_ = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
          if (IndexType == "KDT") {
              auto tempconf = std::make_shared<knowhere::KDTCfg>();
              tempconf->tptnumber = 1;
              tempconf->k = 10;
              conf = tempconf;
          } else {
              auto tempconf = std::make_shared<knowhere::BKTCfg>();
              tempconf->tptnumber = 1;
              tempconf->k = 10;
              conf = tempconf;
          }
  
          Init_with_default();
      }
  
   protected:
      knowhere::Config conf;
      std::shared_ptr<knowhere::CPUSPTAGRNG> index_ = nullptr;
      std::string IndexType;
  };
  
  INSTANTIATE_TEST_CASE_P(SPTAGParameters, SPTAGTest, Values("KDT", "BKT"));
  
  // TODO(lxj): add test about count() and dimension()
  TEST_P(SPTAGTest, sptag_basic) {
      assert(!xb.empty());
  
      auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
      index_->set_preprocessor(preprocessor);
  
      auto model = index_->Train(base_dataset, conf);
      index_->set_index_model(model);
      index_->Add(base_dataset, conf);
      auto result = index_->Search(query_dataset, conf);
      AssertAnns(result, nq, k);
  
      {
          // auto ids = result->array()[0];
          // auto dists = result->array()[1];
          auto ids = result->ids();
          auto dists = result->dist();
  
          std::stringstream ss_id;
          std::stringstream ss_dist;
          for (auto i = 0; i < nq; i++) {
              for (auto j = 0; j < k; ++j) {
                  // ss_id << *ids->data()->GetValues<int64_t>(1, i * k + j) << " ";
                  // ss_dist << *dists->data()->GetValues<float>(1, i * k + j) << " ";
                  ss_id << *((int64_t*)(ids) + i * k + j) << " ";
                  ss_dist << *((float*)(dists) + i * k + j) << " ";
              }
              ss_id << std::endl;
              ss_dist << std::endl;
          }
          std::cout << "id\n" << ss_id.str() << std::endl;
          std::cout << "dist\n" << ss_dist.str() << std::endl;
      }
  }
  
  TEST_P(SPTAGTest, sptag_serialize) {
      assert(!xb.empty());
  
      auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
      index_->set_preprocessor(preprocessor);
  
      auto model = index_->Train(base_dataset, conf);
  
      index_->Add(base_dataset, conf);
      auto binaryset = index_->Serialize();
      auto new_index = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
      new_index->Load(binaryset);
      auto result = new_index->Search(query_dataset, conf);
      AssertAnns(result, nq, k);
      PrintResult(result, nq, k);
      ASSERT_EQ(new_index->Count(), nb);
      ASSERT_EQ(new_index->Dimension(), dim);
      //        ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException);
      //        ASSERT_NO_THROW({ new_index->Seal(); });
  
      {
          int fileno = 0;
          const std::string& base_name = "/tmp/sptag_serialize_test_bin_";
          std::vector<std::string> filename_list;
          std::vector<std::pair<std::string, size_t>> meta_list;
          for (auto& iter : binaryset.binary_map_) {
              const std::string& filename = base_name + std::to_string(fileno);
              FileIOWriter writer(filename);
              writer(iter.second->data.get(), iter.second->size);
  
              meta_list.emplace_back(std::make_pair(iter.first, iter.second->size));
              filename_list.push_back(filename);
              ++fileno;
          }
  
          knowhere::BinarySet load_data_list;
          for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) {
              auto bin_size = meta_list[i].second;
              FileIOReader reader(filename_list[i]);
  
              auto load_data = new uint8_t[bin_size];
              reader(load_data, bin_size);
              auto data = std::make_shared<uint8_t>();
              data.reset(load_data);
              load_data_list.Append(meta_list[i].first, data, bin_size);
          }
  
          auto new_index = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
          new_index->Load(load_data_list);
          auto result = new_index->Search(query_dataset, conf);
          AssertAnns(result, nq, k);
          PrintResult(result, nq, k);
      }
  }
  