import os

import numpy as np
import pytest


def _metrics_floor(x, sr):
    from producer import meters

    return meters.noise_floor_db(x, sr)


@pytest.mark.slow
def test_dfn3_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("df")
    from producer.engines import denoise_dfn

    x = noisy_speech[: sr * 4]
    y, eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu")
    assert "dfn" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0
    assert np.corrcoef(x, y.astype(np.float64))[0, 1] > 0.9


@pytest.mark.slow
def test_zipenhancer_reduces_noise(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("zipenhancer")
    from producer.engines import denoise_zip

    x = noisy_speech[: sr * 4]
    y, eng, _ = denoise_zip.denoise(x, sr, 1.0, "cpu")
    assert "zipenhancer" in eng
    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0


@pytest.mark.slow
def test_mossformer2_enhances(sr, noisy_speech):
    pytest.importorskip("torch")
    pytest.importorskip("clearvoice")
    from producer.engines import enhance_mossformer

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_mossformer.enhance(x, sr, 1.0, "cpu")
    assert "mossformer2" in eng
    assert np.all(np.isfinite(y))


@pytest.mark.slow
def test_resemble_enhance(sr, noisy_speech):
    if not os.environ.get("PRODUCER_TEST_RESEMBLE"):
        pytest.skip("set PRODUCER_TEST_RESEMBLE=1 to run the isolated-venv generative engine")
    from producer.engines import enhance_resemble

    x = noisy_speech[: sr * 4]
    y, eng, _ = enhance_resemble.enhance(x, sr, 1.0, "cpu")
    assert "resemble" in eng
    assert np.all(np.isfinite(y))