"""TIPSv2 DPT head configuration."""

from transformers import PretrainedConfig


_BACKBONE_REPO_BY_HIDDEN_SIZE = {
    768: "google/tipsv2-b14",
    1024: "google/tipsv2-l14",
    1152: "google/tipsv2-so400m14",
    1536: "google/tipsv2-g14",
}


class TIPSv2DPTConfig(PretrainedConfig):
    """Configuration for TIPSv2 DPT dense prediction heads."""

    model_type = "tipsv2_dpt"

    def __init__(
        self,
        backbone_config=None,
        neck_hidden_sizes=(96, 192, 384, 768),
        fusion_hidden_size=256,
        reassemble_factors=(4, 2, 1, 0.5),
        readout_type="project",
        num_depth_bins=256,
        min_depth=1e-3,
        max_depth=10.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        backbone_config = backbone_config or {}
        hidden_size = backbone_config.get("hidden_size", 768)
        out_indices = backbone_config.get("out_indices", (3, 6, 9, 12))
        self.backbone_repo = _BACKBONE_REPO_BY_HIDDEN_SIZE[hidden_size]
        self.embed_dim = hidden_size
        self.channels = fusion_hidden_size
        self.post_process_channels = list(neck_hidden_sizes)
        self.block_indices = [out_index - 1 for out_index in out_indices]
        self.readout_type = readout_type
        self.num_depth_bins = num_depth_bins
        self.min_depth = min_depth
        self.max_depth = max_depth
        self.num_seg_classes = self.num_labels