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
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_roi_x: torch.Tensor) -> ModelIO:

        output = dict(loss=0.0)
        network_output = self.base.forward(mammo_x).softmax(dim=1) 
        network_output_soft = network_output 
        output['pixel_probs'] = network_output_soft

        focal_loss = balanced_focal_cross_entropy_loss(
                network_output_soft,
                mammo_roi_x,
                focal_gamma=self.gamma)

        output['loss'] = focal_loss
        output['org_pixel_labels'] = mammo_roi_x

        return output 