import os

import librosa
import numpy as np

from audiomentations import SpecFrequencyMask
from audiomentations.core.audio_loading_utils import load_sound_file
from demo.demo import DEMO_DIR
from .utils import plot_matrix

DEBUG = False


def test_fill_zeros():
    samples, sample_rate = load_sound_file(
        os.path.join(DEMO_DIR, "acoustic_guitar_0.wav"), sample_rate=None
    )
    magnitude_spectrogram = librosa.feature.melspectrogram(y=samples, sr=sample_rate)

    mask_fraction = 0.05
    transform = SpecFrequencyMask(
        fill_mode="constant",
        fill_constant=0.0,
        min_mask_fraction=mask_fraction,
        max_mask_fraction=mask_fraction,
        p=1.0,
    )
    augmented_spectrogram = transform(magnitude_spectrogram)

    if DEBUG:
        plot_matrix(np.log(augmented_spectrogram))

    with np.testing.assert_raises(AssertionError):
        np.testing.assert_array_equal(augmented_spectrogram, magnitude_spectrogram)

    num_zeroed_frequencies = 0
    for i in range(augmented_spectrogram.shape[0]):
        if sum(augmented_spectrogram[i]) == 0.0:
            num_zeroed_frequencies += 1

    assert num_zeroed_frequencies == int(
        round(magnitude_spectrogram.shape[0] * mask_fraction)
    )


def test_fill_zeros_multichannel():
    samples, sample_rate = load_sound_file(
        os.path.join(DEMO_DIR, "background_noises", "hens.ogg"),
        sample_rate=None,
        mono=False,
    )
    assert samples.shape[0] == 2
    magnitude_spectrogram_chn0 = librosa.feature.melspectrogram(
        y=np.asfortranarray(samples[0, :]), sr=sample_rate
    )
    magnitude_spectrogram_chn1 = librosa.feature.melspectrogram(
        y=np.asfortranarray(samples[1, :]), sr=sample_rate
    )
    multichannel_magnitude_spectrogram = np.zeros(
        shape=(
            magnitude_spectrogram_chn0.shape[0],
            magnitude_spectrogram_chn0.shape[1],
            3,
        ),
        dtype=np.float32,
    )
    multichannel_magnitude_spectrogram[:, :, 0] = magnitude_spectrogram_chn0
    multichannel_magnitude_spectrogram[:, :, 1] = magnitude_spectrogram_chn1
    multichannel_magnitude_spectrogram[:, :, 2] = magnitude_spectrogram_chn1

    mask_fraction = 0.05
    transform = SpecFrequencyMask(
        fill_mode="constant",
        fill_constant=0.0,
        min_mask_fraction=mask_fraction,
        max_mask_fraction=mask_fraction,
        p=1.0,
    )
    augmented_spectrogram = transform(multichannel_magnitude_spectrogram)

    if DEBUG:
        image = (7 + np.log10(augmented_spectrogram + 0.0000001)) / 8
        plot_matrix(image)

    with np.testing.assert_raises(AssertionError):
        np.testing.assert_array_equal(
            augmented_spectrogram, multichannel_magnitude_spectrogram
        )

    num_zeroed_frequencies = 0
    for i in range(augmented_spectrogram.shape[0]):
        if np.sum(augmented_spectrogram[i]) == 0.0:
            num_zeroed_frequencies += 1

    assert num_zeroed_frequencies == int(
        round(multichannel_magnitude_spectrogram.shape[0] * mask_fraction)
    )


def test_fill_mean():
    samples, sample_rate = load_sound_file(
        os.path.join(DEMO_DIR, "acoustic_guitar_0.wav"), sample_rate=None
    )
    magnitude_spectrogram = librosa.feature.melspectrogram(y=samples, sr=sample_rate)

    min_mask_fraction = 0.05
    max_mask_fraction = 0.09
    transform = SpecFrequencyMask(
        fill_mode="mean",
        min_mask_fraction=min_mask_fraction,
        max_mask_fraction=max_mask_fraction,
        p=1.0,
    )
    augmented_spectrogram = transform(magnitude_spectrogram)

    if DEBUG:
        plot_matrix(np.log(augmented_spectrogram))

    num_masked_frequencies = 0
    for i in range(augmented_spectrogram.shape[0]):
        frequency_slice = augmented_spectrogram[i]
        if (
            np.amin(frequency_slice) == np.amax(frequency_slice)
            and sum(frequency_slice) != 0.0
        ):
            num_masked_frequencies += 1

    assert num_masked_frequencies >= int(
        round(magnitude_spectrogram.shape[0] * min_mask_fraction)
    )
    assert num_masked_frequencies <= int(
        round(magnitude_spectrogram.shape[0] * max_mask_fraction)
    )


def test_fill_mean_multichannel():
    samples, sample_rate = load_sound_file(
        os.path.join(DEMO_DIR, "background_noises", "hens.ogg"),
        sample_rate=None,
        mono=False,
    )
    assert samples.shape[0] == 2
    magnitude_spectrogram_chn0 = librosa.feature.melspectrogram(
        y=np.asfortranarray(samples[0, :]), sr=sample_rate
    )
    magnitude_spectrogram_chn1 = librosa.feature.melspectrogram(
        y=np.asfortranarray(samples[1, :]), sr=sample_rate
    )
    multichannel_magnitude_spectrogram = np.zeros(
        shape=(
            magnitude_spectrogram_chn0.shape[0],
            magnitude_spectrogram_chn0.shape[1],
            3,
        ),
        dtype=np.float32,
    )
    multichannel_magnitude_spectrogram[:, :, 0] = magnitude_spectrogram_chn0
    multichannel_magnitude_spectrogram[:, :, 1] = magnitude_spectrogram_chn1
    multichannel_magnitude_spectrogram[:, :, 2] = magnitude_spectrogram_chn1

    mask_fraction = 0.05
    transform = SpecFrequencyMask(
        fill_mode="mean",
        min_mask_fraction=mask_fraction,
        max_mask_fraction=mask_fraction,
        p=1.0,
    )
    augmented_spectrogram = transform(multichannel_magnitude_spectrogram)

    if DEBUG:
        image = (7 + np.log10(augmented_spectrogram + 0.0000001)) / 8
        plot_matrix(image)

    with np.testing.assert_raises(AssertionError):
        np.testing.assert_array_equal(
            augmented_spectrogram, multichannel_magnitude_spectrogram
        )

    num_masked_frequencies = 0
    for i in range(augmented_spectrogram.shape[0]):
        frequency_slice = augmented_spectrogram[i]
        if (
            np.amin(frequency_slice) == np.amax(frequency_slice)
            and np.sum(frequency_slice) != 0.0
        ):
            num_masked_frequencies += 1

    assert num_masked_frequencies == int(
        round(multichannel_magnitude_spectrogram.shape[0] * mask_fraction)
    )
