import numpy as np
from mindspore import log as logger
import mindspore.dataset as ds
import mindspore.dataset.audio.transforms as audio
from mindspore.dataset.audio.utils import ResampleMethod
import librosa
import time
import torch
import torchaudio.functional
import math
from torch.utils.data import Dataset


def testEagerTorchAudio(input_shape, orig_freq, new_freq):
    waveform = np.random.rand(*input_shape) * 100
    orig_freq = orig_freq
    sleep_time = 10
    new_freq = new_freq
    times2 = 100
    print("waveform.shape", waveform.shape)
    print("orig_freq:", orig_freq)
    print("new_freq:", new_freq)
    print("gcd:", math.gcd(orig_freq, new_freq))
    time.sleep(sleep_time)
    start_time = time.time()
    for i in range(times2):
        audio.Resample(orig_freq=orig_freq, new_freq=new_freq)(waveform)
    end_time = time.time()
    output1 = audio.Resample(orig_freq=orig_freq, new_freq=new_freq)(waveform)
    print("minddata time over ", times2, " runs is ", (end_time - start_time) / times2 * 1000 ** 2, "µs")

    time.sleep(sleep_time)
    times = 100
    start_time = time.time()
    for i in range(times):
        librosa.core.resample(waveform, orig_sr=orig_freq, target_sr=new_freq)
    end_time = time.time()
    print("librosa resample average time over", times, "runs is ", (end_time - start_time) / times * 1000 ** 2, "µs")

    time.sleep(sleep_time)
    waveform2 = torch.from_numpy(waveform)
    start_time = time.time()

    times3 = 1
    for i in range(times3):
        torchaudio.functional.resample(waveform2, orig_freq=orig_freq, new_freq=new_freq)
    end_time = time.time()
    output3 = torchaudio.functional.resample(waveform2, orig_freq=orig_freq, new_freq=new_freq)
    output3 = output3.detach().numpy()
    print("pytorch time over ", times3, "runs is ", (end_time - start_time) / times3 * 1000 ** 2, "µs")
    print(abs(output3 - output1))
    assert (np.all(output1 - output3 < 1e-5))

if __name__ == '__main__':
    

if __name__ == '__main__':
    test_data = [{"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 160000},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 48000},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 8000},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 11000},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 500},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 11025},
                 {"waveform": (1000, 1000), "orig_freq": 16000, "new_freq": 11005},
                 {"waveform": (1000, 1000), "orig_freq": 22050, "new_freq": 44100},
                 {"waveform": (1000, 1000), "orig_freq": 22050, "new_freq": 11025},
                 {"waveform": (1000, 1000), "orig_freq": 22050, "new_freq": 500},
                 {"waveform": (1000, 1000), "orig_freq": 22050, "new_freq": 11005},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 160000},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 48000},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 8000},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 11000},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 500},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 11025},
                 {"waveform": (17, 12345), "orig_freq": 16000, "new_freq": 11005},
                 {"waveform": (16560,), "orig_freq": 24000, "new_freq": 48000},
                 {"waveform": (16560,), "orig_freq": 24000, "new_freq": 16000},
                 {"waveform": (16560,), "orig_freq": 24000, "new_freq": 1600},
                 {"waveform": (16560,), "orig_freq": 24000, "new_freq": 11025},
                 {"waveform": (16560,), "orig_freq": 24000, "new_freq": 11005},
                 {"waveform": (7432986,), "orig_freq": 16000, "new_freq": 48000},
                 {"waveform": (7432986,), "orig_freq": 16000, "new_freq": 24000},
                 {"waveform": (7432986,), "orig_freq": 16000, "new_freq": 1600},
                 {"waveform": (7432986,), "orig_freq": 16000, "new_freq": 11025},
                 {"waveform": (241840,), "orig_freq": 16000, "new_freq": 48000},
                 {"waveform": (241840,), "orig_freq": 16000, "new_freq": 24000},
                 {"waveform": (241840,), "orig_freq": 16000, "new_freq": 1600},
                 {"waveform": (241840,), "orig_freq": 16000, "new_freq": 11025},
                 {"waveform": (241840,), "orig_freq": 16000, "new_freq": 11005},
                 {"waveform": (2, 125800), "orig_freq": 22050, "new_freq": 11025},
                 {"waveform": (2, 125800), "orig_freq": 22050, "new_freq": 44100},
                 ]
    for data in test_data:
        testEagerTorchAudio(input_shape=data["waveform"], orig_freq=data["orig_freq"], new_freq=data["new_freq"])