# Copyright 2021 The TensorFlow Authors
#
# 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.

import os
from pathlib import Path
import tempfile
from typing import List, Sequence, Tuple, Union

import nmslib
import tensorflow as tf

from .search import Search
from tensorflow_similarity.distances import Distance, distance_canonicalizer
from tensorflow_similarity.types import FloatTensor


class NMSLibSearch(Search):
    """
    Efficiently find nearest embeddings by indexing known embeddings and make
    them searchable using the  [Approximate Nearest Neigboors Search](https://en.wikipedia.org/wiki/Nearest_neighbor_search)
    search library [NMSLIB](https://github.com/nmslib/nmslib).
    """

    def __init__(
        self,
        distance: Union[Distance, str],
        dims: int,
        algorithm: str = "nmslib_hnsw",
        **kwargs
    ):

        distance_obj: Distance = distance_canonicalizer(distance)

        # convert to nmslib word
        if distance_obj.name == "cosine":
            space = "cosinesimil"
        elif distance_obj.name == "euclidean":
            space = "l2"
        elif distance_obj.name == "manhattan":
            space = "l1"
        else:
            raise ValueError("Unsupported metric space")

        if algorithm == "nmslib_hnsw":
            method = "hnsw"
        else:
            raise ValueError("Unsupported algorithm")

        self._search_index = nmslib.init(method=method, space=space)
        self._search_index.createIndex()

    def add(
        self,
        embedding: FloatTensor,
        idx: int,
        verbose: int = 1,
        build: bool = True,
        **kwargs
    ):
        """Add an embedding to the index

        Args:
            embedding: The embedding to index as computed by
            the similarity model.

            idx: Embedding id as in the index table.
            Returned with the embedding to allow to lookup
            the data associated with a given embedding.

            verbose: Be verbose. Defaults to 1.

            build: Rebuild the index after the addition.
            Required to make the embedding searchable.
            Set to false to save time between successive addition.
            Defaults to True.

        """
        self._search_index.addDataPoint(idx, embedding)
        if build:
            self._build(verbose=verbose)

    def batch_add(
        self,
        embeddings: FloatTensor,
        idxs: Sequence[int],
        verbose: int = 1,
        build: bool = True,
        **kwargs
    ):
        """Add a batch of embeddings to the search index.

        Args:
            embeddings: List of embeddings to add to the index.

            idxs (int): Embedding ids as in the index table. Returned with
            the embeddings to allow to lookup the data associated
            with the returned embeddings.

            verbose: Be verbose. Defaults to 1.

            build: Rebuild the index after the addition. Required to
            make the embeddings searchable. Set to false to save
            time between successive addition. Defaults to True.
        """
        # !addDataPoint and addDataPointBAtch have inverted parameters
        if verbose:
            print("|-Adding embeddings to index.")
        self._search_index.addDataPointBatch(embeddings, idxs)

        if build:
            if verbose:
                print("|-Building index.")
            self._build(verbose=verbose)

    def lookup(
        self, embedding: FloatTensor, k: int = 5
    ) -> Tuple[List[int], List[float]]:
        """Find embedding K nearest neighboors embeddings.

        Args:
            embedding: Query embedding as predicted by the model.
            k: Number of nearest neighboors embedding to lookup. Defaults to 5.
        """
        idxs: List[int] = []
        distances: List[float] = []
        idxs, distances = self._search_index.knnQuery(embedding, k=k)
        return idxs, distances

    def batch_lookup(
        self, embeddings: FloatTensor, k: int = 5
    ) -> Tuple[List[List[int]], List[List[float]]]:
        """Find embeddings K nearest neighboors embeddings.

        Args:
            embedding: Batch of query embeddings as predicted by the model.
            k: Number of nearest neighboors embedding to lookup. Defaults to 5.
        """
        batch_idxs = []
        batch_distances = []

        nn = self._search_index.knnQueryBatch(embeddings, k=k)
        for n in nn:
            batch_idxs.append(n[0])
            batch_distances.append(n[1])
        return batch_idxs, batch_distances

    def save(self, path: str):
        """Serializes the index data on disk

        Args:
            path: where to store the data
        """
        with tempfile.TemporaryDirectory() as tmpdirname:
            fname = self.__make_fname(path)
            tmpidx = os.path.join(tmpdirname, os.path.basename(path))
            self._search_index.saveIndex(tmpidx, save_data=True)

            # read search_index.bin tmp file and write to fname
            with tf.io.gfile.GFile(fname, "w+b") as gfp:
                with open(tmpidx, "rb") as fp:
                    gfp.write(fp.read())

            # read search_index.bin.dat tmp file and write to fname + .dat
            with tf.io.gfile.GFile(fname + ".dat", "w+b") as gfp:
                tmpdat = os.path.join(
                    tmpdirname, os.path.basename(path) + ".dat"
                )
                with open(tmpdat, "rb") as fp:
                    gfp.write(fp.read())

    def load(self, path: str):
        """load index on disk

        Args:
            path: where to store the data
        """
        with tempfile.TemporaryDirectory() as tmpdirname:
            fname = self.__make_fname(path)
            tmpidx = os.path.join(tmpdirname, os.path.basename(path))

            # read fname and write to tmpidx
            with tf.io.gfile.GFile(fname, "rb") as gfp:
                with open(tmpidx, "w+b") as fp:
                    fp.write(gfp.read())

            # read fname + .dat and write to tmpdat
            with tf.io.gfile.GFile(fname + ".dat", "rb") as gfp:
                tmpdat = os.path.join(
                    tmpdirname, os.path.basename(path) + ".dat"
                )
                with open(tmpdat, "w+b") as fp:
                    fp.write(gfp.read())

            self._search_index.loadIndex(tmpidx, load_data=True)

    def _build(self, verbose=0):
        """Build the index this is need to take into account the new points"""
        show = True if verbose else False
        self._search_index.createIndex(print_progress=show)

    def __make_fname(self, path):
        return str(Path(path) / "search_index.bin")
