# -*- mode: python -*-
import logging
import unittest

import numpy as np

from quicksong import core
from quicksong.streaming import (
    STFT,
    IncompatibleBlockSize,
    OverlapSaveConvolver,
)

log = logging.getLogger("quicksong")


class TestConvolution(unittest.TestCase):

    signal = np.random.randn(5000)

    def test_bad_block_size(self):
        window_sizes = (11, 21, 51)
        kernels = core.make_hanning_kernels(window_sizes)
        convolver = OverlapSaveConvolver(kernels)
        dummy = np.zeros(convolver.block_size * 2)
        with self.assertRaises(IncompatibleBlockSize):
            convolver.process(dummy)

    def test_convolve(self):
        """Output of running convolution is the same as np.convolve"""
        window_sizes = (11, 20, 51)
        kernels = core.make_hanning_kernels(window_sizes)
        nt, nk = kernels.shape
        self.assertEqual(nt, max(window_sizes))
        self.assertEqual(nk, len(window_sizes))

        convolver = OverlapSaveConvolver(kernels)
        out = []
        convolver.prefill(self.signal[: convolver.block_size])
        for i in range(convolver.block_size, self.signal.size, convolver.block_size):
            s = self.signal[i : i + convolver.block_size]
            out.append(convolver.process(s))
        out.append(convolver.close())
        out = np.vstack(out)
        nout_t, nout_k = out.shape
        self.assertEqual(nout_k, nk)
        self.assertEqual(nout_t, self.signal.size)
        for i in range(nk):
            y = np.convolve(self.signal, kernels[:, i], mode="same")
            self.assertTrue(np.allclose(out[:, i], y))

    def test_copy_convolver(self):
        from copy import copy

        window_sizes = (11, 21, 51)
        kernels = core.make_hanning_kernels(window_sizes)
        conv1 = OverlapSaveConvolver(kernels)
        conv2 = copy(conv1)
        self.assertTrue(np.allclose(conv1.kernels, conv2.kernels))
        self.assertAlmostEqual(conv2.buffer.sum(), 0.0)


class TestSTFT(unittest.TestCase):

    sampling_rate = 48000
    frequency_range = (500, 20000)
    signal = np.random.randn(50000)

    def test_block_too_big(self):
        speccr = STFT(self.sampling_rate, 0.0020, 0.0005, self.frequency_range)
        dummy = np.zeros(speccr.block_size * 2)
        with self.assertRaises(IncompatibleBlockSize):
            speccr.process(dummy)

    def test_transform(self):
        """Output of running convolution is the same as np.convolve"""
        from math import ceil

        out = []
        speccr = STFT(self.sampling_rate, 0.0020, 0.0005, self.frequency_range)
        speccr.prefill(self.signal[: speccr.prefill_size])
        for i in range(speccr.prefill_size, self.signal.size, speccr.block_size):
            s = self.signal[i : i + speccr.block_size]
            out.append(speccr.process(s))
        out = np.column_stack(out)
        nf, nt = out.shape
        self.assertEqual(nf, speccr.findx.size)
        expected = ceil((self.signal.size - speccr.prefill_size) / speccr.block_size)
        self.assertEqual(nt, expected)


