from copy import deepcopy
import math
from random import uniform
from typing import List
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.ent_losses import entropy_loss_normalized
from ....utils.losses.tverskyLoss import tversky_loss
import time
from mlassistant.core import Model, ModelIO
import torch
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.dynamic_num_generator.num_generator import init_num_generator1, init_num_generator2
from ....utils.dynamic_num_generator.wrapper import dynamic_num_generator
import numpy as np

def random_dropout(min, max):
    def assign(module: nn.Dropout,_: torch.Tensor) -> None:
        module.p = uniform(min, max)
    
    return assign

def static_dropout(prob: float):
    def assign(module: nn.Dropout, _: torch.Tensor) -> None:
        module.p = prob
    return assign

class UncertaintyAwareStudentMeanTeacherNet(Model):
    """ adapted from the paper 'Uncertainty-aware Self-ensembling Model for Semi-supervised 
    3D Left Atrium Segmentation'. link: https://arxiv.org/pdf/1907.07034.pdf

    Args:
        alpha (float): exponential moving average decay
        total_num_iteration (int): total number of iteration this run will have
        base_model (Model): base networks for both student and teacher 
        x_arg (str): argument name of `forward` indicating images, shape: B ...
        x_roi_arg (str): argument name of `forward` indicating ground-truth ROIs, shape: B ...
        x_annotated_arg (str): argument name of `forward` indicating being labeled or not for inputs, shape: B
        uncertainty_ratio_thd (float): the initial ratio threshold for uncertainty pruning
        std (float): standard deviation of gaussian noise to be applied to the input
        mean (float): mean of gaussian noise to be applied to the input
        T (int): number of perturbations for teacher net 
        lambda_coef (float): lambda coefficient for calculating unsupervised loss
        use_supervised_in_consistency_checking (bool): if true then use supervised samples in unsupervised task (consistency checking)
    """
    def train(self,mode : bool=True):
        super().train(mode)
        self.epochNUM += 1
        # with torch.no_grad():
        #     if self.epochNUM % 2 == 0:
        #         self.update_teacher_params()
        self.mode = mode
        print(mode)
        
        
    def __init__(self,
            alpha: float,
            base_model: Model,
            std: float = math.sqrt(0.0001),
            mean: float = 0.0,
            T: int=8,
            gamma = 2, smooth=1, alphat=0.7, beta=0.3) -> None:
            
        super(UncertaintyAwareStudentMeanTeacherNet, self).__init__()

        self.epochNUM = 1
        assert alpha >= 0 and alpha <= 1, f'alpha should be in range [0, 1]; got {alpha} instead.'
        assert T > 0, f'number of perturbations to be applied for teacher net should be at least one; got {T} instead.'

        self._T: int = T
        self._alpha: float = alpha
        self._mean: float = mean
        self._std: float = std
        self.smooth = smooth
        self.alphatt = alphat
        self.beta = beta

        self._student: Model = base_model  
        self._teacher: Model = deepcopy(self._student)  
        self.gamma = gamma
        self.uncercoef = lambda t : torch.tensor(0.1*np.exp(-5 * (1-t/200)**2))

        self.H = np.log(2)
        self.threshold = lambda t : torch.tensor((self.H/4)*np.exp(-5 * (1-t/200)**2) + self.H * 3/4)
        self.flag = True
        self.mode = True

    def _update_teacher_params(self,
            _: Model,
            __: torch.Tensor) -> None:
        """a backward hook to update teacher network parameters, using exponential moving 
        average method."""

        
        for s_p, t_p in zip(self._student.parameters(), self._teacher.parameters()):
            t_p.data = self._alpha * t_p.data + (1 - self._alpha) * s_p.data
    def update_teacher_params(self) -> None:
        """a backward hook to update teacher network parameters, using exponential moving 
        average method."""

        
        for s_p, t_p in zip(self._student.parameters(), self._teacher.parameters()):
            t_p.data = self._alpha * t_p.data + (1 - self._alpha) * s_p.data

    def _apply_gaussian_noise(self,
            input ) -> torch.Tensor:
        inp = input
        noise = torch.randn(inp.size()) * self._std + self._mean
        noise = noise.to(inp.device)
        inp = inp + (noise * inp)
        input = inp
        return input

    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. """
        with torch.no_grad():
            if self.mode and self.epochNUM > 1:
                self.update_teacher_params()
        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._student.forward(mammo_x).softmax(dim=1) # logsoftmax # B 
        network_output_soft = network_output #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)
        ce_loss = F.binary_cross_entropy(dummy_yhat, dummy_y)
        tversky = tversky_loss(inputs = dummy_yhat[:,1,...], targets = dummy_y[:,1,...], smooth=self.smooth, alpha = self.alphatt,  beta= self.beta)
        

        with torch.no_grad():
            teacher_outs: List[torch.Tensor] = list()
            for _ in range(self._T): 
                teacher_outs.append(self._teacher.forward(
                    self._apply_gaussian_noise(mammo_x)).softmax(dim=1))
            out_t_u = (torch.stack(teacher_outs, dim=2)).mean(dim=2)   # B C H' W'
            
            
        uncertainty = -1 * torch.sum(out_t_u * torch.log(out_t_u), 1)       # B H' W'
        mse = torch.sum((network_output_soft - out_t_u) ** 2, dim=1)
        
        # consider only most certain pixels for mse loss
        mse_ = torch.zeros_like(mse).to(mammo_x.device)
        UThreshhold = self.threshold(int(self.epochNUM/2))
        mask_mse = uncertainty < UThreshhold
        mse_[mask_mse] += mse[mask_mse]  
        
        active_pxs = torch.sum(mask_mse.long(), dim=[1, 2]) + 1                     # B
        mse_loss = torch.sum(mse_ / active_pxs[:, None, None], dim=[1, 2])
        mse_loss = torch.mean(mse_loss)
        mse_coef = self.uncercoef(int(self.epochNUM/2)).to(mammo_x.device)
        output['ce_loss'] = ce_loss
        output['loss'] = focal_loss + mse_coef * mse_loss
        #output['diceloss'] = dice
        #output['tverskyloss'] = tversky
        output['focalloss'] = focal_loss
        output['suploss'] = focal_loss
        output['org_pixel_labels'] = mammo_loss_and_gt
        output['mse_loss'] = mse_loss
        output['mse_coef_loss'] = torch.tensor(mse_coef)
        output['Uthreshhold_loss'] = torch.tensor(UThreshhold)
        output['totol_loss'] = output['loss']
        return output 