import numpy as np
import pytest


def _build_map_and_wdm(seed=20260223):
    from wdm_wavelet.wdm import WDM

    rng = np.random.default_rng(seed)
    signal = rng.standard_normal(4096)

    wdm = WDM(16, 32, 6, 8)
    tf_map = wdm.t2w(signal, sample_rate=4096.0, t0=0.0)
    return wdm, tf_map


def test_set_td_filter_shapes_and_metadata():
    wdm, _ = _build_map_and_wdm()
    n_coeffs = 8
    L = 2
    td = wdm.set_td_filter(n_coeffs=n_coeffs, L=L)

    J = wdm.M * L
    assert td.M == wdm.M
    assert td.L == L
    assert td.n_coeffs == n_coeffs
    assert td.T0.shape == (2 * J + 1, 2 * n_coeffs + 1)
    assert td.Tx.shape == (2 * J + 1, 2 * n_coeffs + 1)
    assert td.sinTD.shape == (2 * J,)
    assert td.cosTD.shape == (2 * J,)
    assert td.sinTDx.shape == (4 * J,)


def test_get_td_amp_power_consistency():
    wdm, tf_map = _build_map_and_wdm()
    wdm.set_td_filter(n_coeffs=8, L=2)

    n = tf_map.n_time // 2
    m = wdm.M // 2
    pixel_index = n * (wdm.M + 1) + m

    for delay_index in (-11, -3, 0, 5, 12):
        a00 = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="a")
        a90 = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="A")
        power = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="p")
        np.testing.assert_allclose(power, 0.5 * (a00 * a00 + a90 * a90), rtol=1e-12, atol=1e-12)


def test_get_td_vec_matches_scalar_api():
    wdm, tf_map = _build_map_and_wdm()
    wdm.set_td_filter(n_coeffs=8, L=2)

    n = tf_map.n_time // 2
    m = wdm.M // 2
    pixel_index = n * (wdm.M + 1) + m
    K = 7

    vec_power = wdm.get_td_vec(tf_map, pixel_index, K, mode="p")
    assert vec_power.shape == (2 * K + 1,)

    delays = list(range(-K, K + 1))
    for idx, delay_index in enumerate(delays):
        scalar = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="p")
        np.testing.assert_allclose(vec_power[idx], scalar, rtol=1e-12, atol=1e-12)

    vec_amp = wdm.get_td_vec(tf_map, pixel_index, K, mode="a")
    assert vec_amp.shape == (2 * (2 * K + 1),)
    vec_a00 = vec_amp[: 2 * K + 1]
    vec_a90 = vec_amp[2 * K + 1 :]

    for idx, delay_index in enumerate(delays):
        a00 = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="a")
        a90 = wdm.get_td_amp(tf_map, pixel_index, delay_index, mode="A")
        np.testing.assert_allclose(vec_a00[idx], a00, rtol=1e-12, atol=1e-12)
        np.testing.assert_allclose(vec_a90[idx], a90, rtol=1e-12, atol=1e-12)


def test_td_rejects_power_only_tf_map():
    from wdm_wavelet.wdm import WDM

    rng = np.random.default_rng(20260223)
    signal = rng.standard_normal(2048)

    wdm = WDM(16, 32, 6, 8)
    tf_map_power = wdm.t2w(signal, sample_rate=4096.0, t0=0.0, MM=0)
    wdm.set_td_filter(n_coeffs=8, L=2)

    pixel_index = (tf_map_power.n_time // 2) * (wdm.M + 1) + (wdm.M // 2)
    with pytest.raises(ValueError, match="require 00/90 quadratures"):
        wdm.get_td_amp(tf_map_power, pixel_index, 0, mode="p")


def test_set_td_filter_rejects_odd_m():
    from wdm_wavelet.wdm import WDM

    wdm = WDM(15, 32, 6, 8)
    with pytest.raises(ValueError, match="odd M"):
        wdm.set_td_filter(n_coeffs=8, L=2)
