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 ..base_line.Resunet.Heavy.Models.resunet import Encoder
from ..base_line.Resunet.Heavy.Models.resunet import Decoder


class BIAS_MLP(Model):
    def __init__(self) -> None:
        super(BIAS_MLP, self).__init__()
        self.MLP = nn.ModuleList([
            nn.Conv2d(2,16,1),
            nn.ReLU(),
            nn.Conv2d(16,8,1),
            nn.ReLU(),
            nn.Conv2d(8,2,1),
        ])
    def forward(self,x):
        out = x
        for module in self.MLP:
            out = module(out)
        return out



class RESUNET(Model):
    def __init__(self,
            gamma = 4 , smooth=1, alpha=0.7, beta=0.3, params = None) -> None:
        super(RESUNET, self).__init__()
        self.gamma = gamma
        self.smooth = smooth
        self.alpha = alpha
        self.beta = beta
        print(self.gamma)
        print("gamma")
        self.MLP = BIAS_MLP()
        self.encoder = Encoder(params)
        self.decoder = Decoder(params)
        self.epochNUM = 0
        self.forward_call = 0
        self.mode = True

    def train(self,mode : bool=True):
        super().train(mode)
        self.epochNUM += 1
        self.mode = mode
        if mode:
            self.forward_call = 0

    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,:,:]
        if self.mode:
            if (not self.mode) or (int(self.forward_call) % 50 == 0) and(self.forward_call > 10):
                print("changing")
                print(list(self.encoder.parameters())[0][0])
                network_output = self.decoder(self.encoder(mammo_x)) 
                network_output_soft = network_output.softmax(dim=1) #torch.exp(network_output).to(mammo_x.device) # softmax output of model

                bias_output = self.MLP(network_output) + network_output
                bias_output = bias_output.softmax(dim=1)

                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_l = mask_channels.clone().to(mammo_x.device)
                dummy_net_out_l = bias_output.clone().to(mammo_x.device)
                dummy_mask_l[loss_channels == 0] = 0.0
                dummy_shape_l = dummy_mask_l.shape
                dummy_mask_l = torch.cat((dummy_mask_l,torch.ones((dummy_shape_l[0],1,dummy_shape_l[2],dummy_shape_l[2])).to(mammo_x.device)), 1)
                dummy_net_out_l = torch.cat((dummy_net_out_l,torch.ones((dummy_shape_l[0],1,dummy_shape_l[2],dummy_shape_l[2])).to(mammo_x.device)), 1)
                focal_loss_l = balanced_focal_cross_entropy_loss_semi(
                        dummy_net_out_l,
                        dummy_mask_l,
                        focal_gamma=self.gamma)
                
                with torch.no_grad():
                    dummy_mask_u = (bias_output >= 0.5).clone().to(mammo_x.device) # TODO sanity check
                dummy_net_out_u = bias_output.clone().to(mammo_x.device)
                dummy_mask_u[loss_channels == 1.0] = 0.0
                dummy_shape_u = dummy_mask_u.shape
                dummy_mask_u = torch.cat((dummy_mask_u,torch.ones((dummy_shape_u[0],1,dummy_shape_u[2],dummy_shape_u[2])).to(mammo_x.device)), 1)
                dummy_net_out_u = torch.cat((dummy_net_out_u,torch.ones((dummy_shape_u[0],1,dummy_shape_u[2],dummy_shape_u[2])).to(mammo_x.device)), 1)
                focal_loss_u = balanced_focal_cross_entropy_loss_semi(
                        dummy_net_out_u,
                        dummy_mask_u,
                        focal_gamma=self.gamma) 
                # now two loss for eq 3 are here focal_loss_u & focal_loss_l

                # choose balance set             
                output['loss'] = focal_loss_u + focal_loss_l 
                output['MYloss'] = output["loss"]
                output['focalloss_l'] = focal_loss_l
                output['focalloss_u'] = focal_loss_u
                output['suploss'] = focal_loss_l
                output['unsuploss'] = focal_loss_u
                output['org_pixel_labels'] = mammo_loss_and_gt
            else:
                with torch.no_grad(): # TODO check sanity 
                    #self.encoder.
                    print("torch no grad")
                    features = self.encoder(mammo_x)
                    print(list(self.encoder.parameters())[0][0])
                network_output = self.decoder(features)
                network_output_soft = network_output.softmax(dim=1)
                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_l = mask_channels.clone().to(mammo_x.device)
                dummy_net_out_l = network_output_soft.clone().to(mammo_x.device)
                balance_mask = loss_channels.clone().to(mammo_x.device)
                #balance_mask[loss_channels == 0] = 2 # problem here sometimes all of it are supeprvised so we apply it in all of image.
                main_shape = balance_mask.shape
                number_1 = torch.sum(balance_mask == 1).to(mammo_x.device)
                number_0 = torch.sum(balance_mask == 0).to(mammo_x.device)
                random_matrix = torch.rand(main_shape).to(mammo_x.device)
                #print(number_0, number_1)
                if number_1 > number_0: #choosing balanced version
                    balance_mask[(balance_mask == 1) & (random_matrix >= number_0/number_1)] = 2
                else:
                    balance_mask[(balance_mask == 0) & (random_matrix >= number_1/number_0)] = 2

                dummy_mask_l[balance_mask == 2] = 0.0
                dummy_shape_l = dummy_mask_l.shape
                dummy_mask_l = torch.cat((dummy_mask_l,torch.ones((dummy_shape_l[0],1,dummy_shape_l[2],dummy_shape_l[2])).to(mammo_x.device)), 1)
                dummy_net_out_l = torch.cat((dummy_net_out_l,torch.ones((dummy_shape_l[0],1,dummy_shape_l[2],dummy_shape_l[2])).to(mammo_x.device)), 1)
                focal_loss_l = balanced_focal_cross_entropy_loss_semi(
                        dummy_net_out_l,
                        dummy_mask_l,
                        focal_gamma=self.gamma)
                
                output['loss'] = focal_loss_l 
                output['MYloss'] = output["loss"]
                output['focalloss_l'] = focal_loss_l
                output['focalloss_u'] = focal_loss_l * 0
                output['suploss'] = focal_loss_l
                output['unsuploss'] = focal_loss_l * 0
                output['org_pixel_labels'] = mammo_loss_and_gt



        self.forward_call += 1
        return output 