from copy import deepcopy
import math
from random import uniform
from typing import List
import time
from mlassistant.core import Model, ModelIO
import torch
from torch.nn import functional as F
from torch import nn
from .....utils.losses.balanced_focal_cross_entropy import balanced_focal_cross_entropy_loss_semi
from .....utils.generalized_dice import dice_loss
from .....utils.losses.tverskyLoss import tversky_loss

from .....utils.dynamic_num_generator.num_generator import init_num_generator1, init_num_generator2
from .....utils.dynamic_num_generator.wrapper import dynamic_num_generator





class RESUNET(Model):
    def __init__(self,
            base_model: Model,
            gamma = 4 , smooth=1, alpha=0.7, beta=0.3) -> None:
        super(RESUNET, self).__init__()

        
        self.base = base_model
        self.gamma = gamma
        self.smooth = smooth
        self.alpha = alpha
        self.beta = beta
        print(self.gamma)
        print("gamma")

    def forward(self,
            mammo_x: torch.Tensor,
            mammo_loss_and_gt: torch.Tensor) -> ModelIO:
        """first forwards supervised samples through student net, then all samples through both nets. """
        output = dict(loss=0.0)
        assert mammo_loss_and_gt.shape[1] == 2 , "loss and gt is not in expected shape (B,2,size,size)"
        loss_type = mammo_loss_and_gt[:,0,:,:] 
        mask = mammo_loss_and_gt[:,1,:,:]
        network_output = self.base._forward(mammo_x) # logsoftmax # B 
        network_output_soft = torch.exp(network_output).to(mammo_x.device) # softmax output of model
        output['pixel_probs'] = network_output_soft

        main_shape = mask.shape
        network_output_shape = network_output.shape



        loss_channels = loss_type.unsqueeze(dim=1).expand(main_shape[0],network_output_shape[1],main_shape[1], main_shape[2])
        mask = mask.unsqueeze(dim = 1)
        mask_channels = torch.cat((1-mask, mask), 1)

        dummy_mask = mask_channels.clone().to(mammo_x.device)
        dummy_net_out = network_output_soft.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 = network_output_soft * loss_channels
        dummy_y = mask_channels * loss_channels
        focal_loss = balanced_focal_cross_entropy_loss_semi(
                dummy_net_out,
                dummy_mask,
                focal_gamma=self.gamma)
        dice = dice_loss(dummy_net_out,dummy_mask)
        tversky = tversky_loss(inputs = dummy_yhat[:,1,...], targets = dummy_y[:,1,...], smooth=self.smooth, alpha = self.alpha,  beta= self.beta)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)
        output['ce_loss'] = ce_loss
        output['loss'] = focal_loss 
        output['diceloss'] = dice
        output['focalloss'] = focal_loss
        output['suploss'] = output['loss'] 
        output['tverskyloss'] = tversky
        output['org_pixel_labels'] = mammo_loss_and_gt
        return output 