import sys

import torch
import torch.nn as nn

from transformers import PreTrainedModel

from .ProbUNet_model import InjectionConvEncoder2D, InjectionUNet2D, InjectionConvEncoder3D, InjectionUNet3D, ProbabilisticSegmentationNet
from .PULASkiConfigs import ProbUNetConfig

class ProbUNet(PreTrainedModel):
    config_class = ProbUNetConfig
    def __init__(self, config):
        super().__init__(config)
        
        if config.dim == 2:
            task_op = InjectionUNet2D
            prior_op = InjectionConvEncoder2D
            posterior_op = InjectionConvEncoder2D
        elif config.dim == 3:
            task_op = InjectionUNet3D
            prior_op = InjectionConvEncoder3D
            posterior_op = InjectionConvEncoder3D
        else:
            sys.exit("Invalid dim! Only configured for dim 2 and 3.")
            
        if config.latent_distribution == "normal":
            latent_distribution = torch.distributions.Normal
        else:
            sys.exit("Invalid latent_distribution. Only normal has been implemented.")
            
        self.model = ProbabilisticSegmentationNet(in_channels=config.in_channels, 
                                                    out_channels=config.out_channels, 
                                                    num_feature_maps=config.num_feature_maps,
                                                    latent_size=config.latent_size,
                                                    depth=config.depth,
                                                    latent_distribution=latent_distribution,
                                                    task_op=task_op,
                                                    task_kwargs={"output_activation_op": nn.Identity if config.no_outact_op else nn.Sigmoid, 
                                                                    "activation_kwargs": {"inplace": True}, "injection_at":  config.prob_injection_at},  
                                                    prior_op=prior_op,
                                                    prior_kwargs={"activation_kwargs": {"inplace": True}, "norm_depth": 2}, 
                                                    posterior_op=posterior_op,
                                                    posterior_kwargs={"activation_kwargs": {"inplace": True}, "norm_depth": 2},
                                                )
    def forward(self, x):
        return self.model(x)