import torch 
import torch as tf
from torch.nn import functional as F
import torch.nn as nn
from .deeplabv2 import DeepLabV2
from .losses import *
from mlassistant.core import Model
import sys
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_semi
from ....utils.dice import dice_loss
images_path = './../../images/'
import time
class EntropyMinimization(Model):

    def train(self,mode : bool=True):
        super().train(mode)
        self.epochNUM += 1
        print(self.epochNUM)

    def __init__(self, num_classes=2, _lambda=1.0, input_size=64 , gamma=2.0,ent_factor = 0.001, **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)
        self.epochNUM = 0
        self.ent_factor = ent_factor
        self.time_c = time.time()
    
    def interpolate(self,shapes,arr):
        segmentation_label = arr
        output = dict()

        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)
        return final_segmentation_label, output


    def forward(self,
                mammo_x: torch.Tensor,
                mammo_loss_and_gt: torch.Tensor):
        """


        """

        #print(mammo_roi_x.shape)
        #mammo_loss_and_gt = torch.cat((mammo_roi_x[:,0,:,:].unsqueeze(1), mammo_roi_x[:,0,:,:].unsqueeze(1)), 1)
        #print(mammo_loss_and_gt.shape)
        #mammo_loss_and_gt = 
        #start = time.time()
        #print("between two forwards = ", start - self.time_c, flush= True)
        output = dict()
        #print()
        #print(mammo_x.shape,flush=True)
        if mammo_loss_and_gt is not None:
            assert mammo_loss_and_gt.shape[1] == 2 , "loss and gt is not in expected shape (B,2,size,size)"

            loss_type_lable = mammo_loss_and_gt[:,0,:,:] # 1 means supervised and zero means unsupervised
            mask = mammo_loss_and_gt[:,1,:,:] # 1 means cancer and zero means background
            #deel_time = time.time()
            logits , network_output, shapes = self.base_network(mammo_x)
            #print("base model took= ",time.time() - deel_time, flush=True)
            
            output['pixel_probs'] = self.softmax(logits) 
            network_output_soft = self.softmax(network_output)

            loss_type, o1 = self.interpolate(shapes, loss_type_lable)
            
            mask, o2 = self.interpolate(shapes, mask)

            main_shape = loss_type.shape
            network_output_shape = network_output_soft.shape
            dup_loss_type = loss_type.unsqueeze(dim=1).expand(main_shape[0],network_output_shape[1],main_shape[1])
            mask = mask.unsqueeze(dim = 1)
            dup_mask = torch.cat((1-mask, mask), 1)

            output = {**output, **o1, **o2} 
            supervised_loss = torch.tensor(0.0).to(mammo_x.device)

            yhat = (network_output_soft * dup_loss_type).unsqueeze(-1)
            y = (dup_mask * dup_loss_type).unsqueeze(-1)


            focal_mask = dup_mask.clone().to(mammo_x.device)
            focal_net_out = torch.exp(torch.log(network_output_soft).to(mammo_x.device)).to(mammo_x.device)
            focal_mask[dup_loss_type == 0] = 0.0
            #focal_net_out[dup_loss_type == 0] = 0.0
            focal_shape = focal_mask.shape
            focal_mask = torch.cat((focal_mask,torch.ones((focal_shape[0],1,focal_shape[2])).to(mammo_x.device)), 1)
            focal_net_out = torch.cat((focal_net_out,torch.ones((focal_shape[0],1,focal_shape[2])).to(mammo_x.device)), 1)
            supervised_loss += balanced_focal_cross_entropy_loss_semi(
                  focal_net_out.unsqueeze(-1),
                   focal_mask.unsqueeze(-1),
                    focal_gamma=self.gamma)


            
            entropy_loss_val = entropy_loss_normalized(network_output_soft)* (1-loss_type)
            entropy_loss_val =  torch.mean(entropy_loss_val.flatten()[((1-loss_type) == 1).flatten()])

            net_out_sup = yhat.flatten()[(dup_loss_type == 1).flatten()]
            mask_sup = mask.flatten()[(loss_type == 1).flatten()]
            dice_loss_val = dice_loss((net_out_sup.reshape(2,-1).unsqueeze(dim  = 0).unsqueeze(dim  = -1)),mask_sup.unsqueeze(dim = 0).unsqueeze(dim  = -1).to(torch.int64))
            
            output['suploss'] = supervised_loss
            output['focalloss'] = supervised_loss
            output['dice_loss'] = dice_loss_val
            output["loss"] = supervised_loss + self.ent_factor * entropy_loss_val 

            output['org_pixel_labels'] = mammo_loss_and_gt
            output['ent_loss'] = entropy_loss_val 
            output['ce_loss'] = F.binary_cross_entropy(yhat,y)
            output['ce_0_loss'] = F.binary_cross_entropy(yhat * (1-y) + y,y) # y = 1 -> y / y = 0 -> yhat
            output['ce_1_loss'] = F.binary_cross_entropy(yhat * y,y) # y = 0 -> y / y = 1 -> yhat
            #print(output)

        else:
            output['loss'] = torch.tensor(0.0)
            output['org_pixel_labels'] = mammo_loss_and_gt
            output['ent_loss'] = torch.tensor(0.0)
            output['ce_loss'] = torch.tensor(0.0)
            output['ce_0_loss'] = torch.tensor(0.0)
            output['ce_1_loss'] = torch.tensor(0.0)
            output['dice_loss'] = torch.tensor(0.0)
            output['suploss'] = torch.tensor(0.0)
            output['focalloss'] = torch.tensor(0.0)

        #print("whole forward took =", time.time() - start, flush = True)
        #self.time_c = time.time()
        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))

