#!/usr/bin/python3
# -*- coding: utf-8 -*-
from toolbox.torchaudio.configuration_utils import PretrainedConfig


class DemucsConfig(PretrainedConfig):
    def __init__(self,
                 sample_rate: int = 8000,

                 in_channels: int = 1,
                 out_channels: int = 1,
                 hidden_channels: int = 48,

                 depth: int = 5,
                 kernel_size: int = 8,
                 stride: int = 4,

                 causal: bool = True,
                 resample: int = 4,
                 growth: int = 2,

                 max_hidden: int = 10_000,
                 do_normalize: bool = True,
                 rescale: float = 0.1,
                 floor: float = 1e-3,

                 **kwargs
                 ):
        super(DemucsConfig, self).__init__(**kwargs)
        self.sample_rate = sample_rate

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.hidden_channels = hidden_channels

        self.depth = depth
        self.kernel_size = kernel_size
        self.stride = stride

        self.causal = causal
        self.resample = resample
        self.growth = growth

        self.max_hidden = max_hidden
        self.do_normalize = do_normalize
        self.rescale = rescale
        self.floor = floor


if __name__ == "__main__":
    pass