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 ....utils.mask_processes.resize_binarize_roi import resize_binarize_roi_torch
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 , gamma=1.0, **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.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_roi_x: torch.Tensor = None):
        output = dict()
        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_label = sum_roi_4ch_to_1ch(mammo_roi_x, ROIAggregation.mass_vs_others).squeeze()
            if len(segmentation_label.shape) == 2:
                segmentation_label = segmentation_label.unsqueeze(0)
            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))
            # apply localization losses
            supervised_outputs, localizations = 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), localizations.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(n_classes):
    return MSC(
        base=DeepLabV2(n_classes=n_classes, n_blocks=[3, 8, 20, 3], atrous_rates=[3, 6, 9, 12]),
        scales=[0.5, 0.75])


def balanced_focal_cross_entropy_localization_loss(
        probs:torch.Tensor, gt: torch.Tensor, focal_gamma: float = 1) -> torch.Tensor:
    """
        Calculates class balanced cross entropy loss weighted with the style of Focal loss.
        mean_over_classes(1/mean(focal_weights) * (focal_weights * cross_entropy_loss))
        focal_weight: (1 - p_of_the_related_class)^gamma

        Args:
            probs (torch.Tensor): Probabilities assigned by the model of shape B C ...
            gt (torch.Tensor): The ground truth containing the number of the class, whether of shape B ... or B C ..., make sure the ... matches in both probs and gt
            focal_gamma (float): The power factor used in the weight of focal loss. Default is one which means the typical balanced cross entropy.

        Returns:
            torch.Tensor: The calculated  
    """
    localization_mask = (gt == 0 ).to(torch.int64)
    probs = probs * localization_mask # just compute loss on outside regions
    gt = torch.zeros(gt.shape).to(torch.int64).to(probs.device) # make int zero
    # if channel is one in the pixel probs, convert it to the binary mode!
    if probs.shape[1] == 1:
        probs = torch.cat([1 - probs, probs], dim=1)

    assert gt.max() <= probs.shape[1] - 1, f'Expected {probs.shape[1]} classes according to pixel probs\' shape but found {gt.max()} in GT mask'
    
    if len(gt.shape) == len(probs.shape):
        assert (gt.shape[1] == 1) or (gt.shape[1] == probs.shape[1]), f'Expected the channel dim to have either one channel or the same as probs while it is of shape {gt.shape} and probs is of shape {probs.shape}'
        if gt.shape[1] == 1:
            gt = gt[:, 0, ...]
        else:
            gt = torch.argmax(gt, dim=1)  # transferring one-hot to count data

    assert len(gt.shape) == (len(probs.shape) - 1), f'Expected GT labels to be of shape B ... and pixel probs of shape B C ..., but received {gt.shape} and {probs.shape} instead.'

    # if probabilities and ground truth have different shapes, resize the ground truth
    if probs.shape[-2:] != gt.shape[-2:]:

        # convert gt to one-hot it before interpolation to prevent mistakes
        gt = F.one_hot(gt.long(), probs.shape[1])  # B H W C
        gt = torch.permute(gt, [0, -1, 1, 2])

        # binarize
        gt = resize_binarize_roi_torch(gt.float(), probs.shape[-2:])

        # convert one-hot to numbers with argmax
        gt = torch.argmax(gt, dim=1)  # Eliminating the channel
    
    gt = gt.long()  # B H W

    # flattening and bringing class channel o index 0
    probs = probs.transpose(0, 1).flatten(1)  # C N
    gt = gt.flatten(0)  # N

    c = probs.shape[0]

    gt_related_prob = probs[gt, torch.arange(gt.shape[0])]

    losses = []

    for ci in range(c):
        c_probs = gt_related_prob[gt == ci]

        if torch.numel(c_probs) > 0:
            
            if focal_gamma == 1:
                losses.append(F.binary_cross_entropy(
                    c_probs, 
                    torch.ones_like(c_probs)))
            else:
                w = torch.pow(torch.abs(1 - c_probs.detach()), focal_gamma)
                losses.append(
                    (1.0 / (torch.mean(w) + 1e-4)) * 
                    F.binary_cross_entropy(
                        c_probs, 
                        torch.ones_like(c_probs), 
                        weight=w))

    return torch.mean(torch.stack(losses))

