import torch
import torch.nn as nn
from .deeplabv2 import DeepLabV2
from .losses import *
from mlassistant.core import Model
from ...full_segmentation.utils import sum_roi_4ch_to_1ch
from torchvision.utils import draw_segmentation_masks
from ....enums import *
from typing import List
from .msc import MSC
from ....utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss
from ....augmentations.random_tumor_augmentor import RandomTumorAugmentor
from ....augmentations.random_tumor_coloring import RandomTumorColoring

images_path = './../../images/'


class EntropyMinimization(Model):

    def __init__(self,
                 num_classes=2,
                 _lambda=1.0,
                 input_size=64,
                 gamma=1.0,
                 multi_head: torch.nn.Module = MSC,
                 main_model: torch.nn.Module = DeepLabV2,
                 atrous_rates: List[int] = [3, 6, 9, 12],
                 **kwargs):
        super().__init__()
        self.multi_head = multi_head
        self.main_model = main_model
        self.augmentor = RandomTumorAugmentor(probability=0.5)
        self.tumor_coloring = RandomTumorColoring(probability=0.5)
        self.atrous_rates = atrous_rates
        self.base_network = self.DeepLabV2S_ResNet101_MSC(n_classes=num_classes)
        self.num_classes = num_classes
        self._lambda = _lambda
        self.batch_count = 0
        self.input_size = input_size
        self.gamma = gamma
        # self.activation = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_roi_x_exist: torch.Tensor,
                mammo_y_breast_type:torch.Tensor,
                mammo_roi_x: torch.Tensor):
        output = dict()
        if mammo_roi_x is not None:
            segmentation_label = sum_roi_4ch_to_1ch(mammo_roi_x,
                                                    ROIAggregation.tumor_vs_norm).squeeze()
            segmentation_label[segmentation_label>0] = 1
            if len(segmentation_label.shape) == 2:
                segmentation_label = segmentation_label.unsqueeze(0)
            mammo_x ,segmentation_label = self.tumor_coloring.forward(mammo_x,segmentation_label)
            for i in range(len(mammo_x)):
                mammo_x[i] ,segmentation_label[i] = self.augmentor.forward(mammo_x[i], segmentation_label[i], int(mammo_y_breast_type[i]))
        logits, network_output, shapes = self.base_network(mammo_x)
        output['pixel_probs'] = self.softmax(logits)
        if mammo_roi_x is not None:
            mammo_roi_x_exist = mammo_roi_x_exist.bool()
            segmentation_labels = []
            segmentation_dict = dict()
            for i, shape in enumerate(shapes):
                start_dim = int(len(shape) - 3)
                key = '_'.join([str(i) for i in shape])
                if key not in segmentation_dict:
                    segmentation = F.interpolate(
                        segmentation_label.unsqueeze(1),
                        size=(shape[-2], shape[-1]),
                        mode="bilinear",
                        align_corners=False).squeeze()
                    if len(segmentation.shape) == 2:
                        segmentation = segmentation.unsqueeze(0)
                    output['org_pixel_labels_' + str(i)] = segmentation.clone()
                    segmentation = torch.flatten(segmentation, start_dim=start_dim)
                    segmentation_dict[key] = segmentation
                else:  # already interpolated!
                    segmentation = segmentation_dict[key]
                segmentation_labels.append(segmentation)
            final_segmentation_label = torch.cat(segmentation_labels, dim=start_dim)
            # ent_loss = torch.tensor(0.0).to(mammo_x.device)
            segmentation_loss = torch.tensor(0.0).to(mammo_x.device)
            network_output = self.softmax(network_output)
            # unsupervised_outputs = network_output[~mammo_roi_x_exist]
            # if len(unsupervised_outputs) > 0:
            # ent_loss = ent_loss + torch.mean(entropy_loss(unsupervised_outputs))
            supervised_outputs, rois = network_output[mammo_roi_x_exist], final_segmentation_label[
                mammo_roi_x_exist]
            if len(supervised_outputs) > 0:
                segmentation_loss = segmentation_loss + balanced_focal_cross_entropy_loss(
                    supervised_outputs.unsqueeze(-1),
                    rois.to(torch.int64).unsqueeze(-1),
                    focal_gamma=self.gamma)
            output['org_pixel_labels'] = output['org_pixel_labels_0']
            loss = segmentation_loss  # + self._lambda * torch.mean(ent_loss)
            output['loss'] = loss
        else:
            output['loss'] = torch.tensor(0.0)
            output['org_pixel_labels'] = output['pixel_probs']
        return output

    def DeepLabV2S_ResNet101_MSC(self, n_classes):
        return self.multi_head(
            base=self.main_model(
                n_classes=n_classes, n_blocks=[3, 4, 23, 3], atrous_rates=self.atrous_rates),
            scales=[0.5, 0.75])
