from os.path import dirname, join
from warnings import warn

import numpy as np
import pysptk
import pytest
from scipy.io import wavfile


@pytest.mark.parametrize("hopsize", [40, 80, 160, 320])
@pytest.mark.parametrize("otype", [0, 1, 2])
@pytest.mark.parametrize("otype_str", ["pitch", "f0", "logf0"])
def test_swipe(hopsize, otype, otype_str):
    def __test(x, fs, hopsize, otype):
        f0 = pysptk.swipe(x, fs, hopsize, otype=otype)
        assert np.all(np.isfinite(f0))
        if otype == 1:
            assert np.all(f0 >= 0)

    np.random.seed(98765)
    fs = 16000
    x = np.random.rand(16000)

    __test(x, fs, hopsize, otype)
    __test(x, fs, 80, otype_str)


def test_swipe_corner_case():
    def __test(x, fs, hopsize, otype):
        pysptk.swipe(x, fs, hopsize, otype=otype)

    np.random.seed(98765)
    fs = 16000
    x = np.random.rand(16000)

    # unsupported otype
    with pytest.raises(ValueError):
        __test(x, fs, 80, -1)
    with pytest.raises(ValueError):
        __test(x, fs, 80, 3)
    with pytest.raises(ValueError):
        __test(x, fs, 80, "ff0")


@pytest.mark.parametrize("hopsize", [40, 80, 160, 320])
@pytest.mark.parametrize("otype", [0, 1, 2])
@pytest.mark.parametrize("otype_str", ["pitch", "f0", "logf0"])
def test_rapt(hopsize, otype, otype_str):
    def __test(x, fs, hopsize, min, max, otype):
        f0 = pysptk.rapt(x, fs, hopsize, min=min, max=max, otype=otype)
        assert np.all(np.isfinite(f0))
        if otype == 1:
            assert np.all(f0 >= 0)

    np.random.seed(98765)
    fs = 16000
    x = np.random.rand(16000).astype(np.float32)

    __test(x, fs, hopsize, 60, 240, otype)
    __test(x, fs, 80, 60, 240, otype_str)


def test_rapt_corner_case():
    def __test(x, fs, hopsize, min, max, otype):
        f0 = pysptk.rapt(x, fs, hopsize, min=min, max=max, otype=otype)
        assert np.all(np.isfinite(f0))
        if otype == 1:
            assert np.all(f0 >= 0)

    np.random.seed(98765)
    fs = 16000
    x = np.random.rand(16000).astype(np.float32)

    # unsupported otype
    with pytest.raises(ValueError):
        __test(x, fs, 80, 60, 240, -1)
    with pytest.raises(ValueError):
        __test(x, fs, 80, 60, 240, 3)
    with pytest.raises(ValueError):
        __test(x, fs, 80, 60, 240, "f00")

    # valid min freq
    __test(x, fs, 80, 10, 240, 0)
    warn("TODO: fix RAPT bug to pass this minfreq lower bound test")
    # __test(x, fs, 80, fs / 10000. + 1, 240, 0)

    # valid max freq
    __test(x, fs, 80, 60, fs // 2 - 1, 0)

    # invalid min/max freq
    with pytest.raises(ValueError):
        __test(x, fs, 80, 60, 60, 0)
    with pytest.raises(ValueError):
        __test(x, fs, 80, 60, fs // 2, 0)
    with pytest.raises(ValueError):
        __test(x, fs, 80, fs / 10000.0, 240, 0)

    # valid frame_period (corner case)
    __test(x, fs, 1600, 60, 240, 0)
    __test(x, fs, 2, 60, 240, 0)

    warn("TODO: pass this corner case test")
    # __test(x, fs, 1, 60, 240, 0)

    # invalid frame_period
    with pytest.raises(ValueError):
        __test(x, fs, 1601, 60, 240, 0)

    # valid input length
    __test(x[:1000], fs, 80, 60, 240, 0)

    # invalid input length (too small)
    with pytest.raises(ValueError):
        __test(x[:100], fs, 80, 60, 240, 0)


def test_rapt_regression():
    # Grund truth data is generated by:
    #
    # $ wav2raw pysptk/example_audio_data/arctic_a0007.wav
    #
    # $ x2x +sf ./pysptk/example_audio_data/arctic_a0007.raw | \
    #    pitch -a 0 -s 16 -p 80 -L 60 -H 240 -o 0 > \
    #    arctic_a007_p16_L60_H240_o0_rapt.pitch
    #
    # $ dmp +f arctic_a007_p16_L60_H240_o0_rapt.pitch | awk '{print $2}' >\
    #    arctic_a007_p16_L60_H240_o0_rapt.txt
    #
    # $ pitch -h
    #  ...
    #
    #  SPTK: version 3.8
    #  CVS Info: $Id: pitch.c,v 1.46 2014/12/11 08:30:43 uratec Exp $

    ground_truth_path = join(
        dirname(__file__), "data", "arctic_a007_p16_L60_H240_o0_rapt.txt"
    )
    with open(ground_truth_path) as f:
        ground_truth = np.asarray([float(s) for s in [line for line in f.readlines()]])
    ground_truth = ground_truth.astype(np.float32)

    fs, x = wavfile.read(pysptk.util.example_audio_file())
    assert fs == 16000

    # Since SPTK might have memory corruption bug and the result might be
    # non-deterministic, test it with multiple time...
    for _ in range(5):
        f0 = pysptk.rapt(
            x.astype(np.float32),
            fs=fs,
            hopsize=80,
            min=60,
            max=240,
            voice_bias=0.0,
            otype=0,
        )
        assert np.allclose(ground_truth, f0)


def test_swipe_regression():
    # Grund truth data is generated by:
    #
    # $ wav2raw pysptk/example_audio_data/arctic_a0007.wav
    #
    # $ x2x +sf ./pysptk/example_audio_data/arctic_a0007.raw | \
    #    pitch -a 1 -s 16 -p 80 -L 60 -H 240 -o 0 > \
    #    arctic_a007_p16_L60_H240_o0_swipe.pitch
    #
    # $ dmp +f arctic_a007_p16_L60_H240_o0_swuoe.pitch | awk '{print $2}' >\
    #    arctic_a007_p16_L60_H240_o0_swipe.txt
    #
    # $ pitch -h
    #  ...
    #
    # SPTK: version 3.10
    # CVS Info: $Id: pitch.c,v 1.53 2016/12/25 05:00:19 uratec Exp $

    ground_truth_path = join(
        dirname(__file__), "data", "arctic_a007_p16_L60_H240_o0_swipe.txt"
    )
    with open(ground_truth_path) as f:
        ground_truth = np.asarray([float(s) for s in [line for line in f.readlines()]])
    ground_truth = ground_truth.astype(np.float32)

    fs, x = wavfile.read(pysptk.util.example_audio_file())
    assert fs == 16000

    # Since SPTK might have memory corruption bug and the result might be
    # non-deterministic, test it with multiple time...
    for _ in range(5):
        f0 = pysptk.swipe(
            x.astype(np.float64), fs=fs, hopsize=80, min=60, max=240, otype=0
        )
        assert np.allclose(ground_truth, f0)
