"""Effects (ackit/fx.py): each block measured back on synthetic signals."""
import numpy as np
import pytest

from ackit import fx

SR = 44100


def tone(freq, seconds=1.0, amp=0.3):
    t = np.arange(int(SR * seconds)) / SR
    return (amp * np.sin(2 * np.pi * freq * t)).astype(np.float32)[:, None]


def impulse(seconds=1.5, at=0.1):
    x = np.zeros((int(SR * seconds), 1), np.float32)
    x[int(SR * at)] = 1.0
    return x


def only(key, **params):
    s = fx.defaults()
    s[key]["on"] = True
    s[key].update(params)
    return s


def peak_freq(x, lo=20, hi=None):
    spec = np.abs(np.fft.rfft(x[:, 0] * np.hanning(len(x))))
    f = np.fft.rfftfreq(len(x), 1 / SR)
    m = (f >= lo) & (f <= (hi or SR / 2))
    return f[m][np.argmax(spec[m])]


def band_db(x, f0, width=0.05):
    spec = np.abs(np.fft.rfft(x[:, 0] * np.hanning(len(x)))) ** 2
    f = np.fft.rfftfreq(len(x), 1 / SR)
    m = (f > f0 * (1 - width)) & (f < f0 * (1 + width))
    return 10 * np.log10(spec[m].sum() + 1e-20)


def test_defaults_off_and_identity():
    s = fx.defaults()
    assert not fx.active(s) and fx.summary(s) == ""
    x = tone(440)
    assert np.array_equal(fx.render(x, SR, s), x)


def test_normalize_fills_clamps_and_drops():
    s = fx.normalize({"eq": {"on": 1, "low_cut": 5, "bogus": 3}, "pitch": {"shift": "x"}, "nope": {}})
    assert s["eq"]["on"] is True and s["eq"]["low_cut"] == 20 and "bogus" not in s["eq"]
    assert s["pitch"]["shift"] == 0.0 and "nope" not in s
    assert s["glitch"]["repeats"] == 2 and isinstance(s["glitch"]["repeats"], int)
    assert fx.summary(s) == "Band-pass EQ"


def test_pitch_shift_octave_up():
    y = fx.render(tone(220), SR, only("pitch", shift=12))
    assert abs(peak_freq(y) - 440) < 10


def test_formant_shift_keeps_pitch():
    t = np.arange(SR) / SR
    x = sum(np.sin(2 * np.pi * 150 * k * t) / k for k in range(1, 20)).astype(np.float32)[:, None] * 0.1
    y = fx.render(x, SR, only("pitch", formant=6))
    assert abs(peak_freq(y, 100, 200) - 150) < 5
    assert not np.allclose(y, x, atol=1e-3)


def test_pitch_drift_moves_over_time():
    y = fx.render(tone(440, 4), SR, only("pitch", drift=4, drift_rate=0.5))
    peaks = [peak_freq(y[i:i + 8192]) for i in range(SR // 2, len(y) - 8192, SR // 4)]
    assert max(peaks) - min(peaks) > 100


def test_octave_layer_below():
    y = fx.render(tone(440), SR, only("octaves", interval1=-12, level1=0))
    assert band_db(y, 220) > band_db(tone(440), 220) + 30


def test_ring_modulator_sidebands():
    y = fx.render(tone(1000), SR, only("robot", ring_freq=100, ring_mix=1.0))
    assert band_db(y, 900, 0.01) > band_db(y, 1000, 0.01) + 20
    assert band_db(y, 1100, 0.01) > band_db(y, 1000, 0.01) + 20


def test_glitch_longer_and_seeded():
    x = tone(300, 3)
    a = fx.render(x, SR, only("glitch", rate=5, length=100, repeats=2, pitch=0, seed=1))
    b = fx.render(x, SR, only("glitch", rate=5, length=100, repeats=2, pitch=0, seed=1))
    c = fx.render(x, SR, only("glitch", rate=5, length=100, repeats=2, pitch=0, seed=2))
    assert np.array_equal(a, b)
    assert len(a) > len(x) and len(a) - len(x) == int(SR * 0.1) * 2 * round(5 * 3 / 10)
    assert not np.array_equal(a, c[:len(a)]) or len(a) != len(c)


def test_eq_band_pass():
    x = (tone(100) + tone(1500) + tone(8000)) / 3
    y = fx.render(x, SR, only("eq", low_cut=300, high_cut=3400, slope=24, mid_gain=0))
    assert band_db(y, 100) < band_db(x, 100) - 20
    assert band_db(y, 8000) < band_db(x, 8000) - 20
    assert abs(band_db(y, 1500) - band_db(x, 1500)) < 0.5


def test_distortion_adds_harmonics_same_loudness():
    x = tone(200)
    y = fx.render(x, SR, only("distortion", drive=20))
    assert band_db(y, 600) > band_db(x, 600) + 20
    assert abs(np.sqrt(np.mean(y ** 2)) - np.sqrt(np.mean(x ** 2))) < 1e-3


@pytest.mark.parametrize("opt", [{"clip": -12, "drive": 0}, {"bits": 4, "drive": 0}, {"gsm": True, "drive": 0}])
def test_distortion_options_change_audio(opt):
    x = tone(200)
    assert not np.allclose(fx.render(x, SR, only("distortion", **opt)), x, atol=1e-3)


def test_noise_level_relative_to_voice():
    x = tone(500, 2)
    y = fx.render(x, SR, only("noise", level=-20, colour=2000))
    n = y - x
    ratio = 20 * np.log10(np.sqrt(np.mean(n ** 2)) / np.sqrt(np.mean(x ** 2)))
    assert abs(ratio + 20) < 1


def test_echo_delay_and_tail():
    x = impulse()
    y = fx.render(x, SR, only("echo", time=160, feedback=0, mix=0.5))
    assert len(y) >= len(x)
    d = np.argsort(np.abs(y[:, 0]))[-2:]
    assert abs(sorted(d)[1] - sorted(d)[0] - int(SR * 0.16)) <= 1      # pedalboard rounds the delay


def test_reverb_tail_and_reverse_swell():
    x = impulse()
    y = fx.render(x, SR, only("reverb", mix=0.5, decay=1.0))
    assert np.sum(y[int(SR * 0.3):] ** 2) > 1e-4
    v = np.zeros((SR * 2, 1), np.float32)
    v[SR:SR + SR // 2] = tone(300, 0.5)
    r = fx.render(v, SR, only("reverb", mix=0.0, reverse=1.0, reverse_every=1.0))
    assert np.sum(r[SR // 2:SR] ** 2) > 1e-3 and np.sum(v[SR // 2:SR] ** 2) == 0


def test_channels_pa_layout():
    x = impulse(1.0)
    s = only("channels", front=-12, surround=-6, surround_delay=160, width=0)
    y = fx.render(x, SR, s, channels=5)                # L R C SL SR
    assert y.shape[1] == 5
    i = int(SR * 0.1)
    assert y[i, 2] == pytest.approx(1.0)
    assert y[i, 0] == pytest.approx(10 ** (-12 / 20), rel=1e-3)
    assert np.argmax(np.abs(y[:, 3])) == i + int(SR * 0.16)
    assert np.max(np.abs(y[:, 3])) == pytest.approx(10 ** (-6 / 20), rel=1e-3)
    assert fx.render(x, SR, s, channels=1).shape[1] == 1     # mono target: block has no effect


def test_channels_width_and_pan():
    y = fx.render(impulse(1.0), SR, only("channels", width=10), channels=2)
    assert np.argmax(np.abs(y[:, 1])) - np.argmax(np.abs(y[:, 0])) == int(SR * 0.01)
    p = fx.render(tone(300, 2), SR, only("channels", pan_rate=1, pan_depth=1), channels=2)
    left = [np.sqrt(np.mean(p[i:i + 2000, 0] ** 2)) for i in range(0, len(p) - 2000, 4410)]
    assert max(left) > 2 * min(left) + 1e-3


def test_all_blocks_together_finite():
    s = fx.defaults()
    for b in fx.BLOCKS:
        s[b.key]["on"] = True
    s["distortion"]["gsm"] = True
    x = tone(220, 2) + tone(1300, 2) * 0.3
    y = fx.render(x, SR, s, channels=4)
    assert y.shape[1] == 4 and np.isfinite(y).all() and len(y) >= len(x)
    assert np.array_equal(y, fx.render(x, SR, s, channels=4))     # deterministic


def rt60(ir):
    """Reverb time from the Schroeder backward integral (-5 to -25 dB, extrapolated to -60)."""
    e = np.cumsum((ir ** 2)[::-1])[::-1]
    db = 10 * np.log10(e / e[0] + 1e-30)
    a, b = np.argmax(db < -5), np.argmax(db < -25)
    return (b - a) / SR * 3


@pytest.mark.parametrize("decay", [0.3, 1.0, 3.0])
def test_room_response_decay_measured_back(decay):
    ir = fx.room_response(SR, 1, decay, predelay=0, early=0.0, damping=0.0)
    assert abs(rt60(ir[:, 0]) - decay) < 0.15 * decay
    assert np.sum(ir ** 2) == pytest.approx(1.0, rel=1e-3)


def test_room_response_damping_predelay_width():
    ir = fx.room_response(SR, 1, 1.5, predelay=0, early=0.0, damping=1.0)[:, 0]
    late = ir[int(0.5 * SR):][:, None]
    assert band_db(late, 8000, 0.3) < band_db(late, 300, 0.3) - 20       # the highs die first
    pre = fx.room_response(SR, 1, 0.5, predelay=40)[:, 0]
    assert np.max(np.abs(pre[:int(0.039 * SR)])) == 0 and np.max(np.abs(pre[int(0.04 * SR):])) > 0
    wide = fx.room_response(SR, 2, 0.5, width=1.0)
    mono = fx.room_response(SR, 2, 0.5, width=0.0)
    assert abs(np.corrcoef(wide.T)[0, 1]) < 0.2 and np.corrcoef(mono.T)[0, 1] > 0.99


def test_presets_and_looks_are_valid():
    for block, items in fx.PRESETS.items():
        keys = {p.key for p in fx.BY_KEY[block].params}
        for name, values in items:
            assert set(values) <= keys, (block, name)
            s = fx.with_preset(fx.defaults(), block, name)
            assert s[block]["on"] and all(s[block][k] == values[k] for k in values), (block, name)
    x = tone(300, 1)
    for name, _ in fx.LOOKS:
        s = fx.look(name)
        assert fx.active(s)
        y = fx.render(x, SR, s, channels=5)
        assert np.isfinite(y).all() and len(y) >= len(x)


def test_reverb_mix_dry_to_wet_and_old_name():
    x = impulse(1.0)
    i = int(SR * 0.1)
    assert np.array_equal(fx.render(x, SR, only("reverb", mix=0.0)), x)                   # 0: dry only
    wet = fx.render(x, SR, only("reverb", mix=1.0, predelay=30, early=0.0))
    assert abs(wet[i, 0]) < 1e-3 and np.max(np.abs(wet[i + int(SR * 0.03):])) > 0       # 1: no dry impulse left
    half = fx.render(x, SR, only("reverb", mix=0.5, predelay=30, early=0.0))
    assert half[i, 0] == pytest.approx(np.cos(np.pi / 4), rel=1e-3)                      # equal power
    assert fx.normalize({"reverb": {"on": True, "wet": 0.4}})["reverb"]["mix"] == 0.4    # older settings