# -*- mode: python -*-
import logging
import unittest
from math import ceil

import numpy as np

from quicksong import core
from quicksong.features import total_power, wiener_entropy
from quicksong.streaming import STFT, OverlapSaveConvolver

log = logging.getLogger("quicksong")


class TestFeatureExtraction(unittest.TestCase):
    window_sizes = (11, 21, 51)
    sampling_rate = 48000
    window_ms = 20.0
    shift_ms = 10.0
    frequency_range = (500, 20000)
    signal = np.random.randn(20000)

    def make_extractor(self):
        speccr = STFT(
            self.sampling_rate,
            self.window_ms / 1000,
            self.shift_ms / 1000,
            self.frequency_range,
        )
        kernels = core.make_hanning_kernels(self.window_sizes)
        convolver = OverlapSaveConvolver(kernels)
        features = (total_power, wiener_entropy)
        return core.FeatureExtractor(speccr, features, convolver)

    def test_signal_duration(self):
        extractor_1 = self.make_extractor()
        out = []
        for block in extractor_1.process(self.signal):
            out.append(block)
        out.append(extractor_1.close())
        features_1 = np.vstack(out)
        expected = ceil(self.signal.size / extractor_1.frame_size)
        self.assertEqual(features_1.shape[0], expected)

    def test_signal_split(self):
        """output should be the same no matter how the input is divided up"""
        extractor_1 = self.make_extractor()
        out = []
        for block in extractor_1.process(self.signal):
            out.append(block)
        out.append(extractor_1.close())
        features_1 = np.vstack(out)

        extractor_2 = self.make_extractor()
        out = []
        chunk_size = self.signal.size // 5
        for i in range(0, self.signal.size, chunk_size):
            chunk = self.signal[i : i + chunk_size]
            for block in extractor_2.process(chunk):
                out.append(block)
        out.append(extractor_2.close())
        features_2 = np.vstack(out)

        self.assertSequenceEqual(features_1.shape, features_2.shape)
        self.assertTrue(np.allclose(features_1, features_2))


class TestIntervalFinder(unittest.TestCase):
    start = np.array(
        [
            59,
            101,
            392,
            479,
            532,
            820,
            836,
            989,
            1087,
            1276,
            1355,
            1670,
            1753,
            1921,
            2053,
            2174,
            2483,
            2711,
            2773,
            2847,
            2935,
            2943,
            3330,
            3639,
            4151,
        ]
    )
    lengths = np.array(
        [
            2,
            8,
            9,
            11,
            11,
            2,
            7,
            5,
            14,
            10,
            7,
            10,
            10,
            28,
            10,
            5,
            7,
            19,
            4,
            3,
            10,
            178,
            300,
            20,
            1,
        ]
    )
    max_gap = 100
    max_interval = 500
    expected = np.array([[2935, 3120], [3330, 3658]])

    def setUp(self):
        self.arr = np.ones(5000, dtype="i")
        for start, len in zip(self.start, self.lengths):
            self.arr[start : start + len] = 0

    def test_find_intervals(self):
        ifd = core.IntervalFinder(100, 500, 10)
        intervals = np.vstack(list(ifd.process(self.arr)))
        self.assertEqual(intervals.shape, self.expected.shape)
        log.debug("expected: %s\n, got: %s", self.expected, intervals)
        self.assertTrue(np.allclose(intervals, self.expected))

    def test_find_intervals_discontinuous(self):
        ifd = core.IntervalFinder(100, 500, 10)
        ivals = []
        block_size = 1000
        for i in range(0, self.arr.size, block_size):
            block = self.arr[i : i + block_size]
            for interval in ifd.process(block):
                ivals.append(interval)
        intervals = np.vstack(ivals)
        self.assertEqual(intervals.shape, self.expected.shape)
        log.debug("expected: %s\n, got: %s", self.expected, intervals)
        self.assertTrue(np.allclose(intervals, self.expected))


class TestPadIntervals(unittest.TestCase):
    def test_pad_single_interval(self):
        intervals = [(1000, 2000)]
        self.assertSequenceEqual(
            list(core.pad_intervals(intervals, 1000, 500)), [(0, 2500)]
        )

    def test_merge_overlapping_intervals(self):
        intervals = [(1000, 2000), (2500, 3000)]
        self.assertSequenceEqual(
            list(core.pad_intervals(intervals, 1000, 500)), [(0, 3500)]
        )

    def test_dont_merge_nonoverlapping_intervals(self):
        intervals = [(1000, 2000), (2500, 3000)]
        self.assertSequenceEqual(
            list(core.pad_intervals(intervals, 100, 200)), [(900, 2200), (2400, 3200)]
        )
