// 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.

package indexparamcheck

import (
	"strconv"
	"testing"
)

// TODO: add more test cases which `ConfAdapter.CheckTrain` return false,
//       for example, maybe we can copy test cases from regression test later.

func invalidIVFParamsMin() map[string]string {
	invalidIVFParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(MinNList - 1),
		Metric: L2,
	}
	return invalidIVFParams
}

func invalidIVFParamsMax() map[string]string {
	invalidIVFParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(MaxNList + 1),
		Metric: L2,
	}
	return invalidIVFParams
}

func copyParams(original map[string]string) map[string]string {
	result := make(map[string]string)
	for key, value := range original {
		result[key] = value
	}
	return result
}

func TestBaseConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		Metric: L2,
	}
	paramsWithoutDim := map[string]string{
		Metric: L2,
	}
	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{paramsWithoutDim, false},
	}

	adapter := newBaseConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("BaseConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

// IVFConfAdapter checks if an ivf index can be built.
func TestIVFConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(1024),
		Metric: L2,
	}

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{invalidIVFParamsMin(), false},
		{invalidIVFParamsMax(), false},
	}

	adapter := newIVFConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("IVFConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

func TestIVFPQConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(1024),
		IVFM:   strconv.Itoa(4),
		NBITS:  strconv.Itoa(8),
		Metric: L2,
	}

	validParamsWithoutNbits := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(1024),
		IVFM:   strconv.Itoa(4),
		Metric: L2,
	}

	validParamsWithoutDim := map[string]string{
		NLIST:  strconv.Itoa(1024),
		IVFM:   strconv.Itoa(4),
		NBITS:  strconv.Itoa(8),
		Metric: L2,
	}

	invalidParamsDim := copyParams(validParams)
	invalidParamsDim[DIM] = "NAN"

	invalidParamsNbits := copyParams(validParams)
	invalidParamsNbits[NBITS] = "NAN"

	invalidParamsWithoutIVF := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(1024),
		NBITS:  strconv.Itoa(8),
		Metric: L2,
	}

	invalidParamsIVF := copyParams(validParams)
	invalidParamsIVF[IVFM] = "NAN"

	invalidParamsM := copyParams(validParams)
	invalidParamsM[DIM] = strconv.Itoa(65536)

	invalidParamsMzero := copyParams(validParams)
	invalidParamsMzero[IVFM] = "0"

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{validParamsWithoutNbits, true},
		{invalidIVFParamsMin(), false},
		{invalidIVFParamsMax(), false},
		{validParamsWithoutDim, false},
		{invalidParamsDim, false},
		{invalidParamsNbits, false},
		{invalidParamsWithoutIVF, false},
		{invalidParamsIVF, false},
		{invalidParamsM, false},
		{invalidParamsMzero, false},
	}

	adapter := newIVFPQConfAdapter()
	for i, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Log("i:", i, "params", test.params)
			t.Errorf("IVFPQConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

func TestIVFSQConfAdapter_CheckTrain(t *testing.T) {
	getValidParams := func(withNBits bool) map[string]string {
		validParams := map[string]string{
			DIM:    strconv.Itoa(128),
			NLIST:  strconv.Itoa(100),
			NBITS:  strconv.Itoa(8),
			Metric: L2,
		}
		if withNBits {
			validParams[NBITS] = strconv.Itoa(DefaultNBits)
		}
		return validParams
	}
	validParams := getValidParams(false)
	validParamsWithNBits := getValidParams(true)
	paramsWithInvalidNBits := getValidParams(false)
	paramsWithInvalidNBits[NBITS] = strconv.Itoa(DefaultNBits + 1)
	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{validParamsWithNBits, true},
		{paramsWithInvalidNBits, false},
		{invalidIVFParamsMin(), false},
		{invalidIVFParamsMax(), false},
	}

	adapter := newIVFSQConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("IVFSQConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

// BinIDMAPConfAdapter checks if a bin id map index can be built.
func TestBinIDMAPConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		Metric: JACCARD,
	}
	paramsWithoutDim := map[string]string{
		Metric: JACCARD,
	}
	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{paramsWithoutDim, false},
	}

	adapter := newBinIDMAPConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("BinIDMAPConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

// BinIVFConfAdapter checks if a bin ivf index can be built.
func TestBinIVFConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NLIST:  strconv.Itoa(100),
		IVFM:   strconv.Itoa(4),
		NBITS:  strconv.Itoa(8),
		Metric: JACCARD,
	}
	paramsWithoutDim := map[string]string{
		NLIST:  strconv.Itoa(100),
		IVFM:   strconv.Itoa(4),
		NBITS:  strconv.Itoa(8),
		Metric: JACCARD,
	}

	invalidParams := copyParams(validParams)
	invalidParams[Metric] = L2

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{paramsWithoutDim, false},
		{invalidIVFParamsMin(), false},
		{invalidIVFParamsMax(), false},
		{invalidParams, false},
	}

	adapter := newBinIVFConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("BinIVFConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

// HNSWConfAdapter checks if a hnsw index can be built.
func TestHNSWConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:            strconv.Itoa(128),
		HNSWM:          strconv.Itoa(16),
		EFConstruction: strconv.Itoa(200),
		Metric:         L2,
	}

	invalidEfParamsMin := copyParams(validParams)
	invalidEfParamsMin[EFConstruction] = strconv.Itoa(HNSWMinEfConstruction - 1)

	invalidEfParamsMax := copyParams(validParams)
	invalidEfParamsMax[EFConstruction] = strconv.Itoa(HNSWMaxEfConstruction + 1)

	invalidMParamsMin := copyParams(validParams)
	invalidMParamsMin[HNSWM] = strconv.Itoa(HNSWMinM - 1)

	invalidMParamsMax := copyParams(validParams)
	invalidMParamsMax[HNSWM] = strconv.Itoa(HNSWMaxM + 1)

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{invalidEfParamsMin, false},
		{invalidEfParamsMax, false},
		{invalidMParamsMin, false},
		{invalidMParamsMax, false},
	}

	adapter := newHNSWConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("HNSWConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

// ANNOYConfAdapter checks if an annoy index can be built
func TestANNOYConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		NTREES: strconv.Itoa(4),
		Metric: L2,
	}

	invalidTreeParamsMin := copyParams(validParams)
	invalidTreeParamsMin[NTREES] = strconv.Itoa(MinNTrees - 1)

	invalidTreeParamsMax := copyParams(validParams)
	invalidTreeParamsMax[NTREES] = strconv.Itoa(MaxNTrees + 1)

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{invalidTreeParamsMin, false},
		{invalidTreeParamsMax, false},
	}

	adapter := newANNOYConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("ANNOYConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}

func TestDiskAnnConfAdapter_CheckTrain(t *testing.T) {
	validParams := map[string]string{
		DIM:    strconv.Itoa(128),
		Metric: L2,
	}
	validParamsBigDim := copyParams(validParams)
	validParamsBigDim[DIM] = strconv.Itoa(2048)

	invalidParamsWithoutDim := map[string]string{
		Metric: L2,
	}

	invalidParamsSmallDim := copyParams(validParams)
	invalidParamsSmallDim[DIM] = strconv.Itoa(15)

	cases := []struct {
		params map[string]string
		want   bool
	}{
		{validParams, true},
		{validParamsBigDim, true},
		{invalidParamsWithoutDim, false},
		{invalidParamsSmallDim, false},
	}

	adapter := newDISKANNConfAdapter()
	for _, test := range cases {
		if got := adapter.CheckTrain(test.params); got != test.want {
			t.Errorf("DiskAnnConfAdapter.CheckTrain(%v) = %v", test.params, test.want)
		}
	}
}
