"""DLC SEDD config for Hugging Face."""

import transformers


class DLCSEDDGraphConfig(transformers.PretrainedConfig):
    def __init__(self, graph_type="absorb", **kwargs):
        super().__init__(**kwargs)
        self.type = graph_type
        self.file = "data"
        self.report_all = False

class DLCSEDDNoiseConfig(transformers.PretrainedConfig):
    def __init__(self, noise_type="loglinear", sigma_min=0.0001, sigma_max=20, **kwargs):
        super().__init__(*kwargs)
        self.type = noise_type
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max

class DLCSEDDConfig(transformers.PretrainedConfig):
    """Hugging Face configuration class for dlc sedd."""

    model_type = "dlc_sedd"

    def __init__(
        self,
        length: int = 512,
        hidden_size: int = 1024,
        cond_dim: int = 128,
        n_blocks: int = 24,
        n_heads: int = 16,
        dropout: float = 0.1,
        time_conditioning: bool = True,
        scale_by_sigma: bool = True,
        shared_embedding: bool = True,
        graph_type: str = "absorb",
        tokens: float = 256,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.length = length
        self.hidden_size = hidden_size
        self.cond_dim = cond_dim
        self.n_blocks = n_blocks
        self.n_heads = n_heads
        self.dropout = dropout
        self.time_conditioning = time_conditioning
        self.scale_by_sigma = scale_by_sigma
        self.shared_embedding = shared_embedding
        self.graph = DLCSEDDGraphConfig()
        self.noise = DLCSEDDNoiseConfig()
        self.tokens = tokens