# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

import torch
from mmcls.models.builder import CLASSIFIERS
from mmcls.models.classifiers.image import ImageClassifier
from torch.nn.functional import sigmoid, softmax


@CLASSIFIERS.register_module()
class SupConClassifier(ImageClassifier):
    def __init__(self, backbone, neck=None, head=None, pretrained=None, **kwargs):
        self.multilabel = kwargs.pop("multilabel", False)
        self.hierarchical = kwargs.pop("hierarchical", False)
        super().__init__(backbone, neck=neck, head=head, pretrained=pretrained, **kwargs)

    def forward_train(self, img, gt_label, **kwargs):
        # concatenate the different image views along the batch size
        if len(img.shape) == 5:
            img = torch.cat([img[:, d, :, :, :] for d in range(img.shape[1])], dim=0)
        x = self.extract_feat(img)
        losses = dict()
        if self.multilabel or self.hierarchical:
            loss = self.head.forward_train(x, gt_label, **kwargs)
        else:
            gt_label = gt_label.squeeze(dim=1)
            loss = self.head.forward_train(x, gt_label)
        losses.update(loss)
        return losses
