from omegaconf import OmegaConf

_ATTENTION_DICT = {
    "num_heads": 8,
    "attn_dropout": 0.0,
    "bias_qkv": False,
    "N": 0,
    "rel_pos_num_buckets": 32,
    "rel_pos_max_distance": 64,
    "use_rel_pos": True,
    "Nproj": 8,
}

_BLIND_BWE_OPTIMIZATION = {
    "max_iter": 100,
    "grad_clip": 1,
    "tol": [1e-1, 1e-1],
    "clamp_fc": True,
    "clamp_A": True,
    "block_low_freq": False,
    "only_negative_Ap": False,
    "last_slope_fixed": True,
    "first_slope_fixed": True,
}

_BLIND_BWE_INITIAL_CONDITIONS = {
    "fref": 1000,
    "fc_p": [1500, 2000],
    "A_p": [0, 0, -80],
    "fc_m": [500, 50],
    "A_m": [0, 0, 80],
}

_COLLAPSE_REGULARIZATION = {
    "use": True,
    "beta": 0.1,
    "gamma": 1,
    "lambda_reg": 10,
}

PIANO_CONFIG = OmegaConf.create({
    "architecture": "piano",
    "exp": {
        "sample_rate": 22050,
        "audio_len": 184184,
    },
    "network": {
        "emb_dim": 256,
        "use_norm": True,
        "use_fencoding": False,
        "Ns": [64, 96, 96, 128, 128, 256, 256],
        "Ss": [2, 2, 2, 2, 2, 2, 2],
        "num_dils": [2, 3, 4, 5, 6, 7, 7],
        "attention_layers": [0, 0, 0, 0, 0, 0, 0, 0],
        "bottleneck_type": "res_dil_convs",
        "num_bottleneck_layers": 1,
        "cqt": {"num_octs": 7, "bins_per_oct": 64, "window": "kaiser", "beta": 1},
        "attention_dict": _ATTENTION_DICT,
    },
    "diff_params": {
        "sigma_data": 0.063, "sigma_min": 1e-5, "sigma_max": 10,
        "P_mean": -1.2, "P_std": 1.2, "ro": 13, "ro_train": 10,
        "Schurn": 5, "Snoise": 1, "Stmin": 0, "Stmax": 50,
    },
    "tester": {
        "T": 51,
        "order": 2,
        "filter_out_cqt_DC_Nyq": True,
        "posterior_sampling": {
            "xi": 1.0,
            "data_consistency": False,
            "annealing_y": {"use": True, "mode": "fixed", "sigma_min": 0.25},
            "SNR_observations": "None",
            "start_sigma": 0.5,
            "normalization": "grad_norm",
            "freq_weighting_filter": "sqrt",
        },
        "diff_params": {
            "same_as_training": False,
            "sigma_data": 0.063, "sigma_min": 4e-5, "sigma_max": 10,
            "P_mean": -1.2, "P_std": 1.2, "ro": 13, "ro_train": 13,
            "Schurn": 10, "Snoise": 1.0, "Stmin": 0, "Stmax": 50,
        },
        "blind_bwe": {
            "lr_filter": 10,
            "NFFT": 4096,
            "LTAS_fft": 2048,
            "sigma_norm": 0.07,
            "fcmin": 11, "fcmax": "nyquist", "Amin": -40, "Amax": 40, "Alim": 80,
            "initial_conditions": _BLIND_BWE_INITIAL_CONDITIONS,
            "optimization": _BLIND_BWE_OPTIMIZATION,
        },
        "collapse_regularization": _COLLAPSE_REGULARIZATION,
        "evaluation": {
            "segment_length": 184184,
            "segment_placement": "middle",
            "num_segments_batch": 1,
            "LTAS_init": True,
            "LTAS_as_y": True,
            "overlap": 0.1,
            "process_complete_mode": "Block-Autoregressive",
        },
    },
})

SINGING_CONFIG = OmegaConf.create({
    "architecture": "singing",
    "exp": {
        "sample_rate": 44100,
        "audio_len": 262144,
    },
    "network": {
        "emb_dim": 256,
        "use_norm": True,
        "use_fencoding": False,
        "Ns": [32, 32, 64, 64, 128, 128, 256, 256],
        "Ss": [2, 2, 2, 2, 2, 2, 2, 2],
        "num_dils": [1, 3, 4, 5, 5, 6, 6, 7],
        "attention_layers": [0, 0, 0, 0, 0, 0, 0, 0],
        "bottleneck_type": "res_dil_convs",
        "num_bottleneck_layers": 1,
        "cqt": {"num_octs": 8, "bins_per_oct": 32, "window": "kaiser", "beta": 1},
        "attention_dict": _ATTENTION_DICT,
    },
    "diff_params": {
        "sigma_data": 1, "sigma_min": 1e-3, "sigma_max": 100,
        "P_mean": -1.2, "P_std": 1.2, "ro": 10, "ro_train": 10,
        "Schurn": 5, "Snoise": 1, "Stmin": 0, "Stmax": 50,
    },
    "tester": {
        "T": 51,
        "order": 2,
        "filter_out_cqt_DC_Nyq": True,
        "posterior_sampling": {
            "xi": 0.5,
            "data_consistency": False,
            "annealing_y": {"use": True, "mode": "fixed", "sigma_min": 1},
            "SNR_observations": "None",
            "start_sigma": 10,
            "normalization": "grad_norm",
            "freq_weighting_filter": "sqrt",
        },
        "diff_params": {
            "same_as_training": False,
            "sigma_data": 1, "sigma_min": 1e-3, "sigma_max": 100,
            "P_mean": -1.2, "P_std": 1.2, "ro": 13, "ro_train": 13,
            "Schurn": 10, "Snoise": 1.0, "Stmin": 0, "Stmax": 50,
        },
        "blind_bwe": {
            "lr_filter": 10,
            "NFFT": 4096,
            "LTAS_fft": 4096,
            "sigma_norm": 1,
            "fcmin": 11, "fcmax": "nyquist", "Amin": -40, "Amax": 40, "Alim": 80,
            "initial_conditions": _BLIND_BWE_INITIAL_CONDITIONS,
            "optimization": _BLIND_BWE_OPTIMIZATION,
        },
        "collapse_regularization": _COLLAPSE_REGULARIZATION,
        "evaluation": {
            "segment_length": 262144,
            "segment_placement": "middle",
            "num_segments_batch": 1,
            "LTAS_init": True,
            "LTAS_as_y": True,
            "overlap": 0.1,
            "process_complete_mode": "Block-Autoregressive",
        },
    },
})

# display label -> (checkpoint filename on the HF Hub, architecture name)
CHECKPOINTS = {
    "Piano (MAESTRO)": ("MAESTRO_22kHz_8s-850kits.pt", PIANO_CONFIG),
    "Singing Voice — Generic": ("singing_voice_pretrain_44kHz_6s-325kits.pt", SINGING_CONFIG),
    "Singing Voice — Enrico Caruso": ("VocalSet_male2_44kHz_6s-8kits.pt", SINGING_CONFIG),
    "Singing Voice — Beniamino Gigli": ("VocalSet_male11_44kHz_6s-5kits.pt", SINGING_CONFIG),
    "Singing Voice — Nellie Melba": ("VocalSet_female1_44kHz_6s-8kits.pt", SINGING_CONFIG),
    "Singing Voice — Adelina Patti": ("VocalSet_female5_44kHz_6s-8kits.pt", SINGING_CONFIG),
}

CHECKPOINTS_REPO = "teamup-tech/BABE2-checkpoints"