"""Implicitly padded mixed-radix DFT and linear convolution."""
import numpy as np


def implicit_dft(x, M, m):
    """Positive-sign DFT of x zero-padded to M, without making padded x.

    Returns the convention F[k]=sum_j exp(+2pi i*j*k/M)x[j].
    M must be divisible by m; x may have arbitrary length <= M.
    """
    x = np.asarray(x)
    L = x.shape[-1]
    if L > M or M % m:
        raise ValueError("require len(x)<=M and M divisible by m")
    p, q = (L + m - 1) // m, M // m
    # Tiles are bounded and masking is done at load time. This is the only
    # input-sized temporary; no length-M padded input is allocated.
    tiles = np.zeros(x.shape[:-1] + (p, m), dtype=np.result_type(x, np.complex128))
    for t in range(p):
        lo, hi = t*m, min((t+1)*m, L)
        tiles[..., t, :hi-lo] = x[..., lo:hi]
    r = np.arange(q)[:, None]
    t = np.arange(p)[None, :]
    s = np.arange(m)[None, :]
    zq = np.exp(2j*np.pi*r*t/q)
    zqm = np.exp(2j*np.pi*r*s/(q*m))
    # U[...,r,s] = sum_t zeta_q^(rt) x[t,s]
    U = np.einsum('rt,...ts->...rs', zq, tiles) * zqm
    # positive-sign m DFT in s; output is indexed q*ell+r
    zm = np.exp(2j*np.pi*np.arange(m)[:, None]*np.arange(m)[None, :]/m)
    F_rsell = np.einsum('ls,...rs->...rl', zm, U)
    return np.transpose(F_rsell, tuple(range(F_rsell.ndim-2)) + (F_rsell.ndim-1, F_rsell.ndim-2)).reshape(x.shape[:-1] + (M,))


def implicit_idft(F, L, m):
    """Inverse of implicit_dft, returning the first L entries."""
    F = np.asarray(F)
    M = F.shape[-1]
    if M % m:
        raise ValueError("M must be divisible by m")
    q, p = M // m, (L + m - 1) // m
    # F[...,q*ell+r] -> F[...,r,ell]
    A = F.reshape(F.shape[:-1] + (m, q)).swapaxes(-2, -1)
    r = np.arange(q)[:, None]
    t = np.arange(p)[None, :]
    s = np.arange(m)[None, :]
    zq = np.exp(-2j*np.pi*r*t/q)
    zqm = np.exp(-2j*np.pi*r*s/(q*m))
    zm = np.exp(-2j*np.pi*np.arange(m)[:, None]*np.arange(m)[None, :]/m)
    # sum_l zeta_m^(-s*l) F[r,l], indexed r,s
    H = np.einsum('sl,...rl->...rs', zm, A) * zqm
    out = np.einsum('rt,...rs->...ts', zq, H) / M
    return out.reshape(out.shape[:-2] + (p*m,))[..., :L]


def implicit_convolve(x, g, m):
    M = len(x) + len(g) - 1
    # Select a convenient divisible transform length, while preserving exactness.
    M2 = 1
    while M2 < M or M2 % m:
        M2 += 1
    y = implicit_idft(implicit_dft(x, M2, m) * implicit_dft(g, M2, m), M2, m)
    return y[:M].real