from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import paddle.fluid as fluid
from utility import get_gpu_num


class NpairsLoss():
    def __init__(self,
                 train_batch_size=160,
                 samples_each_class=2,
                 reg_lambda=0.01):
        self.samples_each_class = samples_each_class
        assert (self.samples_each_class == 2)
        self.train_batch_size = train_batch_size
        num_gpus = get_gpu_num()
        assert (train_batch_size % num_gpus == 0)
        self.cal_loss_batch_size = train_batch_size // num_gpus
        assert (self.cal_loss_batch_size % samples_each_class == 0)
        self.reg_lambda = reg_lambda

    def loss(self, input, label=None):
        reg_lambda = self.reg_lambda
        samples_each_class = self.samples_each_class
        batch_size = self.cal_loss_batch_size
        num_class = batch_size // samples_each_class
        fea_dim = input.shape[1]

        input = fluid.layers.reshape(input, shape=[-1, fea_dim])
        feature = fluid.layers.reshape(
            input, shape=[-1, samples_each_class, fea_dim])
        label = fluid.layers.reshape(label, shape=[-1, samples_each_class])
        label = fluid.layers.cast(label, dtype='float32')
        if samples_each_class == 2:
            anchor_fea, positive_fea = fluid.layers.split(
                feature, num_or_sections=2, dim=1)
            anchor_lab, positive_lab = fluid.layers.split(
                label, num_or_sections=2, dim=1)
        else:
            anchor_fea, positive_fea = fluid.layers.split(
                feature, num_or_sections=[1, samples_each_class - 1], dim=1)
            anchor_lab, positive_lab = fluid.layers.split(
                label, num_or_sections=[1, samples_each_class - 1], dim=1)

        anchor_fea = fluid.layers.reshape(anchor_fea, shape=[-1, fea_dim])
        positive_fea = fluid.layers.reshape(positive_fea, shape=[-1, fea_dim])
        positive_fea_trans = fluid.layers.transpose(positive_fea, perm=[1, 0])
        similarity_matrix = fluid.layers.mul(anchor_fea, positive_fea_trans)

        anchor_lab = fluid.layers.expand(
            x=anchor_lab, expand_times=[1, batch_size - num_class])
        positive_lab_tran = fluid.layers.transpose(positive_lab, perm=[1, 0])
        positive_lab_tran = fluid.layers.expand(
            x=positive_lab_tran, expand_times=[num_class, 1])
        label_remapped = fluid.layers.equal(anchor_lab, positive_lab_tran)
        label_remapped = fluid.layers.cast(
            label_remapped, dtype='float32') / (samples_each_class - 1)
        label_remapped.stop_gradient = True

        out = fluid.layers.softmax(input=similarity_matrix, use_cudnn=False)
        xentloss = fluid.layers.cross_entropy(
            input=out, label=label_remapped, soft_label=True)
        xentloss = fluid.layers.mean(x=xentloss)

        reg = fluid.layers.reduce_mean(
            fluid.layers.reduce_sum(fluid.layers.square(input), dim=1))
        l2loss = 0.5 * reg_lambda * reg
        return xentloss + l2loss
