"""Tests for matched filtering module."""

import numpy as np
import pytest

from matched import (matched_filter, matched_filter_fft, make_chirp,
                     simulate_bat_echo, detect_echo)


# --- Basic matched filter ---

def test_perfect_match_peak_at_zero():
    """A template matched against itself should peak at the template position."""
    template = np.array([1.0, 0.5, -0.5, -1.0])
    output = matched_filter(template, template, normalize=True)
    peak = np.argmax(output)
    # Peak should be at position len(template) - 1 (full correlation)
    assert peak == len(template) - 1

def test_normalized_peak_is_one():
    """Normalized matched filter output should peak at 1.0 for a perfect match."""
    template = np.array([1.0, -1.0, 1.0, -1.0])
    output = matched_filter(template, template, normalize=True)
    assert abs(np.max(output) - 1.0) < 1e-10

def test_detects_shifted_template():
    """Should detect a template embedded at a known offset."""
    template = np.array([1.0, 0.5, -0.5, -1.0])
    signal = np.zeros(100)
    offset = 40
    signal[offset:offset + len(template)] = template
    output = matched_filter(template, signal, normalize=True)
    peak = np.argmax(output)
    expected_peak = offset + len(template) - 1
    assert peak == expected_peak

def test_detects_in_noise():
    """Should detect a template buried in moderate noise."""
    rng = np.random.default_rng(0)
    template = np.sin(2 * np.pi * np.linspace(0, 1, 50))
    signal = np.zeros(500)
    offset = 200
    signal[offset:offset + len(template)] = template
    signal += 0.5 * rng.standard_normal(500)  # moderate noise
    output = matched_filter(template, signal, normalize=True)
    peak = np.argmax(output)
    expected_peak = offset + len(template) - 1
    assert abs(peak - expected_peak) <= 1


# --- FFT-based matched filter ---

def test_fft_matches_direct():
    """FFT-based output should match direct correlation (approximately)."""
    rng = np.random.default_rng(1)
    template = rng.standard_normal(20)
    signal = rng.standard_normal(200)
    direct = matched_filter(template, signal, normalize=True)
    fft_out = matched_filter_fft(template, signal, normalize=True)
    # FFT version returns len(signal) samples; direct returns len(signal)+len(template)-1
    # Compare the overlapping portion
    np.testing.assert_allclose(fft_out[:len(signal)], direct[len(template)-1:], atol=1e-10)

def test_fft_detects_shifted():
    """FFT matched filter should detect shifted template."""
    template = np.array([1.0, -1.0, 1.0])
    signal = np.zeros(100)
    signal[50:53] = template
    output = matched_filter_fft(template, signal, normalize=True)
    peak = np.argmax(output)
    assert peak == 50  # FFT version peaks at the start of the match


# --- Chirp generation ---

def test_chirp_length():
    t, chirp = make_chirp(100, 200, 0.01, 10000)
    assert len(chirp) == 100
    assert len(t) == 100

def test_chirp_amplitude():
    _, chirp = make_chirp(100, 200, 0.1, 10000)
    assert np.max(np.abs(chirp)) <= 1.0 + 1e-10


# --- Bat echolocation simulation ---

def test_bat_echo_delay():
    """Matched filter should estimate the correct echo delay."""
    fs = 250000
    target_distance = 2.0
    expected_delay = 2 * target_distance / 343.0

    t, tx, rx, true_delay = simulate_bat_echo(
        fs=fs, target_distance=target_distance, snr_db=0.0,
        rng=np.random.default_rng(42)
    )

    # Use the chirp portion as template
    n_chirp = int(0.002 * fs)
    template = tx[:n_chirp]
    est_delay, _ = detect_echo(template, rx, fs)

    # Allow 2-sample tolerance
    assert abs(est_delay - expected_delay) < 2 / fs + 1e-6

def test_bat_echo_negative_snr():
    """Should still detect echo at -10 dB SNR."""
    fs = 250000
    target_distance = 3.0
    expected_delay = 2 * target_distance / 343.0

    t, tx, rx, true_delay = simulate_bat_echo(
        fs=fs, target_distance=target_distance, snr_db=-10.0,
        rng=np.random.default_rng(99)
    )

    n_chirp = int(0.002 * fs)
    template = tx[:n_chirp]
    est_delay, _ = detect_echo(template, rx, fs)

    # Looser tolerance at low SNR
    assert abs(est_delay - expected_delay) < 5 / fs + 1e-6

def test_bat_echo_components():
    """Verify simulation output shapes and delay value."""
    t, tx, rx, delay = simulate_bat_echo(target_distance=1.5)
    assert len(t) == len(tx) == len(rx)
    assert delay == pytest.approx(2 * 1.5 / 343.0, rel=1e-3)


# --- Edge cases ---

def test_zero_template():
    """Zero template should not crash (normalize handles zero energy)."""
    template = np.zeros(10)
    signal = np.random.default_rng(0).standard_normal(50)
    output = matched_filter(template, signal, normalize=True)
    assert np.all(output == 0)

def test_single_sample_template():
    template = np.array([1.0])
    signal = np.array([0.0, 0.0, 3.0, 0.0])
    output = matched_filter(template, signal, normalize=True)
    assert np.argmax(output) == 2