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

# %% auto 0
__all__ = ['batched_all_pole_direct_form_I', 'all_pole_direct_form_I', 'allpole_bwd', 'allpole_fwd', 'allpole']

# %% ../nbs/filter.ipynb 2
import jax
import jax.numpy as jnp
from jax import custom_vjp, custom_jvp


# %% ../nbs/filter.ipynb 3
def all_pole_direct_form_I(
    x: jnp.ndarray,  # (time, n_coeffs,)
    a: jnp.ndarray,  # (time,)
    zi: jnp.ndarray = None,  # (n_coeffs,) initial state
):
    """
    Apply an all-pole filter defined by the coefficients `a` to the input signal `x`.
    """

    def f(carry, input):
        x, a = input
        y = x - jnp.sum(jnp.flip(carry) * a)
        # we always move in a stride of 1 so we stack all the
        # values after the first one and the current value
        curr = jnp.stack([carry[1:].squeeze(0), y], axis=0)
        return curr, y

    _, y = jax.lax.scan(f=f, init=jnp.flip(zi, axis=-1), xs=(x, a))

    return y

batched_all_pole_direct_form_I = jax.vmap(all_pole_direct_form_I)

def allpole_bwd(res, grad_y):
    A, zi, y = res
    B, T, order = A.shape
    
    flipped_A = jnp.flip(A, axis=2)

    transposed_A = jnp.transpose(flipped_A, (0, 2, 1))
    padded_transposed_A = jnp.pad(transposed_A, ((0, 0), (0, 0), (0, order + 1)))

    reshaped_A = jnp.reshape(padded_transposed_A, (B, T + order + 1, order))
    sliced_A = reshaped_A[:, :-1, :]

    shifted_A = jnp.reshape(sliced_A, (B, order, T + order))
    shifted_A = jnp.transpose(shifted_A, (0, 2, 1))
    shifted_A = jnp.flip(shifted_A, axis=2)
    
    if zi is None:
        shifted_A = shifted_A[:, order:, :]
        padded_grad_y = grad_y
    else:
        padded_grad_y = jnp.pad(grad_y, ((0, 0), (order, 0)), mode="constant")

    flipped_padded_grad_y = jnp.flip(padded_grad_y, axis=1)
    flipped_shifted_A = jnp.flip(shifted_A, axis=1).conj()

    flipped_grad_x = batched_all_pole_direct_form_I(
        flipped_padded_grad_y,
        flipped_shifted_A,
        jnp.zeros_like(zi),
    )

    grad_zi = flipped_grad_x[:, -order:] if zi is not None else None
    flipped_grad_x = flipped_grad_x[:, :-order] if zi is not None else flipped_grad_x

    grad_x = jnp.flip(flipped_grad_x, axis=1) if zi is not None else flipped_grad_x

    valid_y = y[:, :-1]
    padded_y = jnp.concatenate([jnp.flip(zi, axis=1), valid_y], axis=1)

    # unfold the y signal
    unfolded_y = jnp.flip(jax.lax.conv_general_dilated_patches(
        padded_y[..., None],
        filter_shape=(order,),
        window_strides=(1,),
        padding="VALID",
        dimension_numbers=("NTC", "OIT", "NTC"),
    ), axis=2)
    
    grad_A = unfolded_y.conj() * -jnp.flip(flipped_grad_x, axis=1)[..., None]
    return grad_x, grad_A, grad_zi


def allpole_fwd(x, A, zi):
    y = batched_all_pole_direct_form_I(x, A, zi)
    return y, (A, zi, y)

@custom_vjp
def allpole(
    x: jnp.ndarray,  # (B, time)
    A: jnp.ndarray,  # (B, time, n_coeffs)
    zi: jnp.ndarray = None,  # (B, n_coeffs) initial state
):
    """
    Apply an all-pole filter defined by the coefficients `A` to the input signal `x`.
    """
    return allpole_fwd(x, A, zi)[0]


# The jvp will work for fwd mode but not for rev mode due to the scan in the allpole function
# def allpole_jvp(primals, tangents):
#     x, A, zi = primals
#     dx, dA, dzi = tangents
#     y = allpole(x, A, zi)

#     B, T, order = A.shape

#     # Calculate the forward mode Jacobian-vector product
#     fwd_zi = dzi if dzi is not None else jnp.zeros_like(zi)
#     fwd_x = dx if dx is not None else jnp.zeros_like(y)

#     if dA is not None:
#         padded_y = jnp.concatenate([jnp.flip(zi, axis=1), y[:, :-1]], axis=1)

#         # unfold the y signal
#         unfolded_y = jnp.flip(jax.lax.conv_general_dilated_patches(
#             padded_y[..., None],
#             filter_shape=(order,),
#             window_strides=(1,),
#             padding="VALID",
#             dimension_numbers=("NTC", "OIT", "NTC"),
#         ), axis=2)
        
#         fwd_A = -jnp.sum(unfolded_y * dA, axis=2)
#         fwd_x = dx + fwd_A

        
#     return y, allpole(fwd_x, A, fwd_zi)


# allpole.defjvp(allpole_jvp)
allpole.defvjp(allpole_fwd, allpole_bwd)
