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

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

# %% ../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, p=2):
    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(p)
        ]
    ).T.reshape(1, samples, p)
    
    A = np.tile(A, (batch_size, 1, 1))
    x = rng.normal(size=(batch_size, samples)).astype(dtype)
    zi = rng.normal(size=(batch_size, p)).astype(dtype)
    return x, A, zi

def allpole_np(
    x: np.ndarray, # (B, time,)
    A: np.ndarray, # (B, time, order)
    zi: np.ndarray, # (B, order,)
) -> np.ndarray:
    """
    Reference implementation of the allpole filter.
    """
    B, T = x.shape
    order = zi.shape[-1]
    padded_y = np.empty((B, T + order), dtype=x.dtype)
    padded_y[:, :order] = zi[:, ::-1]
    padded_y[:, order:] = x

    for b in range(B):
        for t in range(T):
            ref = padded_y[b, t + order] - A[b, t][::-1] @ padded_y[b, t:t + order]
            padded_y[b, t + order] = ref

    return padded_y[:, order:]

# %% ../nbs/test_grad.ipynb 4
@pytest.mark.parametrize(
    "samples",
    [32],
)
@pytest.mark.parametrize(
    "cmplx",
    [False],
)
@pytest.mark.parametrize(
    "p",
    [2, 4, 8, 16],
)
def test_all(
    samples: int,
    cmplx: bool,
    p: int,
):
    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, p)
    check_grads(allpole, (x, A, zi), order=2, modes=("rev"))

    y = allpole(x, A, zi)
    y_np = allpole_np(x, A, zi)

    assert np.allclose(y, y_np, atol=1e-6)
