import io
import unittest
import numpy as np
from pluto_fm.dsp import Modulator, AUDIO_RATE, IQ_RATE
from pluto_fm.cli import blocks


class StreamingTests(unittest.TestCase):
    def test_chunk_continuity_and_deviation(self):
        audio = 0.01 * np.sin(2 * np.pi * 1000 * np.arange(9600) / AUDIO_RATE)
        whole = Modulator().process(audio)
        mod = Modulator()
        split = np.concatenate([mod.process(audio[:4800]), mod.process(audio[4800:])])
        np.testing.assert_allclose(whole, split, atol=0.002)
        frequency = np.angle(split[1:] * split[:-1].conj()) * IQ_RATE / (2 * np.pi)
        self.assertLessEqual(np.max(np.abs(frequency)), 75001)
        np.testing.assert_allclose(np.abs(split), 0.8 * 2**14, atol=0.003)
        # Demodulated tone has the requested frequency, after filter settling.
        spectrum = np.abs(np.fft.rfft(frequency[10000:]))
        peak = np.argmax(spectrum) * IQ_RATE / len(frequency[10000:])
        self.assertLess(abs(peak - 1000), 10)

    def test_stereo_separation_and_pilot(self):
        from scipy import signal
        t = np.arange(AUDIO_RATE) / AUDIO_RATE
        audio = np.column_stack((0.2 * np.sin(2 * np.pi * 1000 * t),
                                 0.2 * np.sin(2 * np.pi * 2300 * t)))
        iq = Modulator(volume=1, emphasis=0, stereo=True).process(audio)
        mpx = np.angle(iq[1:] * iq[:-1].conj()) * IQ_RATE / (2 * np.pi * 75000)
        clock = np.arange(1, len(iq)) / IQ_RATE
        lowpass = signal.butter(8, 15000, fs=IQ_RATE, output='sos')
        summed = signal.sosfiltfilt(lowpass, mpx)
        difference = signal.sosfiltfilt(lowpass, 2 * mpx * np.sin(2 * np.pi * 38000 * clock))
        left, right = summed + difference, summed - difference
        section = slice(IQ_RATE // 10, -IQ_RATE // 10)
        def amplitude(samples, frequency):
            return abs(2 * np.mean(samples[section] * np.exp(-2j * np.pi * frequency * clock[section])))
        self.assertAlmostEqual(amplitude(mpx, 19000), 0.1, places=4)
        self.assertLess(amplitude(mpx, 38000), 0.001)
        self.assertGreater(amplitude(left, 1000), 0.17)
        self.assertGreater(amplitude(right, 2300), 0.17)
        self.assertLess(amplitude(left, 2300), 0.001)
        self.assertLess(amplitude(right, 1000), 0.001)
        self.assertLessEqual(np.max(np.abs(mpx)), 1.0001)

    def test_stereo_chunk_continuity(self):
        audio = np.random.default_rng(3).normal(0, 0.05, (9600, 2))
        expected = Modulator(stereo=True).process(audio)
        mod = Modulator(stereo=True)
        actual = np.concatenate([mod.process(audio[:4799]), mod.process(audio[4799:])])
        np.testing.assert_allclose(actual, expected, atol=0.002)

    def test_stereo_final_padding(self):
        audio = np.array([[1, 2], [3, 4]], dtype='<f4')
        block = next(blocks(io.BytesIO(audio.tobytes()), channels=2))
        self.assertEqual(block.shape, (4800, 2))
        np.testing.assert_array_equal(block[:2], audio)
        self.assertEqual(np.count_nonzero(block[2:]), 0)

    def test_final_block_padding(self):
        data = list(blocks(io.BytesIO(np.ones(4801, dtype='<f4').tobytes())))
        self.assertEqual(len(data), 2)
        self.assertEqual(data[1][0], 1)
        self.assertEqual(np.count_nonzero(data[1]), 1)


if __name__ == '__main__':
    unittest.main()