import warnings
from typing import Optional, Union, TYPE_CHECKING, Callable, List, Dict, Tuple

from functools import wraps

import numpy as np
from collections import defaultdict

from docarray.score import NamedScore

if TYPE_CHECKING:  # pragma: no cover
    from docarray import Document, DocumentArray
    from docarray.array.mixins.embed import CollateFnType
    from docarray.typing import ArrayType, AnyDNN


def _evaluate_deprecation(f):
    """Raises a deprecation warning if the user executes the evaluate function with
    the old interface and adjust the input to fit the new interface."""

    @wraps(f)
    def func(*args, **kwargs):
        if len(args) > 1:
            if not (
                isinstance(args[1], Callable)
                or isinstance(args[1], str)
                or isinstance(args[1], list)
            ):
                kwargs['ground_truth'] = args[1]
                args = [args[0]] + list(args[2:])
                warnings.warn(
                    'The `other` attribute in `evaluate()` is transfered from a '
                    'positional attribute into the keyword attribute `ground_truth`.'
                    'Using it as a positional attribute is deprecated and will be removed '
                    'soon.',
                    DeprecationWarning,
                )
        for old_key, new_key in zip(
            ['other', 'metric', 'metric_name'],
            ['ground_truth', 'metrics', 'metric_names'],
        ):
            if old_key in kwargs:
                kwargs[new_key] = kwargs[old_key]
                warnings.warn(
                    f'`{old_key}` is renamed to `{new_key}` in `evaluate()`, the '
                    f'usage of `{old_key}` is deprecated and will be removed soon.',
                    DeprecationWarning,
                )

        # transfer metrics and metric_names into lists
        list_warning_msg = (
            'The attribute `%s` now accepts a list instead of a '
            'single element. Passing a single element is deprecated and will soon not '
            'be supported anymore.'
        )
        if len(args) > 1:
            if type(args[1]) is str:
                args = list(args)
                args[1] = [args[1]]
                warnings.warn(list_warning_msg % 'metrics', DeprecationWarning)
        for key in ['metrics', 'metric_names']:
            if key in kwargs and type(kwargs[key]) is str:
                kwargs[key] = [kwargs[key]]
                warnings.warn(list_warning_msg % key, DeprecationWarning)
        return f(*args, **kwargs)

    return func


class EvaluationMixin:
    """A mixin that provides ranking evaluation functionality to DocumentArrayLike objects"""

    @_evaluate_deprecation
    def evaluate(
        self,
        metrics: List[Union[str, Callable[..., float]]],
        ground_truth: Optional['DocumentArray'] = None,
        hash_fn: Optional[Callable[['Document'], str]] = None,
        metric_names: Optional[List[str]] = None,
        strict: bool = True,
        label_tag: str = 'label',
        **kwargs,
    ) -> Dict[str, float]:
        """
        Compute ranking evaluation metrics for a given `DocumentArray` when compared
        with a ground truth.

        If one provides a `ground_truth` DocumentArray that is structurally identical
        to `self`, this function compares the `matches` of `documents` inside the
        `DocumentArray` to this `ground_truth`.
        Alternatively, one can directly annotate the documents by adding labels in the
        form of tags with the key specified in the `label_tag` attribute.
        Those tags need to be added to `self` as well as to the documents in the
        matches properties.

        This method will fill the `evaluations` field of Documents inside this
        `DocumentArray` and will return the average of the computations

        :param metrics: List of metric names or metric functions to be computed
        :param ground_truth: The ground_truth `DocumentArray` that the `DocumentArray`
            compares to.
        :param hash_fn: For the evaluation against a `ground_truth` DocumentArray,
            this function is used for generating hashes which are used to compare the
            documents. If not given, ``Document.id`` is used.
        :param metric_names: If provided, the results of the metrics computation will be
            stored in the `evaluations` field of each Document with this names. If not
            provided, the names will be derived from the metric function names.
        :param strict: If set, then left and right sides are required to be fully
            aligned: on the length, and on the semantic of length. These are preventing
            you to evaluate on irrelevant matches accidentally.
        :param label_tag: Specifies the tag which contains the labels.
        :param kwargs: Additional keyword arguments to be passed to the metric
            functions.
        :return: A dictionary which stores for each metric name the average evaluation
            score.
        """
        if len(self) == 0:
            raise ValueError('It is not possible to evaluate an empty DocumentArray')
        if ground_truth and len(ground_truth) > 0 and ground_truth[0].matches:
            ground_truth_type = 'matches'
        elif label_tag in self[0].tags:
            if ground_truth:
                warnings.warn(
                    'A ground_truth attribute is provided but does not '
                    'contain matches. The labels are used instead and '
                    'ground_truth is ignored.'
                )
            ground_truth = self
            ground_truth_type = 'labels'

        else:
            raise RuntimeError(
                'Could not find proper ground truth data. Either labels or the '
                'ground_truth attribute with matches is required'
            )

        if strict:
            self._check_length(len(ground_truth))

        if hash_fn is None:
            hash_fn = lambda d: d.id

        metric_fns = []
        for metric in metrics:
            if callable(metric):
                metric_fns.append(metric)
            elif isinstance(metric, str):
                from docarray.math import evaluation

                metric_fns.append(getattr(evaluation, metric))

        if not metric_names:
            metric_names = [metric_fn.__name__ for metric_fn in metric_fns]

        if len(metric_names) != len(metrics):
            raise ValueError(
                'Could not match metric names to the metrics since the number of '
                'metric names does not match the number of metrics'
            )

        results = defaultdict(list)
        caller_max_rel = kwargs.pop('max_rel', None)
        for d, gd in zip(self, ground_truth):
            max_rel = caller_max_rel or len(gd.matches)
            if strict and hash_fn(d) != hash_fn(gd):
                raise ValueError(
                    f'Document {d} from the left-hand side and '
                    f'{gd} from the right-hand are not hashed to the same value. '
                    f'This means your left and right DocumentArray may not be aligned; or it means your '
                    f'`hash_fn` is badly designed.'
                )
            if not d.matches or not gd.matches:
                raise ValueError(
                    f'Document {d!r} or {gd!r} has no matches, please check your Document'
                )

            targets = gd.matches[:max_rel]

            if ground_truth_type == 'matches':
                desired = {hash_fn(m) for m in targets}
                if len(desired) != len(targets):
                    warnings.warn(
                        f'{hash_fn!r} may not be valid, as it maps multiple Documents into the same hash. '
                        f'Evaluation results may be affected'
                    )
                binary_relevance = [
                    1 if hash_fn(m) in desired else 0 for m in d.matches
                ]
            elif ground_truth_type == 'labels':
                binary_relevance = [
                    1 if m.tags[label_tag] == d.tags[label_tag] else 0 for m in targets
                ]
            else:
                raise RuntimeError(
                    'Could not identify which kind of ground truth'
                    'information is provided to evaluate the matches.'
                )
            for metric_name, metric_fn in zip(metric_names, metric_fns):
                if 'max_rel' in metric_fn.__code__.co_varnames:
                    kwargs['max_rel'] = max_rel
                r = metric_fn(binary_relevance, **kwargs)
                d.evaluations[metric_name] = NamedScore(
                    value=r, op_name=str(metric_fn), ref_id=d.id
                )
                results[metric_name].append(r)
        return {
            metric_name: float(np.mean(values))
            for metric_name, values in results.items()
        }

    def embed_and_evaluate(
        self,
        metrics: List[Union[str, Callable[..., float]]],
        index_data: Optional['DocumentArray'] = None,
        ground_truth: Optional['DocumentArray'] = None,
        metric_names: Optional[str] = None,
        strict: bool = True,
        label_tag: str = 'label',
        embed_models: Optional[Union['AnyDNN', Tuple['AnyDNN', 'AnyDNN']]] = None,
        embed_funcs: Optional[Union[Callable, Tuple[Callable, Callable]]] = None,
        device: str = 'cpu',
        batch_size: Union[int, Tuple[int, int]] = 256,
        collate_fns: Union[
            Optional['CollateFnType'],
            Tuple[Optional['CollateFnType'], Optional['CollateFnType']],
        ] = None,
        distance: Union[
            str, Callable[['ArrayType', 'ArrayType'], 'np.ndarray']
        ] = 'cosine',
        limit: Optional[Union[int, float]] = 20,
        normalization: Optional[Tuple[float, float]] = None,
        exclude_self: bool = False,
        use_scipy: bool = False,
        num_worker: int = 1,
        match_batch_size: int = 100_000,
        query_sample_size: int = 1_000,
        **kwargs,
    ) -> Optional[Union[float, List[float]]]:  # average for each metric
        """
        Computes ranking evaluation metrics for a given `DocumentArray`. This
        function does embedding and matching in the same turn. Thus, you don't need to
        call ``embed`` and ``match`` before it. Instead, it embeds the documents in
        `self` (and `index_data` when provided`) and compute the nearest neighbour
        itself. This might be done in batches for the `index_data` object to reduce
        the memory consumption of the evlauation process. The evaluation itself can be
        done against a `ground_truth` DocumentArray or on the basis of labels like it
        is possible with the :func:``evaluate`` function.

        :param metrics: List of metric names or metric functions to be computed
        :param index_data: The other DocumentArray  to match against, if not given,
            `self` will be matched against itself. This means that every document in
            will be compared to all other documents in `self` to determine the nearest
            neighbors.
        :param ground_truth: The ground_truth `DocumentArray` that the `DocumentArray`
            compares to.
        :param metric_names: If provided, the results of the metrics computation will be
            stored in the `evaluations` field of each Document with this names. If not
            provided, the names will be derived from the metric function names.
        :param strict: If set, then left and right sides are required to be fully
            aligned: on the length, and on the semantic of length. These are preventing
            you to evaluate on irrelevant matches accidentally.
        :param label_tag: Specifies the tag which contains the labels.
        :param embed_models: One or two embedding model written in Keras / Pytorch /
            Paddle for embedding `self` and `index_data`.
        :param embed_funcs: As an alternative to embedding models, custom embedding
            functions can be provided.
        :param device: the computational device for `embed_models`, and the matching
            can be either `cpu` or `cuda`.
        :param batch_size: Number of documents in a batch for embedding.
        :param collate_fns: For each embedding function the respective collate
            function creates a mini-batch of input(s) from the given `DocumentArray`.
            If not provided a default built-in collate_fn uses the `tensors` of the
            documents to create input batches.
        :param distance: The distance metric.
        :param limit: The maximum number of matches, when not given defaults to 20.
        :param normalization: A tuple [a, b] to be used with min-max normalization,
            the min distance will be rescaled to `a`, the max distance will be
            rescaled to `b` all values will be rescaled into range `[a, b]`.
        :param exclude_self: If set, Documents in ``index_data`` with same ``id``
            as the left-hand values will not be considered as matches.
        :param use_scipy: if set, use ``scipy`` as the computation backend. Note,
            ``scipy`` does not support distance on sparse matrix.
        :param num_worker: Specifies the number of workers for the execution of the
            match function.
        :parma match_batch_size: The number of documents which are embedded and
            matched at once. Set this value to a lower value, if you experience high
            memory consumption.
        :param kwargs: Additional keyword arguments to be passed to the metric
            functions.
        :param query_sample_size: For a large number of documents in `self` the
            evaluation becomes infeasible, especially, if `index_data` is large.
            Therefore, queries are sampled if the number of documents in `self` exceeds
            `query_sample_size`. Usually, this has only small impact on the mean metric
            values returned by this function. To prevent sampling, you can set
            `query_sample_size` to None.
        :return: A dictionary which stores for each metric name the average evaluation
            score.
        """

        from docarray import Document, DocumentArray

        if not query_sample_size:
            query_sample_size = len(self)

        query_data = self
        only_one_dataset = not index_data
        apply_sampling = len(self) > query_sample_size

        if only_one_dataset:
            # if the user does not provide a separate set of documents for indexing,
            # the matching is done on the documents itself
            copy_flag = (
                apply_sampling
                or (type(embed_funcs) is tuple)
                or ((embed_funcs is None) and (type(embed_models) is tuple))
            )
            index_data = DocumentArray(self, copy=True) if copy_flag else self

        if apply_sampling:
            rng = np.random.default_rng()
            query_data = DocumentArray(
                rng.choice(self, size=query_sample_size, replace=False)
            )

        if ground_truth and apply_sampling:
            ground_truth = DocumentArray(
                [ground_truth[d.id] for d in query_data if d.id in ground_truth]
            )
            if len(ground_truth) != len(query_data):
                raise ValueError(
                    'The DocumentArray provided in the ground_truth attribute does '
                    'not contain all the documents in self.'
                )

        index_data_labels = None
        if not ground_truth:
            if not label_tag in query_data[0].tags:
                raise ValueError(
                    'Either a ground_truth `DocumentArray` or labels are '
                    'required for the evaluation.'
                )
            if not label_tag in index_data[0].tags:
                raise ValueError(
                    'The `DocumentArray` provided in `index_data` misses ' 'labels.'
                )
            index_data_labels = dict()
            for id_value, tags in zip(index_data[:, 'id'], index_data[:, 'tags']):
                index_data_labels[id_value] = tags[label_tag]

        if embed_funcs is None:
            # derive embed function from embed model
            if embed_models is None:
                raise RuntimeError(
                    'For embedding the documents you need to provide either embedding '
                    'model(s) or embedding function(s)'
                )
            else:
                if type(embed_models) is not tuple:
                    embed_models = (embed_models, embed_models)
                embed_args = [
                    {
                        'embed_model': model,
                        'device': device,
                        'batch_size': batch_size,
                        'collate_fn': collate_fns[i]
                        if type(collate_fns) is tuple
                        else collate_fns,
                    }
                    for i, (model, docs) in enumerate(
                        zip(embed_models, (query_data, index_data))
                    )
                ]
        else:
            if type(embed_funcs) is not tuple:
                embed_funcs = (
                    embed_funcs,
                    embed_funcs,
                )  # use the same embedding function for queries and index

        # embed queries:
        if embed_funcs:
            embed_funcs[0](query_data)
        else:
            query_data.embed(**embed_args[0])

        for doc in query_data:
            doc.matches.clear()

        local_queries = DocumentArray(
            [Document(id=doc.id, embedding=doc.embedding) for doc in query_data]
        )

        def fuse_matches(global_matches: DocumentArray, local_matches: DocumentArray):
            global_matches.extend(local_matches)
            global_matches = sorted(
                global_matches,
                key=lambda x: x.scores[distance].value,
            )[:limit]
            return DocumentArray(global_matches)

        for batch in index_data.batch(match_batch_size):
            if (
                apply_sampling
                or (batch.embeddings is None)
                or (batch[0].embedding[0] == 0)
            ):
                if embed_funcs:
                    embed_funcs[1](batch)
                else:
                    batch.embed(**embed_args[1])

            local_queries.match(
                batch,
                limit=limit,
                metric=distance,
                normalization=normalization,
                exclude_self=exclude_self,
                use_scipy=use_scipy,
                num_worker=num_worker,
                device=device,
                batch_size=int(len(batch) / num_worker) if num_worker > 1 else None,
                only_id=True,
            )

            for doc in local_queries:
                query_data[doc.id, 'matches'] = fuse_matches(
                    query_data[doc.id].matches,
                    doc.matches,
                )

            batch.embeddings = None
        # set labels if necessary
        if not ground_truth:
            for i, doc in enumerate(query_data):
                new_matches = DocumentArray()
                for m in doc.matches:
                    m.tags = {label_tag: index_data_labels[m.id]}
                    new_matches.append(m)
                query_data[doc.id, 'matches'] = new_matches

        metrics_resp = query_data.evaluate(
            ground_truth=ground_truth,
            metrics=metrics,
            metric_names=metric_names,
            strict=strict,
            label_tag=label_tag,
            **kwargs,
        )

        return metrics_resp
