from copy import deepcopy
import math
from random import uniform
import mlassistant
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 ....models.semi_supervised.base_line.Resunet.Heavy.Models.resunet import Encoder
from ....models.semi_supervised.base_line.Resunet.Heavy.Models.resunet import Decoder
from .utils import *
from .decoders import *

class CCT(Model):
    def __init__(self,
            gamma = 4 , smooth=1, alpha=0.7, beta=0.3, params = None , Epochs = 200, iter = 200) -> None:
        super(CCT, self).__init__()
        self.gamma = gamma
        self.smooth = smooth
        self.alpha = alpha
        self.beta = beta
        print(self.gamma)
        print("gamma")
        self.encoder = Encoder(params)
        self.main_decoder = Decoder(params)
        self.conf = Config["model"]
        self.conft = Config

        rampup_ends = int(self.conft['ramp_up'] * Epochs)
        cons_w_unsup = consistency_weight(final_w=self.conft['unsupervised_w'], iters_per_epoch=iter,
                                         rampup_ends=rampup_ends)

        self.unsup_loss_w = cons_w_unsup 

        self.unsuper_loss = softmax_mse_loss
        self.sup_loss_w = self.conf['supervised_w']
        self.softmax_temp = self.conf['softmax_temp']
        self.sup_type = self.conf['sup_loss']

        # Use weak labels
        self.use_weak_lables = self.conft['use_weak_lables']
        self.weakly_loss_w = self.conft['weakly_loss_w']
        # pair wise loss (sup mat)
        self.aux_constraint = self.conf['aux_constraint']
        self.aux_constraint_w = self.conf['aux_constraint_w']
        # confidence masking (sup mat)
        self.confidence_th = self.conf['confidence_th']
        self.confidence_masking = self.conf['confidence_masking']

        upscale = 8
        num_out_ch = 512
        decoder_in_ch = num_out_ch 
        num_classes = 2

        # The auxilary decoders
        vat_decoder = [VATDecoder(upscale, decoder_in_ch, num_classes, xi=self.conf['xi'],
                                    eps=self.conf['eps']) for _ in range(self.conf['vat'])]
        drop_decoder = [DropOutDecoder(upscale, decoder_in_ch, num_classes,
                                    drop_rate=self.conf['drop_rate'], spatial_dropout=self.conf['spatial'])
                                    for _ in range(self.conf['drop'])]
        cut_decoder = [CutOutDecoder(upscale, decoder_in_ch, num_classes, erase=self.conf['erase'])
                                    for _ in range(self.conf['cutout'])]
        context_m_decoder = [ContextMaskingDecoder(upscale, decoder_in_ch, num_classes)
                                    for _ in range(self.conf['context_masking'])]
        object_masking = [ObjectMaskingDecoder(upscale, decoder_in_ch, num_classes)
                                    for _ in range(self.conf['object_masking'])]
        feature_drop = [FeatureDropDecoder(upscale, decoder_in_ch, num_classes)
                                    for _ in range(self.conf['feature_drop'])]
        feature_noise = [FeatureNoiseDecoder(upscale, decoder_in_ch, num_classes,
                                    uniform_range=self.conf['uniform_range'])
                                    for _ in range(self.conf['feature_noise'])]

        # self.aux_decoders = nn.ModuleList([*vat_decoder, *drop_decoder, *cut_decoder,
        #                         *context_m_decoder, *object_masking, *feature_drop, *feature_noise])
        self.aux_decoders = nn.ModuleList([*drop_decoder,*vat_decoder,*feature_noise,
                                *context_m_decoder, *object_masking, *feature_drop])
        self.iter = 0
        self.mode = True


    def train(self,mode : bool=True):
        super().train(mode)
        self.mode = mode
        self.iter = 0
        print("epoch:",mlassistant.context.epoch.EpochContext.epoch)
    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. """
        self.iter = 1
        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,:,:]
        curr_iter= self.iter 
        epoch= mlassistant.context.epoch.EpochContext.epoch
        network_output = self.main_decoder(self.encoder(mammo_x)) 
        network_output_soft = network_output.softmax(dim=1)
        output['pixel_probs'] = network_output_soft
        output_ul = F.interpolate(network_output_soft, size=(256,256), mode='bilinear')


        # Get auxiliary predictions
        outputs_ul = [aux_decoder(self.encoder(mammo_x)[-1], output_ul.detach()) for aux_decoder in self.aux_decoders]
        targets = F.softmax(output_ul.detach(), dim=1)

        # Compute unsupervised loss
        loss_unsup = sum([self.unsuper_loss(inputs=u, targets=targets, \
                        conf_mask=self.confidence_masking, threshold=self.confidence_th, use_softmax=False)
                        for u in outputs_ul])
        
        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)
        loss_unsup = (loss_unsup / len(outputs_ul))

        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)
        focal_loss = balanced_focal_cross_entropy_loss_semi(
            dummy_net_out,
            dummy_mask,
            focal_gamma=self.gamma)

        loss_sup = focal_loss 
        weight_u = self.unsup_loss_w(epoch=epoch, curr_iter=curr_iter)


        output['loss'] =  loss_sup * self.sup_loss_w + loss_unsup * weight_u 
        output['MYloss'] = output["loss"]
        output['focalloss'] = focal_loss
        output['suploss'] = loss_sup
        output['unsuploss'] = loss_unsup
        output['sup_loss_wloss'] = torch.tensor(self.sup_loss_w)
        output['lossweight_u'] = torch.tensor(weight_u)
        output['org_pixel_labels'] = mammo_loss_and_gt

        self.iter += 1
        return output 