# Copyright (c) 2013-2017 Louisiana State University
#               2017-2025 Cardiff University
#
# This file is part of GWpy.
#
# GWpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GWpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GWpy.  If not, see <http://www.gnu.org/licenses/>.

"""Tests for :mod:`gwpy.signal.spectral._ui`."""

import numpy
import pytest
from astropy import units
from scipy import signal

from ...testing.utils import (
    assert_array_equal,
    assert_quantity_sub_equal,
)
from ...timeseries import TimeSeries
from ..spectral import _ui as fft_ui


@pytest.mark.parametrize(("secs", "rate", "result"), [
    # all floats
    (4, 256, 1024),
    # duration as a quantity
    (1 * units.minute, 16, 960),
    # duration and rate as quantities
    (4 * units.second, 16.384 * units.kiloHertz, 65536),
])
def test_seconds_to_samples(secs, rate, result):
    """Test :func:`gwpy.signal.spectral.ui.seconds_to_samples`."""
    assert fft_ui.seconds_to_samples(secs, rate) == result


def test_normalize_fft_params():
    """Test :func:`gwpy.signal.spectral.ui.normalize_fft_params`."""
    assert fft_ui.normalize_fft_params(
        TimeSeries(numpy.zeros(1024), sample_rate=256),
    ) == {
        "nfft": 1024,
        "noverlap": 0,
    }


def test_normalize_fft_params_window_str():
    ftp = fft_ui.normalize_fft_params(
        TimeSeries(numpy.zeros(1024), sample_rate=256),
        {"window": "hann"},
    )
    win = signal.get_window("hann", 1024)
    assert ftp.pop("nfft") == 1024
    assert ftp.pop("noverlap") == 512
    assert_array_equal(ftp.pop("window"), win)
    assert not ftp


def test_normalize_fft_params_window_array():
    win = signal.get_window(("kaiser", 14), 1024)
    ftp = fft_ui.normalize_fft_params(
        TimeSeries(numpy.zeros(1024), sample_rate=256),
        {"window": win},
    )
    assert_array_equal(ftp.pop("window"), win)


@pytest.mark.requires("lal")
@pytest.mark.parametrize("win", [
    "hann",
    signal.get_window(("kaiser", 14), 1024),
])
def test_normalize_fft_params_window_lal(win):
    import lal

    from gwpy.signal.spectral._lal import welch

    ftp = fft_ui.normalize_fft_params(
        TimeSeries(numpy.zeros(1024, dtype="float32"), sample_rate=256),
        kwargs={"window": win},
        func=welch,
    )
    assert isinstance(ftp.pop("window"), lal.REAL4Window)


def test_chunk_timeseries():
    """Test :func:`gwpy.signal.spectral.ui._chunk_timeseries`."""
    a = TimeSeries(numpy.arange(400))
    chunks = list(fft_ui._chunk_timeseries(a, 100, 50))
    for i, (idxa, idxb) in enumerate([
        (None, 150),
        (75, 225),
        (175, 325),
        (250, 400),
    ]):
        assert_quantity_sub_equal(chunks[i], a[idxa:idxb])
