# -*- coding: utf-8 -*-
# @Time    : 2022/03/23 15:25
# @Author  : Jianing Wang
# @Email   : lygwjn@gmail.com
# @File    : TripletLoss.py
# !/usr/bin/env python
# coding=utf-8

from enum import Enum
import torch
from torch import nn, Tensor
import torch.nn.functional as F
from transformers.models.bert.modeling_bert import BertModel
from transformers import BertTokenizer, BertConfig

class TripletDistanceMetric(Enum):
    """
    The metric for the triplet loss
    """
    COSINE = lambda x, y: 1 - F.cosine_similarity(x, y)
    EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2)
    MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1)

class TripletLoss(nn.Module):
    """
    This class implements triplet loss. Given a triplet of (anchor, positive, negative),
    the loss minimizes the distance between anchor and positive while it maximizes the distance
    between anchor and negative. It compute the following loss function:

    loss = max(||anchor - positive|| - ||anchor - negative|| + margin, 0).

    Margin is an important hyperparameter and needs to be tuned respectively.

    @:param distance_metric: The distance metric function
    @:param triplet_margin: (float) The margin distance

    Input example of forward function:
        rep_anchor: [[0.2, -0.1, ..., 0.6], [0.2, -0.1, ..., 0.6], ..., [0.2, -0.1, ..., 0.6]]
        rep_candidate: [[0.3, 0.1, ...m -0.3], [-0.8, 1.2, ..., 0.7], ..., [-0.9, 0.1, ..., 0.4]]
        label: [0, 1, ..., 1]

    Return example of forward function:
        0.015 (averged)
        2.672 (sum)

    """
    def __init__(self, distance_metric=TripletDistanceMetric.EUCLIDEAN, triplet_margin: float = 0.5):
        super(TripletLoss, self).__init__()
        self.distance_metric = distance_metric
        self.triplet_margin = triplet_margin


    def forward(self, rep_anchor, rep_positive, rep_negative):
        # rep_anchor: [batch_size, hidden_dim] denotes the representations of anchors
        # rep_positive: [batch_size, hidden_dim] denotes the representations of positive, sometimes, it canbe dropout
        # rep_negative: [batch_size, hidden_dim] denotes the representations of negative
        # label: [batch_size, hidden_dim] denotes the label of each anchor - candidate pair
        distance_pos = self.distance_metric(rep_anchor, rep_positive)
        distance_neg = self.distance_metric(rep_anchor, rep_negative)

        losses = F.relu(distance_pos - distance_neg + self.triplet_margin)
        return losses.mean()


if __name__ == "__main__":
    # configure for huggingface pre-trained language models
    config = BertConfig.from_pretrained("bert-base-cased")
    # tokenizer for huggingface pre-trained language models
    tokenizer = BertTokenizer.from_pretrained("bert-base-cased")
    # pytorch_model.bin for huggingface pre-trained language models
    model = BertModel.from_pretrained("bert-base-cased")
    # obtain two batch of examples, each corresponding example is a pair
    anchor_example = ["I am an anchor, which is the source example sampled from corpora."] # anchor sentence
    positive_example = [
        "I am an anchor, which is the source example.",
        "I am the source example sampled from corpora."
    ] # positive, which randomly dropout or noise from anchor
    negative_example = [
        "It is different with the anchor.",
        "My name is Jianing Wang, please give me some stars, thank you!"
    ] # negative, which randomly sampled from corpora
    # convert each example for feature
    # {"input_ids": xxx, "attention_mask": xxx, "token_tuype_ids": xxx}
    anchor_feature = tokenizer(anchor_example, add_special_tokens=True, padding=True)
    positive_feature = tokenizer(positive_example, add_special_tokens=True, padding=True)
    negative_feature = tokenizer(negative_example, add_special_tokens=True, padding=True)
    # padding and convert to feature batch
    max_seq_lem = 24
    anchor_feature = {key: torch.Tensor([value + [0] * (max_seq_lem - len(value)) for value in values]).long() for key, values in anchor_feature.items()}
    positive_feature = {key: torch.Tensor([value + [0] * (max_seq_lem - len(value)) for value in values]).long() for key, values in positive_feature.items()}
    negative_feature = {key: torch.Tensor([value + [0] * (max_seq_lem - len(value)) for value in values]).long() for key, values in negative_feature.items()}
    # obtain sentence embedding by averaged pooling
    rep_anchor = model(**anchor_feature)[0] # [1, max_seq_len, hidden_dim]
    rep_positive = model(**positive_feature)[0] # [batch_size, max_seq_len, hidden_dim]
    rep_negative = model(**negative_feature)[0] # [batch_size, max_seq_len, hidden_dim]
    # repeat
    rep_anchor = torch.mean(rep_anchor, -1) # [1, hidden_dim]
    rep_positive = torch.mean(rep_positive, -1) # [batch_size, hidden_dim]
    rep_negative = torch.mean(rep_negative, -1) # [batch_size, hidden_dim]
    # obtain contrastive loss
    loss_fn = TripletLoss()
    loss = loss_fn(rep_anchor=rep_anchor, rep_positive=rep_positive, rep_negative=rep_negative)
    print(loss) # tensor(0.5001, grad_fn=<MeanBackward0>)