import torch
from torch import nn
import torch.nn.functional as F
from .loss import unsupervised_loss
from .utils import UpSampler, DownSampler
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.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.iter_counter = 13400

        self.down_sampler = DownSampler(self.input_channels, self.scale_channels)
        self.up_sampler = UpSampler(self.scale_channels, self.class_num, self.shape)

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

        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, 2)
        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) / 400)
        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


# python main.py semi_supervised.segmentation.entrypoint_urpc -phase train -save_dir /home/dadbeh/Results -device cuda:2 -console_file false
# python metric.py -L ../Results/TransUNetTopK/log.csv -M train_Seg_Sens_0 val_Seg_Sens_0 -M train_Seg_Sens_1 val_Seg_Sens_1 -M train_Seg_AvgSens val_Seg_AvgSens -M train_Obj_FP/P-1 val_Obj_FP/P-1 -M train_Obj_Sen-T-1 val_Obj_Sen-T-1 -M train_ls_Loss val_ls_Loss -M train_ls_focal_loss val_ls_focal_loss -M train_ls_ce_loss train_ls_ce_loss -M train_ls_dice_loss val_ls_dice_loss -M train_Obj_Prec-T-1 val_Obj_Prec-T-1
# python main.py semi_supervised.segmentation.entrypoint_urpc -phase saveModelOutput -samples_dir liveE15 -device cuda:2 -save_dir /home/dadbeh/Results -load_dir /home/dadbeh/TempResults -epoch 113 -console_file false
# python main.py semi_supervised.segmentation.entrypoint_urpc -phase saveModelOutput -samples_dir liveE15 -device cuda:2 -save_dir /home/dadbeh/Results/TopKFocalCEWeightResults -load_dir /home/dadbeh/Results/TopKFocalCEWeightResults -epoch 37 -console_file false