from dataclasses import dataclass
from typing import Hashable, Optional, Sequence, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
import torch.distributed as dist
import src.custom_utils as utils
import pandas as pd
import numpy as np
from src.train.gather_utils import all_gather_with_grad
from transformers.modeling_outputs import ModelOutput

logger = utils.get_logger(__name__)


def get_labels_contrastive_loss(
    Nq: int,
    Np: int,
    dtype,
    device,
    label_ids: Sequence[Hashable]
) -> Tuple[Tensor, bool]:
    all_label_ids = [None] * dist.get_world_size()
    dist.all_gather_object(all_label_ids, label_ids)
    all_label_ids = [y for x in all_label_ids for y in x]
    unique_ids = len(set(all_label_ids))
    labels = torch.zeros((Nq, Np), dtype=dtype,
                         device=device)
    if len(all_label_ids) > unique_ids:
        utils.rank0_log(
            f"Found batch with {unique_ids}/{len(all_label_ids)} unique entities.", logger.warning)
        all_label_ids = pd.Series(all_label_ids)
        label2idx = {k: torch.from_numpy(
            np.array(grp.index)) for k, grp in all_label_ids.groupby(by=all_label_ids, sort=False)}
        for i, label_id in enumerate(label_ids):
            labels[i, label2idx[label_id]] = 1
        labels = labels / labels.sum(1, keepdim=True)
    else:
        rank = dist.get_rank()
        rank_start = rank * Nq
        rank_end = rank_start + Nq
        labels[:, rank_start:rank_end] = labels[:,
                                                rank_start:rank_end].fill_diagonal_(1)

    return labels


@dataclass
class ContrastiveLossModelOutput(ModelOutput):
    loss: Optional[Tensor] = None
    loss_logit_pos: Optional[Tensor] = None
    r_at_1: Optional[float] = None
    logit_pos: Optional[float] = None
    logit_neg: Optional[float] = None


class ContrastiveLossModel(nn.Module):

    def get_logits(
        self,
        query_feats: Tensor,
        passage_feats: Tensor,
        logit_scale: Optional[Tensor] = None
    ):
        logits = torch.matmul(
            query_feats, passage_feats.transpose(-1, -2))
        if logit_scale is not None:
            logits = logits * logit_scale
        return logits

    def forward(
        self,
        query_feats: Tensor,
        passage_feats: Tensor,
        logit_scale: Optional[Tensor] = None,
        label_ids: Optional[Sequence[Hashable]] = []
    ) -> ContrastiveLossModelOutput:
        rank = dist.get_rank()
        bsz = query_feats.size(0)
        gather_fn = all_gather_with_grad

        passage_feats = gather_fn(passage_feats)
        Nq, Np = bsz, passage_feats.size(0)

        labels = get_labels_contrastive_loss(
            Nq, Np, query_feats.dtype, query_feats.device, label_ids)
        logits = self.get_logits(query_feats, passage_feats, logit_scale)
        loss = F.cross_entropy(logits, labels)
        
        with torch.no_grad():
            r_at_1 = (labels[:, logits.argmax(1)] != 0).any(-1)
            all_r_at_1 = None
            if rank == 0:
                all_r_at_1 = [torch.zeros_like(r_at_1)
                            for _ in range(dist.get_world_size())]
            dist.gather(r_at_1, gather_list=all_r_at_1, dst=0)
            if rank == 0:
                all_r_at_1 = torch.concat(all_r_at_1)
                r_at_1 = all_r_at_1.to(torch.float16).mean().item()

            local_idxs = torch.arange(Nq, device=logits.device) + rank * Nq
            # logit_pos is an approximation, as it does not account for multiple positives for the same query
            logit_pos = logits.gather(1, local_idxs.unsqueeze(1)).squeeze(1)
            logit_neg = (logits.sum(-1) - logit_pos).sum()
            dist.reduce(logit_pos, dst=0, op=dist.ReduceOp.SUM)
            dist.reduce(logit_neg, dst=0, op=dist.ReduceOp.SUM)
            logit_pos /= (dist.get_world_size() * Nq)
            logit_neg /= (dist.get_world_size() * (logits.numel() - Nq))
            logit_pos = logit_pos.item()
            logit_neg = logit_neg.item()

        if rank:
            r_at_1 = None
            logit_pos = None
            logit_neg = None

        return ContrastiveLossModelOutput(
            loss=loss,
            r_at_1=r_at_1,
            logit_pos=logit_pos,
            logit_neg=logit_neg,
        )
