from transformers import PretrainedConfig

class GPUNetConfig(PretrainedConfig):
    model_type = "GPUNet"
    def __init__(
            self,
            in_channels=1, 
            n_classes=3,
            depth=3, 
            wf=6, 
            padding=True,
            batch_norm=False, 
            up_mode="sinc", 
            dropout=True, 
            Relu="Relu", 
            out_act="None",
            **kwargs):
        self.in_channels = in_channels
        self.n_classes = n_classes
        self.depth = depth
        self.wf = wf
        self.padding = padding
        self.batch_norm = batch_norm
        self.up_mode = up_mode
        self.dropout = dropout
        self.Relu = Relu
        self.out_act = out_act
        super().__init__(**kwargs)
        
class GPReconResNetConfig(PretrainedConfig):
    model_type = "GPReconResNet"
    def __init__(
            self,
            in_channels=1, 
            n_classes=3, 
            res_blocks=14, 
            starting_nfeatures=64, 
            updown_blocks=2, 
            is_relu_leaky=True, 
            do_batchnorm=False, 
            res_drop_prob=0.5,
            out_act="None", 
            forwardV=0, 
            upinterp_algo='sinc', 
            post_interp_convtrans=False, 
            is3D=False,
            **kwargs):
        self.in_channels = in_channels
        self.n_classes = n_classes
        self.res_blocks = res_blocks
        self.starting_nfeatures = starting_nfeatures
        self.updown_blocks = updown_blocks
        self.is_relu_leaky = is_relu_leaky
        self.do_batchnorm = do_batchnorm
        self.res_drop_prob = res_drop_prob
        self.out_act = out_act
        self.forwardV = forwardV
        self.upinterp_algo = upinterp_algo
        self.post_interp_convtrans = post_interp_convtrans
        self.is3D = is3D
        super().__init__(**kwargs)

class GPShuffleUNetConfig(PretrainedConfig):
    model_type = "GPShuffleUNet"
    def __init__(
            self,
            d=2, 
            in_ch=1, 
            num_features=64, 
            n_levels=3, 
            out_ch=3, 
            kernel_size=3, 
            stride=1, 
            dropout=True, 
            out_act="None",
            **kwargs):
        self.d = d
        self.in_ch = in_ch
        self.num_features = num_features
        self.n_levels = n_levels
        self.out_ch = out_ch
        self.kernel_size = kernel_size
        self.stride = stride
        self.dropout = dropout
        self.out_act = out_act
        super().__init__(**kwargs)