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 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

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 __init__(self,
            alpha: float,
            base_model: Model,
            std: float = math.sqrt(0.0001),
            mean: float = 0.0,
            T: int=1,
            use_supervised_in_consistency_checking: bool = True,
            uncertainty_thd_calculator=dynamic_num_generator(num_iters=200, func=init_num_generator2()),
            unsup_coef_generator=dynamic_num_generator(num_iters=200, func=init_num_generator1()),
            gamma = 2) -> None:
            
        super(UncertaintyAwareStudentMeanTeacherNet, self).__init__()

        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._use_supervised_in_consistency_checking: bool = use_supervised_in_consistency_checking
        self._uncertainty_thd_calculator = uncertainty_thd_calculator
        self._unsup_coef_generator = unsup_coef_generator

        self._student: Model = base_model  
        self._teacher: Model = deepcopy(self._student)  
        self.gamma = gamma

    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 _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. """
        output = dict(loss=0.0)
        if mammo_loss_and_gt is None:
            output['loss'] = torch.tensor(0.0)
            output['loss'] = torch.tensor(0.0)
            output['mse_loss'] = torch.tensor(0.0)
            output['suploss'] = torch.tensor(0.0)
            output['focalloss'] = torch.tensor(0.0)
            #output['org_pixel_labels'] = output['pixel_probs']
            # output['ce_loss'] = torch.tensor(0.0)
            # output['ce_0_loss'] = torch.tensor(0.0)
            # output['ce_1_loss'] = torch.tensor(0.0)
            # output['dice_loss'] = torch.tensor(0.0)
            return output 
        assert mammo_loss_and_gt.shape[1] == 2 , "loss and gt is not in expected shape (B,2,size,size)"

        loss_type_lable = mammo_loss_and_gt[:,0,:,:] # 1 means supervised and zero means unsupervised
        mask = mammo_loss_and_gt[:,1,:,:] # 1 means cancer and zero means backgroundru     ru            network_output = self._student._forward(mammo_x)
        network_output = self._student._forward(mammo_x) # torch.nn.LogSoftmax has been applied on last layer 
        network_output_soft = torch.exp(network_output).to(mammo_x.device)
        
        output['pixel_probs'] = network_output_soft
        main_shape = loss_type_lable.shape
        network_output_shape = network_output.shape
        dup_loss_type = loss_type_lable.unsqueeze(dim=1).expand(main_shape[0],network_output_shape[1],main_shape[1], main_shape[2])
        mask = mask.unsqueeze(dim = 1)
        dup_mask = torch.cat((1-mask, mask), 1)
        supervised_loss = torch.tensor(0.0).to(mammo_x.device)

        focal_mask = dup_mask.clone().to(mammo_x.device)
        focal_net_out = network_output_soft.clone().to(mammo_x.device)
        focal_mask[dup_loss_type == 0] = 0.0
        focal_shape = focal_mask.shape
        focal_mask = torch.cat((focal_mask,torch.ones((focal_shape[0],1,focal_shape[2],focal_shape[2])).to(mammo_x.device)), 1)
        focal_net_out = torch.cat((focal_net_out,torch.ones((focal_shape[0],1,focal_shape[2],focal_shape[2])).to(mammo_x.device)), 1)

        supervised_loss += balanced_focal_cross_entropy_loss_semi(
                focal_net_out,
                focal_mask,
                focal_gamma=self.gamma)

        dice = dice_loss(focal_net_out,focal_mask)
        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)))
            
            out_t_u = torch.exp(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)
        UThreshhold = self._uncertainty_thd_calculator()
        mask_mse = uncertainty < UThreshhold
        mse_[mask_mse] += mse[mask_mse]  
        
        active_pxs = torch.sum(mask_mse.long(), dim=[1, 2])                     # B
        mse_loss = torch.sum(mse_ / active_pxs[:, None, None], dim=[1, 2])
        mse_loss = torch.mean(mse_loss)
        mse_coef = self._unsup_coef_generator()
        output['loss'] = mse_coef * mse_loss
        output['mse_loss'] = mse_loss
        output['mse_loss'] = mse_coef
        output['Uthreshhold'] = UThreshhold
        output['suploss'] = supervised_loss
        output['focalloss'] = supervised_loss
        output['org_pixel_labels'] = mammo_loss_and_gt
        # output['ce_loss'] = torch.tensor(0.0)
        # output['ce_0_loss'] = torch.tensor(0.0)
        # output['ce_1_loss'] = torch.tensor(0.0)
        # output['dice_loss'] = torch.tensor(0.0)
        return output 