import torch
import torch.nn.functional as F
import torch.nn as nn
import numpy as np

from .utils import init_weight, 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 Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()

        self.branch1 = RESUNETSEGMENTOR(1, 2)
        self.branch2 = RESUNETSEGMENTOR(1, 2)

    def forward(self, data, step=1):
        if step == 1:
            return self.branch1(data)
        elif step == 2:
            return self.branch2(data)


class CPS(Model):
    """
    The main design for CPS method with ResUNet as the
    base network.
    """

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

        self.iter_counter = 0
        self.network = Network()

        init_weight(self.network.branch1.business_layer, nn.init.kaiming_normal_,
                nn.BatchNorm2d, 1e-5, 0.1,
                mode='fan_in', nonlinearity='relu')
        init_weight(self.network.branch2.business_layer, nn.init.kaiming_normal_,
                nn.BatchNorm2d, 1e-5, 0.1,
                mode='fan_in', nonlinearity='relu')

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

        pred1 = F.softmax(self.network(mammo_x, step=1), dim=1)
        pred2 = F.softmax(self.network(mammo_x, step=2), dim=1)

        unsup_loss = unsupervised_loss(pred1, pred2)

        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 = pred1.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_out1 = pred1.clone().to(mammo_x.device)
        dummy_net_out2 = pred2.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_out1 = torch.cat((dummy_net_out1, torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_net_out2 = torch.cat((dummy_net_out2, torch.ones((dummy_shape[0],1,dummy_shape[2],dummy_shape[2])).to(mammo_x.device)), 1)
        dummy_yhat = pred1 * loss_channels
        dummy_y = mask_channels * loss_channels

        focal_loss1 = balanced_focal_cross_entropy_loss_semi(dummy_net_out1, dummy_mask, 4)
        focal_loss2 = balanced_focal_cross_entropy_loss_semi(dummy_net_out2, dummy_mask, 4)
        dice_ls = dice_loss(dummy_net_out1, dummy_mask)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)

        weight = self._get_unsup_weight()

        output = {
            'pixel_probs': pred1,
            'org_pixel_labels': mammo_loss_and_gt,
            'loss': focal_loss1 + focal_loss2 + unsup_loss,
            'ce_loss': ce_loss,
            'dice_loss': dice_ls,
        }

        return output

    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)

