from typing import Tuple
import torch
import torch.nn.functional as F
import numpy as np
from .loss import unsupervised_loss
from .unet 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 ICT(Model):
    """
    The main design for ICT method 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 = 14400

        self.main_unet = UNet(self.input_channels, self.scale_channels, self.class_num, False)
        self.mean_teacher_unet = UNet(self.input_channels, self.scale_channels, self.class_num, True)

    def forward(self,
                mammo_x: torch.Tensor,
                mammo_loss_and_gt: torch.Tensor) -> ModelIO:
        
        self.iter_counter += 1
        self._update_mean_teacher()

        combined_model_out, combined_mean_out = self._mix_up(mammo_x)
        model_out = self.main_unet(mammo_x)

        unsup_loss = unsupervised_loss(combined_mean_out, combined_model_out)

        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 = model_out.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 = model_out.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 = model_out * 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)

        output = {
            'pixel_probs': model_out,
            '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

    
    def _mix_up(self, unlabeled_mammo_x: torch.Tensor) -> Tuple:
        lambda_coeff = np.random.beta(1., 1.)
        shuffled_indices = torch.from_numpy(np.random.permutation(unlabeled_mammo_x.shape[0])).to(unlabeled_mammo_x.device)

        mixed_mammo_x = lambda_coeff * unlabeled_mammo_x + \
                         (1 - lambda_coeff) * unlabeled_mammo_x[shuffled_indices]

        unlabeled_model_out = self.main_unet(mixed_mammo_x)
        with torch.no_grad():
            unlabeled_mean_out = self.mean_teacher_unet(unlabeled_mammo_x)
            unlabeled_mean_out = lambda_coeff * unlabeled_mean_out + \
                                (1 - lambda_coeff) * unlabeled_mean_out[shuffled_indices]

        return unlabeled_model_out, unlabeled_mean_out

    def _update_mean_teacher(self, alpha=0.999):
        alpha = min(1 - 1 / (self.iter_counter + 1), alpha)
        for mean_param, param in zip(self.mean_teacher_unet.parameters(), self.main_unet.parameters()):
            mean_param.data.mul_(alpha).add_(1 - alpha, param.data)

