import torch
import numpy as np
from typing import Tuple
from torch.nn import BCELoss, KLDivLoss
from .....utils.dice import dice_loss


def single_supervised_loss(probability_map: torch.Tensor, ground_truth: torch.Tensor) -> torch.Tensor:
    """
    Computes the supervised loss of a single scale of the model.
    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: overall loss of the supervised 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 supervised_loss(probability_map: list, 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 outputs of the segmentation model from different scales
    :param ground_truth: ground truth of the processed input
    :return: mean loss of the supervised data over all scales
    """

    probability_map0 = probability_map.pop(0)
    probability_map1 = probability_map.pop(0)
    probability_map2 = probability_map.pop(0)
    probability_map3 = probability_map.pop(0)

    total_loss = (single_supervised_loss(probability_map0, ground_truth) +
                  single_supervised_loss(probability_map1, ground_truth) +
                  single_supervised_loss(probability_map2, ground_truth) +
                  single_supervised_loss(probability_map3, ground_truth)) / 4

    return total_loss


def uncertainty_minimization(probability_map: list, mean_prob_map: torch.Tensor) -> Tuple[torch.Tensor, Tuple]:
    """
    Computes the uncertainty minimization term for the unsupervised loss.
    :param mean_prob_map: mean probability map over all the scales
    :param probability_map: probability map outputs of the segmentation model from different scales
    :return: uncertainty minimization term and a tuple of the exponential weights
    """

    kld_loss = KLDivLoss(reduction='none')

    kl_loss_scale0 = torch.sum(kld_loss(torch.log(probability_map[0]),
                               mean_prob_map), dim=1, keepdim=True)
    exp_kl_loss_scale0 = torch.exp(-1 * kl_loss_scale0)

    kl_loss_scale1 = torch.sum(kld_loss(torch.log(probability_map[1]),
                                        mean_prob_map), dim=1, keepdim=True)
    exp_kl_loss_scale1 = torch.exp(-1 * kl_loss_scale1)

    kl_loss_scale2 = torch.sum(kld_loss(torch.log(probability_map[2]), 
                                        mean_prob_map), dim=1, keepdim=True)
    exp_kl_loss_scale2 = torch.exp(-1 * kl_loss_scale2)

    kl_loss_scale3 = torch.sum(kld_loss(torch.log(probability_map[3]), 
                                        mean_prob_map), dim=1, keepdim=True)
    exp_kl_loss_scale3 = torch.exp(-1 * kl_loss_scale3)

    UM = (torch.mean(kl_loss_scale0) + torch.mean(kl_loss_scale1) +
          torch.mean(kl_loss_scale2) + torch.mean(kl_loss_scale3)) / 4

    return UM, (exp_kl_loss_scale0, exp_kl_loss_scale1, exp_kl_loss_scale2, exp_kl_loss_scale3)


def uncertainty_rectification(probability_map: list, mean_prob_map: torch.Tensor,
                              weights: tuple) -> torch.Tensor:
    """
    Computes the uncertainty rectification term for the unsupervised loss.
    :param weights: weights computed by the kl divergences from the UM term
    :param mean_prob_map: mean probability map over all the scales
    :param probability_map: probability map outputs of the segmentation model from different scales
    :return: uncertainty rectification term
    """

    mse_loss_scale0 = torch.mean(((probability_map[0] - mean_prob_map) ** 2) *
                                 weights[0]) / torch.mean(weights[0])
    mse_loss_scale1 = torch.mean(((probability_map[1] - mean_prob_map) ** 2) *
                                 weights[1]) / torch.mean(weights[1])
    mse_loss_scale2 = torch.mean(((probability_map[2] - mean_prob_map) ** 2) *
                                 weights[2]) / torch.mean(weights[2])
    mse_loss_scale3 = torch.mean(((probability_map[3] - mean_prob_map) ** 2) *
                                 weights[3]) / torch.mean(weights[3])
    UR = (mse_loss_scale0 + mse_loss_scale1 + mse_loss_scale2 + mse_loss_scale3) / 4

    return UR


def unsupervised_loss(probability_map: list) -> torch.Tensor:
    """
    Computes the total unsupervised loss for the batch.
    :param probability_map: probability map outputs of the segmentation model from different scales
    :return: mean loss of the unsupervised data over all scales
    """

    mean_prob_map = (probability_map[0] + probability_map[1] + probability_map[2] + probability_map[3]) / 4
    UM, weights = uncertainty_minimization(probability_map, mean_prob_map)
    UR = uncertainty_rectification(probability_map, mean_prob_map, weights)
    loss = UM + UR
    return loss


def total_loss(x: list, ground_truth: torch.Tensor, region_mask: 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 region_mask: regional mask for specifying supervised regions
    :param epoch: current epoch of training
    :return: total loss
    """

    # sup_loss = supervised_loss(labeled_x, ground_truth)
    # if unlabeled_x is None:
    #     unsup_loss = 0
    # else:
    #     unsup_loss = unsupervised_loss(unlabeled_x)
    # weight = unsupervised_weight(epoch)
    # loss = sup_loss * (0.8 - weight) + unsup_loss * (0.2 + weight)

    scales_outputs_sup, scales_outputs_unsup = [], []
    for scale_output in x:
        scales_outputs_sup.append(scale_output * region_mask.to(torch.float64))
        scales_outputs_unsup.append(scale_output * ~region_mask.to(torch.float64))

    # sup_loss = supervised_loss(scales_outputs_sup, ground_truth * region_mask.to(torch.float64))
    unsup_loss = unsupervised_loss(scales_outputs_unsup)
    
    return unsup_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)