# Copyright (c) ModelScope Contributors. All rights reserved.
import numpy as np
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from accelerate.utils import gather_object
from enum import Enum
from torch import nn
from torch.nn import MSELoss
from transformers.utils import strtobool

from swift.sequence_parallel import sequence_parallel
from swift.utils import get_dist_setting
from .base import BaseLoss


# Code borrowed from sentence_transformers
class SiameseDistanceMetric(Enum):
    """The metric for the contrastive loss"""

    EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2)  # noqa
    MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1)  # noqa
    COSINE_DISTANCE = lambda x, y: 1 - F.cosine_similarity(x, y)  # noqa


def _parse_pair_sentence(outputs):
    if isinstance(outputs, dict):
        last_hidden_state = outputs['last_hidden_state']
    else:
        last_hidden_state = outputs
    batch_size = last_hidden_state.shape[0]
    shape_len = len(last_hidden_state.shape)
    first_sentence = list(range(0, batch_size, 2))
    second_sentence = list(range(1, batch_size, 2))
    if shape_len == 3:
        sentence1 = last_hidden_state[first_sentence][:, 0].squeeze(dim=1)
        sentence2 = last_hidden_state[second_sentence][:, 0].squeeze(dim=1)
    else:
        sentence1 = last_hidden_state[first_sentence]
        sentence2 = last_hidden_state[second_sentence]
    return sentence1, sentence2


class CosineSimilarityLoss(BaseLoss):

    def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
        # You need to return a scalar representing the loss.
        cos_score_transformation = nn.Identity()
        loss_fct = MSELoss()
        sentence1, sentence2 = _parse_pair_sentence(outputs)
        output = cos_score_transformation(torch.cosine_similarity(sentence1, sentence2))
        return loss_fct(output, labels.to(output.dtype).view(-1))


class ContrastiveLoss(BaseLoss):

    def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
        sentence1, sentence2 = _parse_pair_sentence(outputs)
        distance_metric = SiameseDistanceMetric.COSINE_DISTANCE
        distances = distance_metric(sentence1, sentence2)
        margin = 0.5
        labels = labels.to(sentence1.dtype)
        losses = 0.5 * (labels * distances.pow(2) + (1 - labels) * F.relu(margin - distances).pow(2))
        return losses.mean()


class OnlineContrastiveLoss(BaseLoss):

    def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
        sentence1, sentence2 = _parse_pair_sentence(outputs)
        distance_metric = SiameseDistanceMetric.COSINE_DISTANCE
        distance_matrix = distance_metric(sentence1, sentence2)
        negs = distance_matrix[labels == 0]
        poss = distance_matrix[labels == 1]

        # select hard positive and hard negative pairs
        negative_pairs = negs[negs < (poss.max() if len(poss) > 1 else negs.mean())]
        positive_pairs = poss[poss > (negs.min() if len(negs) > 1 else poss.mean())]

        positive_loss = positive_pairs.pow(2).sum()
        margin = 0.5
        negative_loss = F.relu(margin - negative_pairs).pow(2).sum()
        loss = positive_loss + negative_loss
        return loss


def _parse_multi_negative_sentences(sentences, labels, hard_negatives=None):
    split_indices = torch.nonzero(labels, as_tuple=False).squeeze().tolist()
    if isinstance(split_indices, int):
        split_indices = [split_indices]
    split_indices.append(len(labels))
    split_indices = np.array(split_indices) + np.array(list(range(len(split_indices))))
    split_tensors = []

    for i in range(len(split_indices) - 1):
        start = split_indices[i]
        end = split_indices[i + 1]
        split_part = sentences[start:end]
        if hard_negatives is not None:
            negatives = len(split_part) - 2
            assert negatives > 0
            if negatives > hard_negatives:
                split_part = split_part[:hard_negatives + 2]
            elif negatives < hard_negatives:
                selected = np.random.choice(list(range(negatives)), size=hard_negatives - negatives, replace=True)
                selected += 1  # skip positive
                split_part = torch.cat((split_part, split_part[selected]), dim=0)
        split_tensors.append(split_part)
    return split_tensors


class InfonceLoss(BaseLoss):

    def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
        temperature = float(os.environ.get('INFONCE_TEMPERATURE', '0.1'))  # temperature
        # calculate CE across the batch, meaning all samples will be negative except the matching positive
        use_batch = strtobool(os.environ.get('INFONCE_USE_BATCH', 'True'))
        hard_negatives = os.environ.get('INFONCE_HARD_NEGATIVES', None)  # how many negative prompts kept in one sample
        # mask out fake negatives
        infonce_mask_fake_negative = strtobool(os.environ.get('INFONCE_MASK_FAKE_NEGATIVE', 'False'))
        fake_neg_margin = float(os.environ.get('INFONCE_FAKE_NEG_MARGIN', '0.1'))
        # enhanced components to align with Qwen3-Embedding denominator; controlled individually
        # defaults set to False for backward compatibility
        infonce_include_qq = strtobool(os.environ.get('INFONCE_INCLUDE_QQ', 'False'))
        infonce_include_dd = strtobool(os.environ.get('INFONCE_INCLUDE_DD', 'False'))
        if hard_negatives is not None:
            hard_negatives = int(hard_negatives)
        if self.is_megatron:
            from megatron.core import mpu
            rank, world_size = mpu.get_data_parallel_rank(), mpu.get_data_parallel_world_size()
        else:
            rank, _, world_size, _ = get_dist_setting()
        # repeat of anchor(1)+positive(1)+negatives(n)
        sentences = outputs['last_hidden_state']

        if world_size > 1 and use_batch:
            if getattr(sequence_parallel, 'dp_group', None) is not None:
                all_sentences = sequence_parallel._gather_object_dp(sentences.unsqueeze(0))
                labels = sequence_parallel._gather_object_dp(labels)
                rank = sequence_parallel.dp_rank
            elif self.is_megatron:
                from megatron.core import mpu
                dp_group = mpu.get_data_parallel_group()
                shapes = [sentences.new_empty((2, ), dtype=torch.long) for _ in range(world_size)]
                dist.all_gather(
                    shapes,
                    sentences.new_tensor(sentences.shape, dtype=torch.long),
                    group=dp_group,
                )
                all_sentences = [sentences.new_empty(shape.tolist()) for shape in shapes]
                dist.all_gather(
                    all_sentences,
                    sentences,
                    group=dp_group,
                )
            else:
                # gather all the sentences and labels across the gpus when calculate loss across all batches of all gpus
                all_sentences = gather_object(sentences.unsqueeze(0))
                labels = gather_object(labels)
            # override the gathered one
            all_sentences[rank] = sentences
            for idx in range(len(all_sentences)):
                if idx == rank:
                    continue
                # we don't calculate grad from other gpus
                all_sentences[idx] = all_sentences[idx].detach().to(sentences.device)
            sentences = torch.cat(all_sentences, dim=0)
            labels = [tensor.to(sentences.device) for tensor in labels]
            labels = torch.stack(labels, dim=0)

        # split tensors into single sample
        # for example: batch_size=2 with tensor anchor(1)+positive(1)+negatives(3) + anchor(1)+positive(1)+negatives(2)
        # labels will be [1,0,0,0,1,0,0], meaning 1 positive, 3 negatives, 1 positive, 2 negatives
        split_tensors = _parse_multi_negative_sentences(sentences, labels, hard_negatives)
        loss = 0
        can_batched = hard_negatives is not None
        if hard_negatives is None and len(set([s.shape[0] for s in split_tensors])) == 1:
            # all tensors have the same batch size
            can_batched = True
        if not use_batch:
            # only calculate loss inside one sample
            if can_batched:
                # negative numbers are equal
                # [B, neg+2, D]
                sentences = torch.stack(split_tensors, dim=0)
                # [B, 1, D] * [B, neg+1, D]
                similarity_matrix = torch.matmul(sentences[:, 0:1], sentences[:, 1:].transpose(1, 2)) / temperature
                # The positive one is the first element
                labels = torch.zeros(len(split_tensors), dtype=torch.int64).to(sentences.device)
                loss = nn.CrossEntropyLoss()(similarity_matrix.squeeze(1), labels)
            else:
                # the negative numbers may be different, use for loop
                for tensor in split_tensors:
                    # [D] * [neg+1, D]
                    similarity_matrix = torch.matmul(tensor[0], tensor[1:].T) / temperature
                    # The positive one is the first element
                    labels = torch.tensor(0).to(tensor.device)
                    loss += nn.CrossEntropyLoss()(similarity_matrix, labels)
                # avg between all batches in one gpu
                loss /= len(split_tensors)
        else:
            if can_batched:
                # [B, neg+2, D]
                sentences = torch.stack(split_tensors, dim=0)
                # base q->d similarities (includes own positive and all in-batch documents)
                queries = sentences[:, 0].squeeze(1)  # [B, D]
                docs_all = sentences[:, 1:].reshape(-1, sentences.size(2))  # [B*(neg+1), D]
                qd_matrix = torch.matmul(queries, docs_all.T)  # [B, B*(neg+1)]
                # target indices: start of each group's document block (its positive)
                labels = torch.tensor(range(0,
                                            sentences.size(0) * (sentences.size(1) - 1),
                                            sentences.size(1) - 1)).view(-1).to(sentences.device)

                logits_list = [qd_matrix]

                if infonce_include_qq:
                    # q->q similarities; exclude self via -inf on diagonal to avoid accidental positives
                    qq_matrix = torch.matmul(queries, queries.T)  # [B, B]
                    qq_matrix = qq_matrix.clone()
                    qq_matrix.fill_diagonal_(float('-inf'))
                    logits_list.append(qq_matrix)

                if infonce_include_dd:
                    # d+ -> d (doc-doc) similarities; exclude self-positive column per row
                    pos_docs = sentences[:, 1].squeeze(1)  # [B, D]
                    dd_matrix = torch.matmul(pos_docs, docs_all.T)  # [B, B*(neg+1)]
                    # mask self positive per row: column index = row_idx * (neg+1)
                    block = sentences.size(1) - 1  # (neg+1)
                    if block > 0:
                        row_idx = torch.arange(dd_matrix.size(0), device=dd_matrix.device)
                        col_idx = row_idx * block
                        dd_matrix[row_idx, col_idx] = float('-inf')
                    logits_list.append(dd_matrix)

                if infonce_mask_fake_negative:
                    # thresholds derived from positive q->d scores per row
                    row_idx = torch.arange(qd_matrix.size(0), device=qd_matrix.device)
                    pos_scores = qd_matrix[row_idx, labels]
                    thresholds = pos_scores.view(-1, 1).detach() + fake_neg_margin

                    # qd block mask
                    qd_block = qd_matrix.clone()
                    qd_mask = qd_block > thresholds
                    qd_block[qd_mask] = float('-inf')

                    components = [qd_block]

                    # qq block mask (if present)
                    if infonce_include_qq:
                        qq_block = qq_matrix.clone()
                        qq_mask = qq_block > thresholds
                        qq_block[qq_mask] = float('-inf')
                        # diagonal already masked unconditionally at construction time
                        components.append(qq_block)

                    # dd block (if present): self-positive column already masked unconditionally
                    if infonce_include_dd:
                        # align with Qwen3-Embedding, no threshold masking for d-d
                        components.append(dd_matrix)

                    similarity_matrix = torch.cat(components, dim=1)
                else:
                    # concatenate all components without masking
                    similarity_matrix = torch.cat(logits_list, dim=1)
                # temperature scaling and CE
                similarity_matrix = similarity_matrix / temperature
                loss = nn.CrossEntropyLoss()(similarity_matrix, labels)
            else:
                all_tensors = []
                for tensor in split_tensors:
                    all_tensors.append(tensor[1:])
                # cat all neg+1 tensors
                sentences = torch.cat(all_tensors, dim=0)
                # prepare query anchors list if q-q is included
                if infonce_include_qq:
                    queries_all = torch.stack([t[0] for t in split_tensors], dim=0)  # [B, D]
                length = 0
                for idx, tensor in enumerate(split_tensors):
                    # [D] * [B*(neg+1), D], neg numbers are different
                    qd_vec = torch.matmul(tensor[0], sentences.T)
                    target = torch.tensor(length).to(tensor.device)
                    logits_parts = []

                    # compute threshold from positive q->d score
                    threshold = (qd_vec[target].detach() + fake_neg_margin)

                    # qd part with masking
                    if infonce_mask_fake_negative:
                        qd_masked = torch.where(qd_vec > threshold, torch.tensor(float('-inf'), device=qd_vec.device),
                                                qd_vec)
                    else:
                        qd_masked = qd_vec
                    logits_parts.append(qd_masked)

                    # qq part
                    if infonce_include_qq:
                        qq_vec = torch.matmul(tensor[0], queries_all.T)  # [B]
                        # exclude self
                        qq_vec = qq_vec.clone()
                        qq_vec[idx] = float('-inf')
                        if infonce_mask_fake_negative:
                            qq_vec = torch.where(qq_vec > threshold, torch.tensor(float('-inf'), device=qq_vec.device),
                                                 qq_vec)
                        logits_parts.append(qq_vec)

                    # dd part
                    if infonce_include_dd:
                        dd_vec = torch.matmul(tensor[1], sentences.T)  # [B*(neg+1)]
                        # mask self positive column for this row only (no threshold masking for d-d)
                        block = split_tensors[idx].size(0) - 1  # (neg+1) for this group
                        dd_vec[length] = float('-inf')
                        logits_parts.append(dd_vec)

                    logits_row = torch.cat(logits_parts, dim=-1)
                    logits_row = logits_row / temperature
                    loss += nn.CrossEntropyLoss()(logits_row.unsqueeze(0), target.unsqueeze(0))
                    # next positive is neg+1
                    length += tensor.size(0) - 1
                loss /= len(split_tensors)
        return loss