from transformers import PretrainedConfig


class CondUNetConfig(PretrainedConfig):
    model_type = "cond_unet"

    def __init__(
        self,
        in_channels=3,
        num_labels=1,
        n_organs=10,
        size=32,
        depth=5,
        attn_start=0,
        use_attn=True,
        image_size=512,
        patch_size=8,
        emb_dim=768,
        n_heads=8,
        use_dwt=False,
        wavelet="haar",
        dwt_bands=None,
        use_shape=False,
        shape_res=32,
        unknown_organ_id=-1,
        keep_aspect_ratio=True,
        self_normalize=True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.in_channels = in_channels
        self.num_labels = num_labels
        self.n_organs = n_organs
        self.size = size
        self.depth = depth
        self.attn_start = attn_start
        self.use_attn = use_attn
        self.image_size = image_size
        self.patch_size = patch_size
        self.emb_dim = emb_dim
        self.n_heads = n_heads
        self.use_dwt = use_dwt
        self.wavelet = wavelet
        self.dwt_bands = dwt_bands or ["LL", "LH", "HL", "HH"]
        self.use_shape = use_shape
        self.shape_res = shape_res
        self.unknown_organ_id = unknown_organ_id
        self.keep_aspect_ratio = keep_aspect_ratio
        self.self_normalize = self_normalize
        self.id2label = {0: "foreground"}
        self.label2id = {"foreground": 0}