import torch
from torch import nn
import torch.nn.functional as F
from .loss import unsupervised_loss
from .resunet import ResunetSegmentor
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 URPC(Model):
    """
    The main design for URPC architecture with UNet as the
    base network.
    """

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

        self.shape = (512, 512)  # Change the probability map shape according to the mammo scans shape
        self.input_channels = 1
        self.class_num = 2
        self.iter_counter = 0

        encoders_conv_kws = [
            {"inp_channel": 1, "out_channel": 32, "n_conv_blocks": 2},
            {"inp_channel": 33, "out_channel": 64, "n_conv_blocks": 2},
            {"inp_channel": 65, "out_channel": 128, "n_conv_blocks": 3},
            {"inp_channel": 129, "out_channel": 256, "n_conv_blocks": 3},
        ]

        bridge_conv_kws = {"inp_channel": 257, "out_channel": 512}

        decoders_conv_kws = [
            {"inp_channel": 512, "enc_inp_channel": 256, "out_channel": 256},
            {"inp_channel": 256, "enc_inp_channel": 128, "out_channel": 64},
            {"inp_channel": 64, "enc_inp_channel": 64, "out_channel": 32},
            {"inp_channel": 32, "enc_inp_channel": 32, "out_channel": 16},
        ]

        self.base_network = ResunetSegmentor(
            seg_num_labels=self.class_num,
            concat_original_image=True,
            encoders_conv_kws=encoders_conv_kws,
            decoders_conv_kws=decoders_conv_kws,
            bridge_conv_kws=bridge_conv_kws,
        )

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_loss_and_gt: torch.Tensor) -> ModelIO:
        
        self.iter_counter += 1
        
        out3, out2, out1, out0 = self.base_network(mammo_x, None)

        loss_type = mammo_loss_and_gt[:, 0, ...]
        ground_truth = mammo_loss_and_gt[:, 1, ...] * loss_type

        main_shape = ground_truth.shape
        out_prob_map_shape = out0.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 = out0.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 = out0 * loss_channels
        dummy_y = mask_channels * loss_channels

        focal_loss = balanced_focal_cross_entropy_loss_semi(dummy_net_out, dummy_mask, 4)
        dice_ls = dice_loss(dummy_net_out,dummy_mask)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)

        weight = 0.6 * ((self.iter_counter // 200) / 200)
        unsup_loss = unsupervised_loss([out0[:, [1], ...], out1[:, [1], ...],
                                         out2[:, [1], ...], out3[:, [1], ...]])

        output = {
            'pixel_probs': out0,
            'org_pixel_labels': mammo_loss_and_gt,
            'loss': focal_loss * (0.8 - weight) + unsup_loss * (0.2 + weight),
            'ce_loss': ce_loss,
            'dice_loss': dice_ls,
        }

        return output
