// Copyright (C) 2019-2020 Zilliz. 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 <folly/FBVector.h>
#include <gtest/gtest.h>
#include <algorithm>
#include <cstddef>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "NamedType/named_type_impl.hpp"
#include "common/EasyAssert.h"
#include "common/QueryResult.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
#include "filemanager/InputStream.h"
#include "gtest/gtest.h"
#include "index/Index.h"
#include "index/VectorIndex.h"
#include "knowhere/comp/index_param.h"
#include "pb/common.pb.h"
#include "pb/plan.pb.h"
#include "query/Plan.h"
#include "query/PlanNode.h"
#include "segcore/SegmentSealed.h"
#include "segcore/Types.h"
#include "segcore/storagev1translator/ChunkTranslator.h"
#include "test_utils/DataGen.h"
#include "test_utils/cachinglayer_test_utils.h"
#include "test_utils/storage_test_utils.h"

using namespace milvus;
using namespace milvus::query;
using namespace milvus::segcore;

TEST(Rescorer, Normal) {
    //0. prepare schema
    int dim = 64;
    auto schema = std::make_shared<Schema>();
    auto vec_fid = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);
    schema->AddDebugField("int8", DataType::INT8);
    schema->AddDebugField("int16", DataType::INT16);
    schema->AddDebugField("int32", DataType::INT32);
    schema->AddDebugField("int64", DataType::INT64);
    auto str_fid = schema->AddDebugField("string", DataType::VARCHAR);
    schema->AddDebugField("bool", DataType::BOOL);
    schema->set_primary_field_id(str_fid);
    size_t N = 50;

    //2. load raw data
    auto raw_data = DataGen(schema, N, 42, 0, 8, 10, false, false);
    auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data);

    //3. load index
    auto vector_data = raw_data.get_col<float>(vec_fid);
    auto indexing = GenVecIndexing(
        N, dim, vector_data.data(), knowhere::IndexEnum::INDEX_HNSW);
    LoadIndexInfo load_index_info;
    load_index_info.field_id = vec_fid.get();
    load_index_info.index_params = GenIndexParams(indexing.get());
    load_index_info.cache_index =
        CreateTestCacheIndex("test", std::move(indexing));
    load_index_info.index_params["metric_type"] = knowhere::metric::L2;
    segment->LoadIndex(load_index_info);

    // no result after search
    {
        ScopedSchemaHandle handle(*schema);
        ScopedSchemaHandle::ScorerParams scorer;
        scorer.weight = 4;
        auto plan = handle.ParseSearchWithScorers(schema,
                                                  "int8 >= 100 and int8 < -1",
                                                  "fakevec",
                                                  10,
                                                  "L2",
                                                  R"({"ef": 50})",
                                                  {scorer});
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    }

    // search result not empty but no boost filter
    {
        ScopedSchemaHandle handle(*schema);
        ScopedSchemaHandle::ScorerParams scorer;
        scorer.weight = 4;
        auto plan = handle.ParseSearchWithScorers(schema,
                                                  "int8 >= -1 and int8 < 100",
                                                  "fakevec",
                                                  10,
                                                  "L2",
                                                  R"({"ef": 50})",
                                                  {scorer});
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    }

    // random function with seed
    {
        ScopedSchemaHandle handle(*schema);
        ScopedSchemaHandle::ScorerParams scorer;
        scorer.type = proto::plan::FunctionTypeRandom;
        scorer.weight = 1;
        scorer.params = {{"seed", "123"}};
        auto plan = handle.ParseSearchWithScorers(schema,
                                                  "int8 >= -1 and int8 < 100",
                                                  "fakevec",
                                                  10,
                                                  "L2",
                                                  R"({"ef": 50})",
                                                  {scorer});
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    }

    // random function with field as random seed
    {
        ScopedSchemaHandle handle(*schema);
        ScopedSchemaHandle::ScorerParams scorer;
        scorer.type = proto::plan::FunctionTypeRandom;
        scorer.weight = 1;
        scorer.params = {{"field", "int64"}};
        auto plan = handle.ParseSearchWithScorers(schema,
                                                  "int8 >= -1 and int8 < 100",
                                                  "fakevec",
                                                  10,
                                                  "L2",
                                                  R"({"ef": 50})",
                                                  {scorer});
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    }

    // random function with field and seed
    {
        ScopedSchemaHandle handle(*schema);
        ScopedSchemaHandle::ScorerParams scorer;
        scorer.type = proto::plan::FunctionTypeRandom;
        scorer.weight = 1;
        scorer.params = {{"seed", "123"}, {"field", "int64"}};
        auto plan = handle.ParseSearchWithScorers(schema,
                                                  "int8 >= -1 and int8 < 100",
                                                  "fakevec",
                                                  10,
                                                  "L2",
                                                  R"({"ef": 50})",
                                                  {scorer});
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);

        auto search_result_same_seed =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);

        // should return same score when use same seed
        for (auto i = 0; i < 10; i++) {
            AssertInfo(search_result->distances_[i] ==
                           search_result_same_seed->distances_[i],
                       "distance not equal %f:%f",
                       search_result->distances_[i],
                       search_result_same_seed->distances_[i]);
        }
    }
}