import torch.nn.functional as F


class DistillationLoss:
    def __init__(self, weight=1.0):
        self.weight = weight

    def __call__(self, student, teacher):
        return self.weight * self._smooth_l1(student, self._detach(teacher))

    def _smooth_l1(self, student, teacher):
        if isinstance(student, (list, tuple)):
            losses = [
                F.smooth_l1_loss(student_item, teacher_item)
                for student_item, teacher_item in zip(student, teacher)
            ]
            return sum(losses) / len(losses)
        return F.smooth_l1_loss(student, teacher)

    def _detach(self, teacher):
        if isinstance(teacher, list):
            return [item.detach() for item in teacher]
        if isinstance(teacher, tuple):
            return tuple(item.detach() for item in teacher)
        return teacher.detach()