from typing import List, Tuple, Optional

import numpy as np
import torch.cuda

from autorag import embedding_models
from autorag.evaluation.metric.util import calculate_cosine_similarity
from autorag.nodes.passagefilter.base import passage_filter_node
from autorag.nodes.passagefilter.similarity_threshold_cutoff import embedding_query_content


@passage_filter_node
def similarity_percentile_cutoff(queries: List[str], contents_list: List[List[str]],
                                 scores_list: List[List[float]], ids_list: List[List[str]],
                                 percentile: float, embedding_model: Optional[str] = None,
                                 batch: int = 128,
                                 ) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
    """
    Re-calculate each content's similarity with the query and filter out the contents that are below the content's
    length times percentile. If This is a filter and does not override scores. The output of scores is not coming from
    query-content similarity.
    If the value of content's length times percentile is less than 1, keep the only one highest similarity content.

    :param queries: The list of queries to use for filtering
    :param contents_list: The list of lists of contents to filter
    :param scores_list: The list of lists of scores retrieved
    :param ids_list: The list of lists of ids retrieved
    :param percentile: The percentile to cut off
    :param embedding_model: The embedding model to use for calculating similarity
        Default is OpenAIEmbedding.
    :param batch: The number of queries to be processed in a batch
        Default is 128.
    :return: Tuple of lists containing the filtered contents, ids, and scores
    """
    if embedding_model is None:
        embedding_model = embedding_models['openai']
    else:
        embedding_model = embedding_models[embedding_model]

    query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model,
                                                                   batch)

    results = list(map(lambda x: similarity_percentile_cutoff_pure(x[0], x[1], x[2], x[3], x[4], percentile),
                       zip(query_embeddings, content_embeddings, contents_list, ids_list, scores_list)))

    remain_content_list = list(map(lambda x: x[0], results))
    remain_ids_list = list(map(lambda x: x[1], results))
    remain_scores_list = list(map(lambda x: x[2], results))

    del embedding_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return remain_content_list, remain_ids_list, remain_scores_list


def similarity_percentile_cutoff_pure(query_embedding: str,
                                      content_embeddings: List[List[float]],
                                      content_list: List[str],
                                      ids_list: List[str],
                                      scores_list: List[float],
                                      percentile: float) -> Tuple[List[str], List[str], List[float]]:
    """
    Return tuple of lists containing the filtered contents, ids, and scores

    :param query_embedding: Query embedding
    :param content_embeddings: Each content embedding
    :param content_list: Each content
    :param ids_list: Each id
    :param scores_list: Each score
    :param percentile: The percentile to cut off
    :return: Tuple of lists containing the filtered contents, ids, and scores
    """
    num_top_k = int(len(content_embeddings) * percentile)

    if num_top_k == 0:
        num_top_k = 1

    similarities = np.array(list(map(lambda x: calculate_cosine_similarity(query_embedding, x),
                                     content_embeddings))).tolist()

    content_id_score_similarity = list(zip(ids_list, content_list, scores_list, similarities))

    sorted_content_id_score_similarity = sorted(content_id_score_similarity, key=lambda x: x[3], reverse=True)[
                                         :num_top_k]

    content_result, id_result, score_result, _ = zip(*sorted_content_id_score_similarity)
    return list(content_result), list(id_result), list(score_result)
