import os
import sys
import math
from array import array
from unittest import TestCase
from genty import genty, genty_dataset
from auditok.io import DATA_FORMAT, AudioParameterError, check_audio_data


if sys.version_info >= (3, 0):
    PYTHON_3 = True
else:
    PYTHON_3 = False


def _sample_generator(*data_buffers):
    """
    Takes a list of many mono audio data buffers and makes a sample generator
    of interleaved audio samples, one sample from each channel. The resulting
    generator can be used to build a multichannel audio buffer.
    >>> gen = _sample_generator("abcd", "ABCD")
    >>> list(gen)
    ["a", "A", "b", "B", "c", "C", "d", "D"]
    """
    frame_gen = zip(*data_buffers)
    return (sample for frame in frame_gen for sample in frame)


def _array_to_bytes(a):
    """
    Converts an `array.array` to `bytes`.
    """
    if PYTHON_3:
        return a.tobytes()
    else:
        return a.tostring()


def _generate_pure_tone(
    frequency, duration_sec=1, sampling_rate=16000, sample_width=2, volume=1e4
):
    """
    Generates a pure tone with the given frequency.
    """
    assert frequency <= sampling_rate / 2
    max_value = (2 ** (sample_width * 8) // 2) - 1
    if volume > max_value:
        volume = max_value
    fmt = DATA_FORMAT[sample_width]
    total_samples = int(sampling_rate * duration_sec)
    step = frequency / sampling_rate
    two_pi_step = 2 * math.pi * step
    data = array(
        fmt,
        (
            int(math.sin(two_pi_step * i) * volume)
            for i in range(total_samples)
        ),
    )
    return data


PURE_TONE_DICT = {
    freq: _generate_pure_tone(freq, 1, 16000, 2) for freq in (400, 800, 1600)
}
PURE_TONE_DICT.update(
    {
        freq: _generate_pure_tone(freq, 0.1, 16000, 2)
        for freq in (600, 1150, 2400, 7220)
    }
)


@genty
class TestIO(TestCase):
    @genty_dataset(
        valid_mono=(b"\0" * 113, 1, 1),
        valid_stereo=(b"\0" * 160, 1, 2),
        invalid_mono_sw_2=(b"\0" * 113, 2, 1, False),
        invalid_stereo_sw_1=(b"\0" * 113, 1, 2, False),
        invalid_stereo_sw_2=(b"\0" * 158, 2, 2, False),
    )
    def test_check_audio_data(self, data, sample_width, channels, valid=True):

        if not valid:
            with self.assertRaises(AudioParameterError):
                check_audio_data(data, sample_width, channels)
        else:
            self.assertIsNone(check_audio_data(data, sample_width, channels))

@@ -50,16 +50,71 @@
 t array%0A
+from tempfile import NamedTemporaryFile%0Aimport filecmp%0A
 from uni
@@ -197,16 +197,22 @@
  import 
+(%0A    
 DATA_FOR
@@ -215,16 +215,20 @@
 _FORMAT,
+%0A   
  AudioPa
@@ -240,16 +240,20 @@
 erError,
+%0A   
  check_a
@@ -261,16 +261,35 @@
 dio_data
+,%0A    _save_wave,%0A)
 %0A%0A%0Aif sy
@@ -2526,8 +2526,738 @@
 nnels))%0A
+%0A    @genty_dataset(%0A        mono=(%22mono_400Hz.wav%22, (400,)),%0A        three_channel=(%223channel_400-800-1600Hz.wav%22, (400, 800, 1600)),%0A    )%0A    def test_save_wave(self, filename, frequencies):%0A        filename = %22tests/data/test_16KHZ_%7B%7D%22.format(filename)%0A        sampling_rate = 16000%0A        sample_width = 2%0A        channels = len(frequencies)%0A        fmt = DATA_FORMAT%5Bsample_width%5D%0A        mono_channels = %5BPURE_TONE_DICT%5Bfreq%5D for freq in frequencies%5D%0A        data = _array_to_bytes(array(fmt, _sample_generator(*mono_channels)))%0A        tmpfile = NamedTemporaryFile()%0A        _save_wave(tmpfile.name, data, sampling_rate, sample_width, channels)%0A        self.assertTrue(filecmp.cmp(tmpfile.name, filename, shallow=False))%0A

Add tests for _save_wave