#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
https://huggingface.co/spaces/alibabasglab/ClearVoice/blob/main/checkpoints/FRCRN_SE_16K/config.yaml
https://huggingface.co/spaces/alibabasglab/ClearVoice/blob/main/config/inference/FRCRN_SE_16K.yaml

"""
from toolbox.torchaudio.configuration_utils import PretrainedConfig


class FRCRNConfig(PretrainedConfig):
    def __init__(self,
                 num_gpus: int = -1,

                 lr: float = 0.001,
                 lr_scheduler: str = "CosineAnnealingLR",
                 lr_scheduler_kwargs: dict = None,

                 max_epochs: int = 100,
                 weight_decay: float = 0.00001,
                 clip_grad_norm: float = 10.,
                 seed: int = 1234,

                 sample_rate: int = 8000,
                 segment_size: int = 32000,
                 nfft: int = 512,
                 win_size: int = 512,
                 hop_size: int = 256,
                 win_type: str = "hann",

                 use_complex_networks: bool = True,
                 model_depth: int = 20,
                 model_complexity: int = 45,

                 min_snr_db: float = -10,
                 max_snr_db: float = 20,

                 num_workers: int = 4,
                 batch_size: int = 4,
                 eval_steps: int = 25000,

                 **kwargs
                 ):
        super(FRCRNConfig, self).__init__(**kwargs)
        self.num_gpus = num_gpus

        self.lr = lr
        self.lr_scheduler = lr_scheduler
        self.lr_scheduler_kwargs = lr_scheduler_kwargs or dict()

        self.max_epochs = max_epochs
        self.weight_decay = weight_decay
        self.clip_grad_norm = clip_grad_norm
        self.seed = seed

        self.sample_rate = sample_rate
        self.segment_size = segment_size
        self.nfft = nfft
        self.win_size = win_size
        self.hop_size = hop_size
        self.win_type = win_type

        self.use_complex_networks = use_complex_networks
        self.model_depth = model_depth
        self.model_complexity = model_complexity

        self.min_snr_db = min_snr_db
        self.max_snr_db = max_snr_db

        self.num_workers = num_workers
        self.batch_size = batch_size
        self.eval_steps = eval_steps


def main():
    config = FRCRNConfig()
    config.to_yaml_file("config.yaml")
    return


if __name__ == "__main__":
    main()