# 파일 위치 예시: mmyolo/models/losses/cross_entropy_loss_mmyolo.py

# -----------------------------------------------------
# Copyright (c) OpenMMLab. All rights reserved.
# -----------------------------------------------------
import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F

# 1) *** mmdet.registry가 아니라, mmyolo에서 사용하는 registry를 import ***
#    mmyolo의 registry.py 혹은 __init__.py 등에서 MODELS를 가져오는 방식 확인
#    가정: mmyolo/registry.py 파일에 MODELS = Registry('model', scope='mmyolo') 등의 코드가 있다고 가정
#    (프로젝트 구조에 따라 다를 수 있습니다.)
from mmyolo.registry import MODELS  

# accuracy, utils 같은 부분은 그대로 두되,
# "from .accuracy_mmyolo import accuracy" 와 같은 식으로 상대 import 처리
from .accuracy_mmyolo import accuracy
from .utils_mmyolo import weight_reduce_loss


def cross_entropy(pred,
                  label,
                  weight=None,
                  reduction='mean',
                  avg_factor=None,
                  class_weight=None,
                  ignore_index=-100,
                  avg_non_ignore=False):
    """Cross Entropy 계산함수.
    (본문 그대로)
    """
    ignore_index = -100 if ignore_index is None else ignore_index
    loss = F.cross_entropy(
        pred,
        label,
        weight=class_weight,
        reduction='none',
        ignore_index=ignore_index
    )
    if (avg_factor is None) and avg_non_ignore and reduction == 'mean':
        avg_factor = label.numel() - (label == ignore_index).sum().item()
    if weight is not None:
        weight = weight.float()
    loss = weight_reduce_loss(loss, weight=weight, reduction=reduction, avg_factor=avg_factor)
    return loss


def binary_cross_entropy(pred,
                         label,
                         weight=None,
                         reduction='mean',
                         avg_factor=None,
                         class_weight=None,
                         ignore_index=-100,
                         avg_non_ignore=False):
    """Binary Cross Entropy 계산함수.
    (본문 그대로)
    """
    ignore_index = -100 if ignore_index is None else ignore_index

    if pred.dim() != label.dim():
        # pred의 채널 수가 다르다면 one-hot 확장
        bin_labels, weight, valid_mask = _expand_onehot_labels(
            label, weight, pred.size(-1), ignore_index
        )
    else:
        valid_mask = ((label >= 0) & (label != ignore_index)).float()
        if weight is not None:
            weight = weight * valid_mask
        else:
            weight = valid_mask

    if (avg_factor is None) and avg_non_ignore and reduction == 'mean':
        avg_factor = valid_mask.sum().item()

    weight = weight.float()
    loss = F.binary_cross_entropy_with_logits(
        pred, label.float(), pos_weight=class_weight, reduction='none'
    )
    loss = weight_reduce_loss(loss, weight, reduction=reduction, avg_factor=avg_factor)
    return loss


def _expand_onehot_labels(labels, label_weights, label_channels, ignore_index):
    """label을 one-hot으로 확장해주는 함수. (본문 그대로)
    """
    bin_labels = labels.new_full((labels.size(0), label_channels), 0)
    valid_mask = (labels >= 0) & (labels != ignore_index)
    inds = torch.nonzero(valid_mask & (labels < label_channels), as_tuple=False)

    if inds.numel() > 0:
        bin_labels[inds, labels[inds]] = 1

    valid_mask = valid_mask.view(-1, 1).expand(labels.size(0), label_channels).float()
    if label_weights is None:
        bin_label_weights = valid_mask
    else:
        bin_label_weights = label_weights.view(-1, 1).repeat(1, label_channels)
        bin_label_weights *= valid_mask

    return bin_labels, bin_label_weights, valid_mask


def mask_cross_entropy(pred,
                       target,
                       label,
                       reduction='mean',
                       avg_factor=None,
                       class_weight=None,
                       ignore_index=None,
                       **kwargs):
    """Mask Cross Entropy 계산함수. (본문 그대로)
    """
    assert ignore_index is None, 'BCE loss does not support ignore_index'
    num_rois = pred.size()[0]
    inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device)
    pred_slice = pred[inds, label].squeeze(1)
    return F.binary_cross_entropy_with_logits(
        pred_slice, target, weight=class_weight, reduction='mean')[None]


@MODELS.register_module()
class CrossEntropyLoss_mmyolo(nn.Module):
    """mmyolo용 기본 CrossEntropyLoss 커스텀 버전.

    Args:
        use_sigmoid (bool): sigmoid 사용 여부
        use_mask (bool): mask cross entropy 사용 여부
        reduction (str): 손실값 축소 방식(none, mean, sum)
        class_weight (list[float] | None): 클래스별 가중치
        ignore_index (int | None): 무시할 라벨
        loss_weight (float): 전체 손실에 곱할 비율
        avg_non_ignore (bool): 무시할 라벨을 제외하고 평균 낼지 여부
    """

    def __init__(self,
                 use_sigmoid=False,
                 use_mask=False,
                 reduction='mean',
                 class_weight=None,
                 ignore_index=None,
                 loss_weight=1.0,
                 avg_non_ignore=False):
        super().__init__()
        assert (use_sigmoid is False) or (use_mask is False), \
            'use_sigmoid와 use_mask는 동시에 True가 될 수 없습니다.'
        self.use_sigmoid = use_sigmoid
        self.use_mask = use_mask
        self.reduction = reduction
        self.loss_weight = loss_weight
        self.class_weight = class_weight
        self.ignore_index = ignore_index
        self.avg_non_ignore = avg_non_ignore
        if ((ignore_index is not None) and not self.avg_non_ignore and self.reduction == 'mean'):
            warnings.warn(
                'ignore_index가 설정되었지만 avg_non_ignore=False이므로, '
                'PyTorch CrossEntropy와 동일하게 무시 라벨은 전체 평균에 포함되지 않습니다.'
            )

        if self.use_sigmoid:
            self.cls_criterion = binary_cross_entropy
        elif self.use_mask:
            self.cls_criterion = mask_cross_entropy
        else:
            self.cls_criterion = cross_entropy

    def extra_repr(self):
        return f'avg_non_ignore={self.avg_non_ignore}'

    def forward(self,
                cls_score,
                label,
                weight=None,
                avg_factor=None,
                reduction_override=None,
                ignore_index=None,
                **kwargs):
        reduction = reduction_override if reduction_override else self.reduction
        if ignore_index is None:
            ignore_index = self.ignore_index

        # class_weight가 있을 경우 텐서로 변환
        if self.class_weight is not None:
            class_weight = cls_score.new_tensor(self.class_weight)
        else:
            class_weight = None

        loss_cls = self.loss_weight * self.cls_criterion(
            cls_score,
            label,
            weight,
            class_weight=class_weight,
            reduction=reduction,
            avg_factor=avg_factor,
            ignore_index=ignore_index,
            avg_non_ignore=self.avg_non_ignore,
            **kwargs
        )
        return loss_cls


@MODELS.register_module()
class CrossEntropyCustomLoss_mmyolo(CrossEntropyLoss_mmyolo):
    """mmyolo용 사용자 커스텀 CrossEntropyLoss.
    
    - YOLO 시리즈(특히 v8 계열)에서 추가로 필요한 메서드들을 오버라이드하여
      `num_classes`, `get_cls_channels`, `get_activation`, `get_accuracy` 등을 설정.

    Args:
        use_sigmoid (bool): sigmoid 사용 여부
        use_mask (bool): mask cross entropy 사용 여부
        reduction (str): 손실값 축소 방식(none, mean, sum)
        num_classes (int): 클래스 개수 (YOLO 등에서 명시적으로 필요)
        class_weight (list[float] | None): 클래스별 가중치
        ignore_index (int | None): 무시할 라벨
        loss_weight (float): 전체 손실에 곱할 비율
        avg_non_ignore (bool): 무시할 라벨을 제외하고 평균 낼지 여부
    """

    def __init__(self,
                 use_sigmoid=False,
                 use_mask=False,
                 reduction='mean',
                 num_classes=-1,
                 class_weight=None,
                 ignore_index=None,
                 loss_weight=1.0,
                 avg_non_ignore=False):
        super().__init__(
            use_sigmoid=use_sigmoid,
            use_mask=use_mask,
            reduction=reduction,
            class_weight=class_weight,
            ignore_index=ignore_index,
            loss_weight=loss_weight,
            avg_non_ignore=avg_non_ignore
        )
        self.num_classes = num_classes
        assert self.num_classes != -1, \
            'num_classes는 반드시 지정되어야 합니다.'

        # 아래 세 속성은 커스텀 헤드에서 사용됨.
        self.custom_cls_channels = True
        self.custom_activation = True
        self.custom_accuracy = True

    def get_cls_channels(self, num_classes):
        """헤드에서 최종 cls_score의 channel 수를 지정."""
        assert num_classes == self.num_classes
        # 일반 softmax면 배경을 +1 해서 반환하기도 하지만,
        # num_classes + 1 반환 여부는 필요에 따라 수정.
        # 여기서는 self.num_classes + 1 형태로 가정
        if not self.use_sigmoid:
            return num_classes + 1
        else:
            return num_classes

    def get_activation(self, cls_score):
        """추론 시점에 활성화 함수를 적용한 최종 scores를 반환."""
        fine_cls_score = cls_score[:, :self.num_classes]
        if not self.use_sigmoid:
            bg_score = cls_score[:, [-1]]
            new_score = torch.cat([fine_cls_score, bg_score], dim=-1)
            scores = F.softmax(new_score, dim=-1)
        else:
            score_classes = fine_cls_score.sigmoid()
            score_neg = 1 - score_classes.sum(dim=1, keepdim=True)
            score_neg = score_neg.clamp(min=0, max=1)
            scores = torch.cat([score_classes, score_neg], dim=1)
        return scores

    def get_accuracy(self, cls_score, labels):
        """정확도 계산. 보통 pos_inds = labels < num_classes로 필터링."""
        fine_cls_score = cls_score[:, :self.num_classes]
        pos_inds = labels < self.num_classes
        acc_classes = accuracy(fine_cls_score[pos_inds], labels[pos_inds])
        acc = dict()
        acc['acc_classes'] = acc_classes
        return acc