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

import numpy as np

from quicksong import features


def total_power(frame: np.ndarray) -> float:
    psd = np.abs(frame) ** 2
    return np.log10(np.sum(psd))


def wiener_entropy(frame: np.ndarray) -> float:
    psd = np.abs(frame) ** 2
    return np.log10(np.exp(np.mean(np.log(psd))) / np.mean(psd))


class TestCythonImpl(unittest.TestCase):
    def test_total_power(self):
        input = np.random.randn(1000) + np.random.randn(1000) * 1j
        self.assertTrue(np.allclose(total_power(input), features.total_power(input)))

    def test_wiener_entropy(self):
        input = np.random.randn(1000) + np.random.randn(1000) * 1j
        self.assertTrue(
            np.allclose(wiener_entropy(input), features.wiener_entropy(input))
        )
