import torch
from torch.nn import BCELoss
from .....utils.dice import dice_loss


def supervised_loss(probability_map: torch.Tensor, ground_truth: torch.Tensor) -> torch.Tensor:
    """
    Computes the total supervised loss of the batch.
    The loss is consisted of dice loss and cross entropy loss.
    :param probability_map: probability map output of the segmentation model
    :param ground_truth: ground truth of the processed input
    :return: loss of the labeled data
    """

    binary_entropy_loss = BCELoss()
    loss = (binary_entropy_loss(probability_map, ground_truth) + 
            dice_loss(probability_map.unsqueeze(dim=0), ground_truth.long())) / 2

    return loss


def unsupervised_loss(mean_teacher_mixed_up: torch.Tensor, model_mixed_up_output: torch.Tensor) -> torch.Tensor:
    """
    Computes the total unsupervised loss for the batch.
    :param mean_teacher_mixed_up: probability map output of the mean teacher model
    :param model_mixed_up_output: probability map output of the main model using the mixed up data
    :return: loss of the unlabeled data
    """

    loss = torch.mean((mean_teacher_mixed_up - model_mixed_up_output) ** 2)

    return loss


def total_loss(labeled_out: torch.Tensor,
               ground_truth: torch.Tensor,
               unlabeled_mean_out: torch.Tensor,
               unlabeled_model_out: torch.Tensor,
               epoch: int) -> torch.Tensor:
    """
    Computes the total loss consisting of both supervised and unsupervised losses.
    :param labeled_x: probability maps belonging to the labeled dataset
    :param unlabeled_x: probability maps belonging to the unlabeled dataset
    :param ground_truth: ground truth belonging to the labeled dataset
    :param epoch: current epoch of training
    :return: total loss
    """

    sup_loss = supervised_loss(labeled_out, ground_truth)
    if unlabeled_mean_out is None:
        unsup_loss = 0
    else:
        unsup_loss = unsupervised_loss(unlabeled_mean_out, unlabeled_model_out)

    weight = unsupervised_weight(epoch)

    loss = sup_loss * (0.8 - weight) + unsup_loss * (0.2 + weight)
    return loss


def unsupervised_weight(epoch: int, max_epoch=200):
    """
    Computes trade-off weight lambda for the unsupervised loss
    :param epoch: current epoch
    :param w_max: a coefficient used in the formula
    :param max_epoch: maximum number of epochs to run
    :return: trade-off weight lambda
    """

    return 0.6 * (epoch / max_epoch)