from typing import (
    TYPE_CHECKING,
    TypeVar,
    Sequence,
    List,
    Dict,
    Optional,
    Union,
)

import numpy as np

from docarray import Document, DocumentArray
from docarray.math import ndarray
from docarray.math.helper import EPSILON
from docarray.math.ndarray import to_numpy_array
from docarray.score import NamedScore

if TYPE_CHECKING:
    import tensorflow
    import torch

    WeaviateArrayType = TypeVar(
        'WeaviateArrayType',
        np.ndarray,
        tensorflow.Tensor,
        torch.Tensor,
        Sequence[float],
    )


class FindMixin:
    def _find_similar_vectors(
        self,
        query: 'WeaviateArrayType',
        limit=10,
        filter: Optional[Dict] = None,
        additional: Optional[List] = None,
        sort: Optional[Union[Dict, List]] = None,
        query_params: Optional[Dict] = None,
    ):
        """Returns a subset of documents by the given vector.

        :param query: input supported to be stored in Weaviate. This includes any from the list '[np.ndarray, tensorflow.Tensor, torch.Tensor, Sequence[float]]'
        :param limit: number of retrieved items
        :param filter: the input filter to apply in each stored document
        :param additional: Optional Weaviate flags for meta data
        :param sort: sort parameters performed on matches performed on results
        :param query_params: additional parameters applied to the query outside of the where clause
        :return: a `DocumentArray` containing the `Document` objects that verify the filter.
        """
        query = to_numpy_array(query)
        is_all_zero = np.all(query == 0)
        if is_all_zero:
            query = query + EPSILON

        query_dict = {'vector': query}

        if query_params:
            query_dict.update(query_params)

        _additional = ['id', 'certainty']
        if additional:
            _additional = _additional + additional

        query_builder = (
            self._client.query.get(self._class_name, '_serialized')
            .with_additional(_additional)
            .with_limit(limit)
            .with_near_vector(query_dict)
        )

        if filter:
            query_builder = query_builder.with_where(filter)

        if sort:
            query_builder = query_builder.with_sort(sort)

        results = query_builder.do()

        docs = []
        if 'errors' in results:
            errors = '\n'.join(map(lambda error: error['message'], results['errors']))
            raise ValueError(
                f'find failed, please check your filter query. Errors: \n{errors}'
            )

        found_results = results.get('data', {}).get('Get', {}).get(self._class_name, [])

        # The serialized document is stored in results['data']['Get'][self._class_name]
        for result in found_results:
            doc = Document.from_base64(result['_serialized'], **self._serialize_config)
            certainty = result['_additional']['certainty']

            doc.scores['weaviate_certainty'] = NamedScore(value=certainty)

            if certainty is None:
                doc.scores['cosine_similarity'] = NamedScore(value=None)
            else:
                doc.scores['cosine_similarity'] = NamedScore(value=2 * certainty - 1)

            doc.tags['wid'] = result['_additional']['id']

            if additional:
                for add in additional:
                    doc.tags[f'{add}'] = result['_additional'][add]

            docs.append(doc)

        return DocumentArray(docs)

    def _filter(
        self,
        filter: Dict,
        limit: Optional[Union[int, float]] = 20,
        additional: Optional[List] = None,
        sort: Optional[Union[Dict, List]] = None,
    ) -> 'DocumentArray':
        """Returns a subset of documents by filtering by the given filter (Weaviate `where` filter).

        :param filter: the input filter to apply in each stored document
        :param limit: number of retrieved items
        :param additional: Optional Weaviate flags for meta data
        :param sort: sort parameters performed on matches performed on results
        :return: a `DocumentArray` containing the `Document` objects that verify the filter.
        """
        if not filter:
            return self

        _additional = ['id']
        if additional:
            _additional = _additional + additional

        query_builder = (
            self._client.query.get(self._class_name, '_serialized')
            .with_additional(_additional)
            .with_where(filter)
            .with_limit(limit)
        )

        if sort:
            query_builder = query_builder.with_sort(sort)

        results = query_builder.do()

        docs = []
        if 'errors' in results:
            errors = '\n'.join(map(lambda error: error['message'], results['errors']))
            raise ValueError(
                f'filter failed, please check your filter query. Errors: \n{errors}'
            )

        found_results = results.get('data', {}).get('Get', {}).get(self._class_name, [])

        # The serialized document is stored in results['data']['Get'][self._class_name]
        for result in found_results:
            doc = Document.from_base64(result['_serialized'], **self._serialize_config)

            doc.tags['wid'] = result['_additional']['id']

            if additional:
                for add in additional:
                    doc.tags[f'{add}'] = result['_additional'][add]

            docs.append(doc)

        return DocumentArray(docs)

    def _find(
        self,
        query: 'WeaviateArrayType',
        limit: int = 10,
        filter: Optional[Dict] = None,
        additional: Optional[List] = None,
        sort: Optional[Union[Dict, List]] = None,
        query_params: Optional[Dict] = None,
        **kwargs,
    ) -> List['DocumentArray']:
        """Returns approximate nearest neighbors given a batch of input queries.
        :param query: input supported to be stored in Weaviate. This includes any from the list '[np.ndarray, tensorflow.Tensor, torch.Tensor, Sequence[float]]'
        :param limit: number of retrieved items
        :param filter: filter query used for pre-filtering
        :param additional: Optional Weaviate flags for meta data
        :param sort: sort parameters performed on matches performed on results
        :param query_params: additional parameters applied to the query outside of the where clause

        :return: DocumentArray containing the closest documents to the query if it is a single query, otherwise a list of DocumentArrays containing
           the closest Document objects for each of the queries in `query`.

        Note: Weaviate returns `certainty` values. To get cosine similarities one needs to use `cosine_sim = 2*certainty - 1` as explained here:
                  https://weaviate.io/developers/weaviate/current/more-resources/faq.html#q-how-do-i-get-the-cosine-similarity-from-weaviates-certainty
        """

        num_rows, _ = ndarray.get_array_rows(query)

        if num_rows == 1:
            return [
                self._find_similar_vectors(
                    query,
                    limit=limit,
                    additional=additional,
                    filter=filter,
                    sort=sort,
                    query_params=query_params,
                )
            ]
        else:
            closest_docs = []
            for q in query:
                da = self._find_similar_vectors(
                    q,
                    limit=limit,
                    additional=additional,
                    filter=filter,
                    sort=sort,
                    query_params=query_params,
                )
                closest_docs.append(da)
            return closest_docs
