from dataclasses import dataclass
from typing import Any

from omegaconf import OmegaConf


@dataclass
class UnfoldConfig:
    patch_size: int = 9
    padding: int = 1


@dataclass
class DataEncoderConfig:
    type: str = "unfold"
    n_channel: int = 3
    trainable: bool = False
    encoder_spec: Any = None

    def __post_init__(self):
        supported_types = {
            "unfold": UnfoldConfig,
        }

        try:
            config_cls = supported_types[self.type]
        except KeyError:
            raise ValueError(f"unsupported DataEncoder type {self.type} (must be in {supported_types.keys()})")

        default_spec = OmegaConf.structured(config_cls())
        if self.encoder_spec is None:
            self.encoder_spec = default_spec
        self.encoder_spec = OmegaConf.merge(default_spec, self.encoder_spec)
