from voicefixer.vocoder.config import Config
from voicefixer.tools.pytorch_util import try_tensor_cuda, check_cuda_availability
import torch
import librosa
import numpy as np


def tr_normalize(S):
    if Config.allow_clipping_in_normalization:
        if Config.symmetric_mels:
            return torch.clip(
                (2 * Config.max_abs_value) * ((S - Config.min_db) / (-Config.min_db))
                - Config.max_abs_value,
                -Config.max_abs_value,
                Config.max_abs_value,
            )
        else:
            return torch.clip(
                Config.max_abs_value * ((S - Config.min_db) / (-Config.min_db)),
                0,
                Config.max_abs_value,
            )

    assert S.max() <= 0 and S.min() - Config.min_db >= 0
    if Config.symmetric_mels:
        return (2 * Config.max_abs_value) * (
            (S - Config.min_db) / (-Config.min_db)
        ) - Config.max_abs_value
    else:
        return Config.max_abs_value * ((S - Config.min_db) / (-Config.min_db))


def tr_amp_to_db(x):
    min_level = torch.exp(Config.min_level_db / 20 * torch.log(torch.tensor(10.0)))
    min_level = min_level.type_as(x)
    return 20 * torch.log10(torch.maximum(min_level, x))


def normalize(S):
    if Config.allow_clipping_in_normalization:
        if Config.symmetric_mels:
            return np.clip(
                (2 * Config.max_abs_value) * ((S - Config.min_db) / (-Config.min_db))
                - Config.max_abs_value,
                -Config.max_abs_value,
                Config.max_abs_value,
            )
        else:
            return np.clip(
                Config.max_abs_value * ((S - Config.min_db) / (-Config.min_db)),
                0,
                Config.max_abs_value,
            )

    assert S.max() <= 0 and S.min() - Config.min_db >= 0
    if Config.symmetric_mels:
        return (2 * Config.max_abs_value) * (
            (S - Config.min_db) / (-Config.min_db)
        ) - Config.max_abs_value
    else:
        return Config.max_abs_value * ((S - Config.min_db) / (-Config.min_db))


def amp_to_db(x):
    min_level = np.exp(Config.min_level_db / 20 * np.log(10))
    return 20 * np.log10(np.maximum(min_level, x))


def tr_pre(npy):
    # conditions = torch.FloatTensor(npy).type_as(npy) # to(device)
    conditions = npy.transpose(1, 2)
    l = conditions.size(-1)
    pad_tail = l % 2 + 4
    zeros = (
        torch.zeros([conditions.size()[0], Config.num_mels, pad_tail]).type_as(
            conditions
        )
        + -4.0
    )
    return torch.cat([conditions, zeros], dim=-1)


def pre(npy):
    conditions = npy
    ## padding tail
    if type(conditions) == np.ndarray:
        conditions = torch.FloatTensor(conditions).unsqueeze(0)
    else:
        conditions = torch.FloatTensor(conditions.float()).unsqueeze(0)
    conditions = conditions.transpose(1, 2)
    l = conditions.size(-1)
    pad_tail = l % 2 + 4
    zeros = torch.zeros([1, Config.num_mels, pad_tail]) + -4.0
    return torch.cat([conditions, zeros], dim=-1)


def load_try(state, model):
    model_dict = model.state_dict()
    try:
        model_dict.update(state)
        model.load_state_dict(model_dict)
    except RuntimeError as e:
        print(str(e))
        model_dict = model.state_dict()
        for k, v in state.items():
            model_dict[k] = v
            model.load_state_dict(model_dict)


def load_checkpoint(checkpoint_path, device):
    checkpoint = torch.load(checkpoint_path, map_location=device)
    return checkpoint


def build_mel_basis():
    return librosa.filters.mel(
        Config.sample_rate,
        Config.n_fft,
        htk=True,
        n_mels=Config.num_mels,
        fmin=0,
        fmax=int(Config.sample_rate // 2),
    )


def linear_to_mel(spectogram):
    _mel_basis = build_mel_basis()
    return np.dot(_mel_basis, spectogram)


if __name__ == "__main__":
    data = torch.randn((3, 5, 100))
    b = normalize(amp_to_db(data.numpy()))
    a = tr_normalize(tr_amp_to_db(data)).numpy()
    print(a - b)
