from turtle import forward
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 .msc import MSC
from ....utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss
images_path = './../../images/'
class EntropyMinimization(Model):

    def __init__(self, num_classes=2, _lambda=1.0, input_size=64, **kwargs):
        super().__init__()
        self.base_network = 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.activation = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_roi_x_exist: torch.Tensor,
                mammo_roi_x: torch.Tensor = None):
        output = dict()
        mammo_roi_x_exist = mammo_roi_x_exist.bool()
        mammo_x = mammo_x.repeat(1, 3, 1, 1)    
        segmentation_label = sum_roi_4ch_to_1ch(mammo_roi_x, ROIAggregation.mass_vs_others).squeeze()
        logits , network_output, shapes = self.base_network(mammo_x)
        # print(shapes)
        segmentation_labels = []
        segmentation_dict = dict()
        for i,shape in enumerate(shapes):
            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()
                output['org_pixel_labels_' + str(i)] = segmentation.clone()
                segmentation = torch.flatten(segmentation,start_dim=1)
                segmentation_dict[key] = segmentation
            else: # already interpolated!
                segmentation = segmentation_dict[key]
            segmentation_labels.append(segmentation)
        final_segmentation_label = torch.cat(segmentation_labels, dim=1)
        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))
        output['pixel_probs'] = self.softmax(logits)
        
        
        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=1)
        # for i in range(len(mammo_x)):
        #     if mammo_roi_x_exist[i]:
        #         for threshold in range(20):
        #             binary_threshold = (threshold + 1) * 0.05
        #             image = (mammo_x[i] * 255).to(torch.uint8).cpu()
        #             if len(image.shape)< 3:
        #                 image = image.unsqueeze(0)
        #             # image 
        #             new_label = (segmentation_label[i]>0.5).bool().unsqueeze(0).cpu()
        #             # new label
        #             new_prediction = (output['pixel_probs'][i][1].clone().detach()>binary_threshold).bool().unsqueeze(0).cpu()
        #             # new_pred
        #             masks = torch.cat([new_label,new_prediction]).cpu()
        #             colors = ["green","red"]
        #             alpha = 0.7
        #             final_image = draw_segmentation_masks(image, masks, alpha, colors).transpose(0,1).transpose(1,2)
        #             Image.fromarray(final_image.cpu().numpy()).save(
        #                 os.path.join(images_path,
        #                             str(self.batch_count) + '_' + str(i) + '_' + str(binary_threshold) + '_image_mask.png'))
        output['org_pixel_labels'] = output['org_pixel_labels_2']
        loss = segmentation_loss + self._lambda * torch.mean(ent_loss)
        output['loss'] = loss
        return output


def DeepLabV2S_ResNet101_MSC(n_classes):
    return MSC(
        base=DeepLabV2(n_classes=n_classes, n_blocks=[3, 8, 5, 3], atrous_rates=[3, 6, 9, 12]),
        scales=[0.5, 8.0])
