import torch
import torch.nn.functional as F
import numpy as np
from .unet_model import UNet
from mlassistant.core import ModelIO, Model
from ......utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss_semi
from ......utils.generalized_dice import dice_loss



class UNetBaseline(Model):
    """
    UNet architecture for baseline evaluation.
    """

    def __init__(self):
        super().__init__()

        self.scale_channels = [64, 128, 128, 256, 512]  # Filter numbers for each scale
        self.shape = (512, 512)  # Change the probability map shape according to the mammo scans shape
        self.input_channels = 1
        self.class_num = 2
        self.unet = UNet(1, 2)

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_loss_and_gt: torch.Tensor) -> ModelIO:
        
        loss_type = mammo_loss_and_gt[:, 0, ...]
        ground_truth = mammo_loss_and_gt[:, 1, ...] * loss_type
        out_prob_map = self.unet(mammo_x)

        main_shape = ground_truth.shape
        out_prob_map_shape = out_prob_map.shape

        loss_channels = loss_type.unsqueeze(dim=1).expand(main_shape[0],out_prob_map_shape[1],main_shape[1], main_shape[2])
        ground_truth = ground_truth.unsqueeze(dim=1)
        mask_channels = torch.cat((1 - ground_truth, ground_truth), 1)

        dummy_mask = mask_channels.clone().to(mammo_x.device)
        dummy_net_out = out_prob_map.clone().to(mammo_x.device)
        dummy_mask[loss_channels == 0] = 0.0
        dummy_shape = dummy_mask.shape
        dummy_mask = torch.cat((dummy_mask,torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_net_out = torch.cat((dummy_net_out, torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_yhat = out_prob_map * loss_channels
        dummy_y = mask_channels * loss_channels

        dice = dice_loss(dummy_net_out, dummy_mask)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)
        focal_loss = balanced_focal_cross_entropy_loss_semi(dummy_net_out, dummy_mask, 2)

        output = {
            'pixel_probs': out_prob_map,
            'org_pixel_labels': mammo_loss_and_gt,
            'loss': focal_loss,
            'ce_loss': ce_loss,
            'dice_loss': dice,
        }

        return output
