import numpy as np
import pytest


def _build_truncated_filter(M=16, K=32, beta_order=6, precision=8, n_max=30000):
    from wdm_wavelet.core.get_filter import get_filter
    from wdm_wavelet.fourier_coeff import load_fourier_coeff

    cos, _, _, cos_size, _, _ = load_fourier_coeff()
    raw_filter = get_filter(M, K, beta_order, n_max, cos, cos_size)

    residual = 1.0 - raw_filter[0] ** 2
    threshold = 10 ** (-precision)
    n = 1
    stride = 2 * M

    while residual > threshold or (n - 1) % stride or n // stride < 3:
        if n >= len(raw_filter):
            raise RuntimeError("n_max is too small for the requested precision.")
        residual -= 2 * raw_filter[n] ** 2
        n += 1

    return raw_filter[:n], residual


def test_get_filter_truncation_constraints():
    M = 16
    precision = 8
    filt, residual = _build_truncated_filter(M=M, K=32, beta_order=6, precision=precision)

    assert np.isfinite(filt).all()
    assert (len(filt) - 1) % (2 * M) == 0
    assert len(filt) // (2 * M) >= 3
    assert residual <= 10 ** (-precision)


def test_t2w_mode_shapes_and_power_relation():
    from wdm_wavelet.core.t2w import t2w_jax

    rng = np.random.default_rng(1234)
    signal = rng.standard_normal(1533)
    M = 16
    filt, _ = _build_truncated_filter(M=M, K=32, beta_order=6, precision=8)

    m_l_full, aligned_len, tf_full = t2w_jax(M, len(filt), signal, filt, -1)
    m_l_power, aligned_len_power, tf_power = t2w_jax(M, len(filt), signal, filt, 0)

    assert m_l_full == len(filt)
    assert m_l_power == 0
    assert aligned_len == aligned_len_power
    assert aligned_len % M == 0
    assert aligned_len >= len(signal)
    assert tf_full.shape == (2, aligned_len // M, M + 1)
    assert tf_power.shape == tf_full.shape

    # In power mode, only the first plane is populated.
    assert np.all(tf_power[1] == 0.0)

    # Eq. mapping consistency:
    # power = re^2 + im^2 = 0.5 * (w_nm^2 + w_hat_nm^2)
    expected_power = 0.5 * (tf_full[0] ** 2 + tf_full[1] ** 2)
    np.testing.assert_allclose(tf_power[0], expected_power, rtol=1e-12, atol=1e-12)


def test_t2w_is_linear():
    from wdm_wavelet.core.t2w import t2w_jax

    rng = np.random.default_rng(99)
    x = rng.standard_normal(1024)
    y = rng.standard_normal(1024)
    M = 8
    filt, _ = _build_truncated_filter(M=M, K=16, beta_order=6, precision=8)

    _, _, tx = t2w_jax(M, len(filt), x, filt, -1)
    _, _, ty = t2w_jax(M, len(filt), y, filt, -1)
    _, _, txy = t2w_jax(M, len(filt), x + y, filt, -1)

    np.testing.assert_allclose(txy, tx + ty, rtol=1e-12, atol=1e-12)


def test_w2t_is_linear():
    from wdm_wavelet.core.w2t import w2t_jax

    rng = np.random.default_rng(7)
    M = 8
    m_layer = M + 1
    n_filter = 33
    n_time = 20
    n_coeffs = n_time * (2 * M + 2)

    coeff_a = rng.standard_normal(n_coeffs)
    coeff_b = rng.standard_normal(n_coeffs)
    filt = np.abs(rng.standard_normal(n_filter))

    ta = w2t_jax(coeff_a, m_layer, filt)
    tb = w2t_jax(coeff_b, m_layer, filt)
    tab = w2t_jax(coeff_a + coeff_b, m_layer, filt)

    np.testing.assert_allclose(tab, ta + tb, rtol=1e-12, atol=1e-12)


def test_t2w_jax_matches_reference():
    from wdm_wavelet.core.t2w import t2w_jax, HAS_JAX as HAS_JAX_T2W

    if not HAS_JAX_T2W:
        pytest.skip("JAX backend is not available.")

    rng = np.random.default_rng(123)
    signal = rng.standard_normal(1536)
    M = 16
    filt, _ = _build_truncated_filter(M=M, K=32, beta_order=6, precision=8)

    # Round-trip: forward transform twice with same inputs should be identical.
    got1 = t2w_jax(M, len(filt), signal, filt, -1)
    got2 = t2w_jax(M, len(filt), signal, filt, -1)

    assert got1[0] == got2[0]
    assert got1[1] == got2[1]
    np.testing.assert_allclose(got2[2], got1[2], rtol=1e-12, atol=1e-12)


def test_w2t_jax_matches_reference():
    from wdm_wavelet.core.w2t import w2t_jax, HAS_JAX as HAS_JAX_W2T

    if not HAS_JAX_W2T:
        pytest.skip("JAX backend is not available.")

    rng = np.random.default_rng(321)
    M = 8
    m_layer = M + 1
    n_filter = 33
    n_time = 20
    n_coeffs = n_time * (2 * M + 2)

    coeffs = rng.standard_normal(n_coeffs)
    filt = np.abs(rng.standard_normal(n_filter))

    # Deterministic: same inputs produce identical outputs.
    got1 = w2t_jax(coeffs, m_layer, filt)
    got2 = w2t_jax(coeffs, m_layer, filt)
    np.testing.assert_allclose(got2, got1, rtol=1e-12, atol=1e-12)


def test_w2t_output_length_truncation():
    from wdm_wavelet.core.w2t import w2t_jax, HAS_JAX as HAS_JAX_W2T

    rng = np.random.default_rng(20260223)
    M = 8
    m_layer = M + 1
    n_filter = 33
    n_time = 20
    n_coeffs = n_time * (2 * M + 2)

    coeffs = rng.standard_normal(n_coeffs)
    filt = np.abs(rng.standard_normal(n_filter))
    target_length = 123

    out = w2t_jax(coeffs, m_layer, filt, output_length=target_length)
    assert out.shape == (target_length,)

    if HAS_JAX_W2T:
        out_jax = w2t_jax(coeffs, m_layer, filt, output_length=target_length)
        np.testing.assert_allclose(out_jax, out, rtol=1e-12, atol=1e-12)


def test_t2w_numba_matches_jax():
    from wdm_wavelet.core.t2w import (
        HAS_JAX as HAS_JAX_T2W,
        HAS_NUMBA as HAS_NUMBA_T2W,
        t2w_jax,
        t2w_numba,
    )

    if not (HAS_JAX_T2W and HAS_NUMBA_T2W):
        pytest.skip("Requires both JAX and Numba backends.")

    rng = np.random.default_rng(20260301)
    signal = rng.standard_normal(1536)
    M = 16
    filt, _ = _build_truncated_filter(M=M, K=32, beta_order=6, precision=8)

    got_jax = t2w_jax(M, len(filt), signal, filt, -1)
    got_numba = t2w_numba(M, len(filt), signal, filt, -1)

    assert got_jax[0] == got_numba[0]
    assert got_jax[1] == got_numba[1]
    np.testing.assert_allclose(got_numba[2], got_jax[2], rtol=1e-12, atol=1e-12)


def test_w2t_numba_matches_jax():
    from wdm_wavelet.core.w2t import (
        HAS_JAX as HAS_JAX_W2T,
        HAS_NUMBA as HAS_NUMBA_W2T,
        w2tQ_jax,
        w2tQ_numba,
        w2t_jax,
        w2t_numba,
    )

    if not (HAS_JAX_W2T and HAS_NUMBA_W2T):
        pytest.skip("Requires both JAX and Numba backends.")

    rng = np.random.default_rng(20260302)
    signal = rng.standard_normal(2048)
    M = 16
    filt, _ = _build_truncated_filter(M=M, K=32, beta_order=6, precision=8)

    from wdm_wavelet.core.t2w import t2w_jax

    _, _, tf_map = t2w_jax(M, len(filt), signal, filt, -1)
    tf_flat = np.asarray(tf_map).reshape(-1)
    m_layer = M + 1

    out_jax = w2t_jax(tf_flat, m_layer, filt)
    out_numba = w2t_numba(tf_flat, m_layer, filt)
    np.testing.assert_allclose(out_numba, out_jax, rtol=1e-12, atol=1e-12)

    outq_jax = w2tQ_jax(tf_flat, m_layer, filt)
    outq_numba = w2tQ_numba(tf_flat, m_layer, filt)
    np.testing.assert_allclose(outq_numba, outq_jax, rtol=1e-12, atol=1e-12)
