from typing import Tuple
import torch
import torch.nn.functional as F
import numpy as np
from copy import deepcopy

from .loss import unsupervised_loss
from .....models.semi_supervised.base_line.Resunet.Heavy.Models.resunet import RESUNet as 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 ICT(Model):
    """
    The main design for ICT method with ResUNet as the
    base network.
    """

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

        self.scale_channels = [64, 128, 128, 256, 512]
        self.shape = (512, 512)
        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.main_resunet = 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,
        )

        self.mean_teacher_resunet = deepcopy(self.main_resunet)

        for param in self.mean_teacher_resunet.parameters():
            param.detach_()

    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_resunet(mammo_x, None)["pixel_probs"]

        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, 4)
        dice_ls = dice_loss(dummy_net_out,dummy_mask)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)

        weight = self._get_unsup_weight()

        output = {
            'pixel_probs': model_out,
            'org_pixel_labels': mammo_loss_and_gt,
            'loss': focal_loss + unsup_loss, # * 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_resunet(mixed_mammo_x, None)["pixel_probs"]
        with torch.no_grad():
            unlabeled_mean_out = self.mean_teacher_resunet(unlabeled_mammo_x, None)["pixel_probs"]
            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_resunet.parameters(), self.main_resunet.parameters()):
            mean_param.data.mul_(alpha).add_(1 - alpha, param.data)

    def _get_unsup_weight(self):
        if self.iter_counter // 200 >= 70:
            return 1.0
        term = 1 - (self.iter_counter // 200) / 70
        return np.exp(-5.0 * term * term)

