// Polyphase rational resampling.
//
// Converting by L/M is, in principle, zero-stuff by L, low-pass, keep every
// Mth sample. Done literally that computes L-1 zeros for every real sample and
// then throws away M-1 of every M results. The polyphase identity removes both:
// output n depends only on filter phase (n*M) mod L and a short run of input
// around floor(n*M/L).
//
// torchaudio expresses this as a conv1d with `new_freq` output channels and
// stride `orig_freq`. That is the same decomposition, but the per-phase filters
// are stored padded to a common width, so every output sample multiplies through
// a kernel row whose support is a small fraction of its length. At 44.1k -> 16k
// with the default filter width the row is 475 taps and the support is about 34,
// and the padding is multiplied and summed like everything else.
//
// Here each phase keeps only its own support, as an offset and a run of taps.

#include <torch/all.h>

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <vector>

namespace {

#if defined(__GNUC__) || defined(__clang__)
#if defined(__x86_64__)
#define RS_VL 8
#else
#define RS_VL 4
#endif
typedef float vf __attribute__((vector_size(RS_VL * sizeof(float))));
#else
#define RS_VL 1
typedef float vf;
#endif

#if defined(__x86_64__) && defined(__GNUC__) && !defined(__clang__)
#define RS_MV __attribute__((target_clones("avx2", "default"), flatten))
#else
#define RS_MV
#endif

inline vf vzero() {
#if RS_VL == 1
  return 0.0f;
#else
  vf v;
  float* p = reinterpret_cast<float*>(&v);
  for (int i = 0; i < RS_VL; i++) p[i] = 0.0f;
  return v;
#endif
}

inline vf vloadu(const float* p) {
  vf v;
  std::memcpy(&v, p, sizeof(vf));
  return v;
}

// Tree reduction rather than a running scalar sum. Phases here carry as few as
// a dozen taps, so the dot product is two or three vector operations and the
// reduction that follows it is a large share of the per-output cost; a chain of
// RS_VL dependent adds is the wrong shape for that.
inline float hsum(const vf& v) {
#if RS_VL == 1
  return v;
#else
  float t[RS_VL];
  std::memcpy(t, &v, sizeof(vf));
  for (int s = RS_VL / 2; s > 0; s >>= 1)
    for (int i = 0; i < s; i++) t[i] += t[i + s];
  return t[0];
#endif
}

// One output block: the L outputs sharing input block `a`. Each reads a
// contiguous run of input, so the tap loop is unit stride in both operands.
RS_MV void resample_block(float* out, const float* x, int64_t T, int64_t a,
                          int64_t L, int64_t M, int64_t n_out,
                          const int32_t* start, const int32_t* len,
                          const int64_t* off, const float* taps) {
  const int64_t abase = a * M;
  for (int64_t r = 0; r < L; r++) {
    const int64_t n = a * L + r;
    if (n >= n_out) return;
    const int64_t base = abase + start[r];
    const int64_t K = len[r];
    const float* h = taps + off[r];

    // Clip the tap run to where the input index is in range rather than
    // testing every tap; the interior blocks take the fast path untouched.
    int64_t m0 = 0, m1 = K;
    if (base < 0) m0 = -base;
    if (base + K > T) m1 = T - base;
    if (m0 >= m1) { out[n] = 0.0f; continue; }

    const float* xp = x + base;
    vf acc = vzero();
    int64_t m = m0;
    for (; m + RS_VL <= m1; m += RS_VL)
      acc = acc + vloadu(h + m) * vloadu(xp + m);
    float s = hsum(acc);
    for (; m < m1; m++) s += h[m] * xp[m];
    out[n] = s;
  }
}

}  // namespace

// out:   [B, n_out] f32
// x:     [B, T] f32 contiguous
// taps:  concatenated per-phase filter supports, f32
// start: [L] int32, input offset of each phase's first tap relative to a*M
// len:   [L] int32, tap count per phase
// off:   [L] int64, index of each phase's first tap within `taps`
void rp_resample(torch::Tensor& out, torch::Tensor const& x,
                 torch::Tensor const& taps, torch::Tensor const& start,
                 torch::Tensor const& len, torch::Tensor const& off,
                 int64_t L, int64_t M) {
  TORCH_CHECK(x.dim() == 2 && x.is_contiguous() &&
                  x.scalar_type() == torch::kFloat32,
              "x must be contiguous f32 [B, T]");
  TORCH_CHECK(out.dim() == 2 && out.is_contiguous(), "out must be [B, n_out]");
  TORCH_CHECK(L >= 1 && M >= 1, "L and M must be positive");
  TORCH_CHECK(start.numel() == L && len.numel() == L && off.numel() == L,
              "phase tables must have L entries");

  const int64_t B = x.size(0), T = x.size(1), n_out = out.size(1);
  const float* xp = x.data_ptr<float>();
  float* op = out.data_ptr<float>();
  const float* tp = taps.data_ptr<float>();
  const int32_t* sp = start.data_ptr<int32_t>();
  const int32_t* lp = len.data_ptr<int32_t>();
  const int64_t* fp = off.data_ptr<int64_t>();

  const int64_t n_blocks = (n_out + L - 1) / L;

  at::parallel_for(0, B * n_blocks, 1, [&](int64_t begin, int64_t end) {
    for (int64_t i = begin; i < end; i++) {
      const int64_t b = i / n_blocks;
      const int64_t a = i - b * n_blocks;
      resample_block(op + b * n_out, xp + b * T, T, a, L, M, n_out,
                     sp, lp, fp, tp);
    }
  });
}