import torch
import torch.nn as nn
import torch.nn.functional as F


class Dice:
    def __init__(self,
                 to_onehot_y: bool = True,
                 to_onehot_x: bool = False,
                 use_softmax: bool = True,
                 use_argmax: bool = False,
                 include_background: bool = False,
                 layout: str = "NCDHW"):
        self.include_background = include_background
        self.to_onehot_y = to_onehot_y
        self.to_onehot_x = to_onehot_x
        self.use_softmax = use_softmax
        self.use_argmax = use_argmax
        self.smooth_nr = 1e-6
        self.smooth_dr = 1e-6
        self.layout = layout

    def __call__(self, prediction, target):
        if self.layout == "NCDHW":
            channel_axis = 1
            reduce_axis = list(range(2, len(prediction.shape)))
        else:
            channel_axis = -1
            reduce_axis = list(range(1, len(prediction.shape) - 1))
        num_pred_ch = prediction.shape[channel_axis]

        if self.use_softmax:
            prediction = torch.softmax(prediction, dim=channel_axis)
        elif self.use_argmax:
            prediction = torch.argmax(prediction, dim=channel_axis)

        if self.to_onehot_y:
            target = to_one_hot(target, self.layout, channel_axis)

        if self.to_onehot_x:
            prediction = to_one_hot(prediction, self.layout, channel_axis)

        if not self.include_background:
            assert num_pred_ch > 1, \
                f"To exclude background the prediction needs more than one channel. Got {num_pred_ch}."
            if self.layout == "NCDHW":
                target = target[:, 1:]
                prediction = prediction[:, 1:]
            else:
                target = target[..., 1:]
                prediction = prediction[..., 1:]

        assert (target.shape == prediction.shape), \
            f"Target and prediction shape do not match. Target: ({target.shape}), prediction: ({prediction.shape})."

        intersection = torch.sum(target * prediction, dim=reduce_axis)
        target_sum = torch.sum(target, dim=reduce_axis)
        prediction_sum = torch.sum(prediction, dim=reduce_axis)

        return (2.0 * intersection + self.smooth_nr) / (target_sum + prediction_sum + self.smooth_dr)


def to_one_hot(array, layout, channel_axis):
    if len(array.shape) >= 5:
        array = torch.squeeze(array, dim=channel_axis)
    array = F.one_hot(array.long(), num_classes=3)
    if layout == "NCDHW":
        array = array.permute(0, 4, 1, 2, 3).float()
    return array


class DiceCELoss(nn.Module):
    def __init__(self, to_onehot_y, use_softmax, layout, include_background):
        super(DiceCELoss, self).__init__()
        self.dice = Dice(to_onehot_y=to_onehot_y, use_softmax=use_softmax, layout=layout,
                         include_background=include_background)
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, y_pred, y_true):
        cross_entropy = self.cross_entropy(y_pred, torch.squeeze(y_true, dim=1).long())
        dice = torch.mean(1.0 - self.dice(y_pred, y_true))
        return (dice + cross_entropy) / 2


class DiceScore:
    def __init__(self, to_onehot_y: bool = True, use_argmax: bool = True, layout: str = "NCDHW",
                 include_background: bool = False):
        self.dice = Dice(to_onehot_y=to_onehot_y, to_onehot_x=True, use_softmax=False,
                         use_argmax=use_argmax, layout=layout, include_background=include_background)

    def __call__(self, y_pred, y_true):
        return torch.mean(self.dice(y_pred, y_true), dim=0)
