from typing import List, Optional
import pandas as pd
import torch
import torch.distributed as dist
from tqdm import tqdm
from src.custom_utils import get_logger
import gc

logger = get_logger(__name__)


def similarity_matrix(
    query_features: torch.Tensor,
    passage_features: torch.Tensor,
    batch_size: Optional[int] = None
):
    similarities = []

    bsz = batch_size if batch_size else len(passage_features)

    for i in tqdm(
        range(0, len(query_features), bsz),
        desc=f"Computing a {len(query_features)}x{len(passage_features)} similarity matrix with block size {bsz}..."
    ):
        chunk_query_features = query_features[:bsz]

        row_similarities = []
        for j in range(0, len(passage_features), bsz):
            sim = torch.matmul(chunk_query_features.to(device=0),
                               passage_features[j:j + bsz].transpose(-1, -2).to(device=0))
            row_similarities.append(sim.cpu())

        similarities.append(torch.concat(row_similarities, dim=1))
        query_features = query_features[bsz:]

    return torch.concat(similarities, dim=0)


def similarity_matrix_recall(
    similarities: torch.Tensor,
    labels: Optional[torch.Tensor] = None,
    K=(1, 2, 3, 5, 10, 20, 50, 100, 200, 300),
):
    if labels is None:
        labels = torch.arange(
            len(similarities), dtype=similarities.dtype, device=similarities.device)[:, None]

    predictions = torch.argsort(similarities, dim=1, descending=True)
    del similarities
    gc.collect()

    ret = {}
    N = len(predictions)
    for k in K:
        ret[f"R@{k}"] = (labels.gather(dim=1, index=predictions[:, :k])
                         != 0).any(-1).sum().item() / N
    ret['total_len'] = N
    return ret


def compute_metrics_retrieval(
    eval_preds,
    query_features: List[torch.Tensor],
    passage_features: List[torch.Tensor],
    eval_labels: List[str],
    K=(1, 2, 3, 5, 10, 20, 50, 100, 200, 300),
    batch_size: Optional[int] = None
):
    torch.cuda.empty_cache()
    rank = dist.get_rank()
    all_query_features = [None] * dist.get_world_size()
    all_passage_features = [None] * dist.get_world_size()
    all_eval_labels = [None] * dist.get_world_size()
    dist.gather_object(
        query_features,
        all_query_features if rank == 0 else None,
        dst=0
    )
    dist.gather_object(
        passage_features,
        all_passage_features if rank == 0 else None,
        dst=0
    )
    dist.gather_object(
        eval_labels,
        all_eval_labels if rank == 0 else None,
        dst=0
    )

    if len(query_mask) > 0:
        all_query_mask = [None] * dist.get_world_size()
        dist.gather_object(
            query_mask,
            all_query_mask if rank == 0 else None,
            dst=0
        )
        del query_mask
    else:
        all_query_mask = None

    if len(passage_mask) > 0:
        all_passage_mask = [None] * dist.get_world_size()
        dist.gather_object(
            passage_mask,
            all_passage_mask if rank == 0 else None,
            dst=0
        )
        del passage_mask
    else:
        all_passage_mask = None

    del query_features
    del passage_features
    del eval_labels
    gc.collect()

    ret = {}

    # assuming there are not duplicates in the eval dataset !
    if dist.get_rank() == 0:
        all_query_features = torch.concat(
            [torch.stack(x) for x in all_query_features])
        all_passage_features = torch.concat(
            [torch.stack(x) for x in all_passage_features])
        if all_query_mask is not None:
            all_query_mask = torch.concat(
                [torch.stack(x) for x in all_query_mask], dim=0)
            all_query_mask = all_query_mask.to(all_query_features.dtype)
        if all_passage_mask is not None:
            all_passage_mask = torch.concat(
                [torch.stack(x) for x in all_passage_mask], dim=0)
            all_passage_mask = all_passage_mask.to(all_query_features.dtype)
        Nq = len(all_query_features)
        Np = len(all_passage_features)
        logger.info(
            f"Computing a {Nq}x{Np} similarity matrix...")

        similarities = similarity_matrix(
            query_features=all_query_features, passage_features=all_passage_features, batch_size=batch_size)

        del all_query_features
        del all_passage_features
        if all_query_mask is not None:
            del all_query_mask
        if all_passage_mask is not None:
            del all_passage_mask
        gc.collect()

        all_eval_labels = [y for x in all_eval_labels for y in x]
        labels = torch.zeros((Nq, Np), dtype=similarities.dtype)
        all_eval_labels = pd.Series(all_eval_labels)
        pid2idxs = {k: grp.index for k, grp in all_eval_labels.groupby(
            all_eval_labels, sort=False)}
        for i, pid in enumerate(all_eval_labels):
            labels[i, pid2idxs[pid]] = 1
        labels = labels / labels.sum(dim=1, keepdim=True)
        ret = similarity_matrix_recall(
            similarities=similarities, labels=labels, K=K)

    return ret
