# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/test_grad.ipynb.

# %% auto 0
__all__ = ['get_random_biquads', 'create_test_inputs', 'test_low_order']

# %% ../nbs/test_grad.ipynb 2
import numpy as np
import pytest

# %% ../nbs/test_grad.ipynb 3
def get_random_biquads(cmplx=False, rng=None):
    if rng is None:
        rng = np.random.default_rng()

    if cmplx:
        mag = rng.random(2)
        phase = rng.random(2) * 2 * np.pi
        roots = mag * np.exp(1j * phase)
        return np.array([-roots[0] - roots[1], roots[0] * roots[1]], dtype=np.complex128)
    mag = rng.random(1)
    phase = rng.random(1) * np.pi
    return np.array([-mag * np.cos(phase) * 2, mag**2], dtype=np.float64)

def create_test_inputs(batch_size, samples, cmplx=False, seed=None):
    rng = np.random.default_rng(seed)

    start_coeffs = get_random_biquads(cmplx, rng)
    end_coeffs = get_random_biquads(cmplx, rng)
    dtype = np.complex128 if cmplx else np.float64

    A = np.stack(
        [
            np.linspace(start_coeffs[i], end_coeffs[i], samples, dtype=dtype)
            for i in range(2)
        ]
    ).T.reshape(1, samples, 2)
    
    A = np.tile(A, (batch_size, 1, 1))
    x = rng.normal(size=(batch_size, samples)).astype(dtype)
    zi = rng.normal(size=(batch_size, 2)).astype(dtype)
    return x, A, zi


# %% ../nbs/test_grad.ipynb 4
@pytest.mark.parametrize(
    "samples",
    [32],
)
@pytest.mark.parametrize(
    "cmplx",
    [False],
)
def test_low_order(
    samples: int,
    cmplx: bool,
):
    from jaxpole.filter import allpole
    import jax
    from jax.test_util import check_grads

    print(f"The jax default backend is {jax.default_backend()}")
    
    batch_size = 4
    x, A, zi = create_test_inputs(batch_size, samples, cmplx)
    check_grads(allpole, (x, A, zi), order=1, modes=("rev"))
