// Copyright 2022 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

$assert SAMPLE_TILE >= 1
#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#include <xnnpack/math.h>
#include <xnnpack/fft.h>


void xnn_cs16_bfly4_ukernel__scalar_x${SAMPLE_TILE}(
    size_t batch,
    size_t samples,
    int16_t* data,
    const int16_t* twiddle,
    size_t stride)
{
  assert(batch != 0);
  assert(samples != 0);
  assert(samples % (sizeof(int16_t) * 2) == 0);
  assert(data != NULL);
  assert(stride != 0);
  assert(twiddle != NULL);

  int16_t* data3 = data;

  do {
    int16_t* data0 = data3;
    int16_t* data1 = (int16_t*) ((uintptr_t) data0 + samples);
    int16_t* data2 = (int16_t*) ((uintptr_t) data1 + samples);
    data3 = (int16_t*) ((uintptr_t) data2 + samples);

    const int16_t* tw1 = twiddle;
    const int16_t* tw2 = twiddle;
    const int16_t* tw3 = twiddle;
    tw1 = (const int16_t*) ((uintptr_t) tw1 + stride);
    tw2 = (const int16_t*) ((uintptr_t) tw2 + stride * 2);
    tw3 = (const int16_t*) ((uintptr_t) tw3 + stride * 3);

    size_t s = samples - sizeof(int16_t) * 2;

    // First sample skips twiddle.
    // Same code as samples=1 but supports stride
    {
      int32_t vout0r = (int32_t) data0[0];
      int32_t vout0i = (int32_t) data0[1];
      int32_t vout1r = (int32_t) data1[0];
      int32_t vout1i = (int32_t) data1[1];
      int32_t vout2r = (int32_t) data2[0];
      int32_t vout2i = (int32_t) data2[1];
      int32_t vout3r = (int32_t) data3[0];
      int32_t vout3i = (int32_t) data3[1];

      // Note 32767 / 4 = 8191.  Should be 8192.
      vout0r = math_asr_s32(vout0r * 8191 + 16384, 15);
      vout0i = math_asr_s32(vout0i * 8191 + 16384, 15);
      vout1r = math_asr_s32(vout1r * 8191 + 16384, 15);
      vout1i = math_asr_s32(vout1i * 8191 + 16384, 15);
      vout2r = math_asr_s32(vout2r * 8191 + 16384, 15);
      vout2i = math_asr_s32(vout2i * 8191 + 16384, 15);
      vout3r = math_asr_s32(vout3r * 8191 + 16384, 15);
      vout3i = math_asr_s32(vout3i * 8191 + 16384, 15);

      const int32_t vtmp5r = vout0r - vout2r;
      const int32_t vtmp5i = vout0i - vout2i;
      vout0r += vout2r;
      vout0i += vout2i;
      const int32_t vtmp3r = vout1r + vout3r;
      const int32_t vtmp3i = vout1i + vout3i;
      const int32_t vtmp4r = vout1i - vout3i;
      const int32_t vtmp4i = -(vout1r - vout3r);  // swap r,i and neg i
      vout2r = vout0r - vtmp3r;
      vout2i = vout0i - vtmp3i;

      vout0r += vtmp3r;
      vout0i += vtmp3i;

      vout1r = vtmp5r + vtmp4r;
      vout1i = vtmp5i + vtmp4i;
      vout3r = vtmp5r - vtmp4r;
      vout3i = vtmp5i - vtmp4i;

      data0[0] = (int16_t) vout0r;
      data0[1] = (int16_t) vout0i;
      data1[0] = (int16_t) vout1r;
      data1[1] = (int16_t) vout1i;
      data2[0] = (int16_t) vout2r;
      data2[1] = (int16_t) vout2i;
      data3[0] = (int16_t) vout3r;
      data3[1] = (int16_t) vout3i;
      data0 += 2;
      data1 += 2;
      data2 += 2;
      data3 += 2;
    }

    $if SAMPLE_TILE > 1:
      for (; s >= ${SAMPLE_TILE} * sizeof(int16_t) * 2; s -= ${SAMPLE_TILE} * sizeof(int16_t) * 2) {
        $for C in range(SAMPLE_TILE):
          int32_t vout0r${C} = (int32_t) data0[${C * 2 + 0}];
          int32_t vout0i${C} = (int32_t) data0[${C * 2 + 1}];
        $for C in range(SAMPLE_TILE):
          int32_t vout1r${C} = (int32_t) data1[${C * 2 + 0}];
          int32_t vout1i${C} = (int32_t) data1[${C * 2 + 1}];
        $for C in range(SAMPLE_TILE):
          int32_t vout2r${C} = (int32_t) data2[${C * 2 + 0}];
          int32_t vout2i${C} = (int32_t) data2[${C * 2 + 1}];
        $for C in range(SAMPLE_TILE):
          int32_t vout3r${C} = (int32_t) data3[${C * 2 + 0}];
          int32_t vout3i${C} = (int32_t) data3[${C * 2 + 1}];

        $for C in range(SAMPLE_TILE):
          const int32_t vtw1r${C} = (const int32_t) tw1[0];
          const int32_t vtw1i${C} = (const int32_t) tw1[1];
          tw1 = (const int16_t*) ((uintptr_t) tw1 + stride);
        $for C in range(SAMPLE_TILE):
          const int32_t vtw2r${C} = (const int32_t) tw2[0];
          const int32_t vtw2i${C} = (const int32_t) tw2[1];
          tw2 = (const int16_t*) ((uintptr_t) tw2 + stride * 2);
        $for C in range(SAMPLE_TILE):
          const int32_t vtw3r${C} = (const int32_t) tw3[0];
          const int32_t vtw3i${C} = (const int32_t) tw3[1];
          tw3 = (const int16_t*) ((uintptr_t) tw3 + stride * 3);

        // Note 32767 / 4 = 8191.  Should be 8192.
        $for C in range(SAMPLE_TILE):
          vout0r${C} = math_asr_s32(vout0r${C} * 8191 + 16384, 15);
          vout0i${C} = math_asr_s32(vout0i${C} * 8191 + 16384, 15);
        $for C in range(SAMPLE_TILE):
          vout1r${C} = math_asr_s32(vout1r${C} * 8191 + 16384, 15);
          vout1i${C} = math_asr_s32(vout1i${C} * 8191 + 16384, 15);
        $for C in range(SAMPLE_TILE):
          vout2r${C} = math_asr_s32(vout2r${C} * 8191 + 16384, 15);
          vout2i${C} = math_asr_s32(vout2i${C} * 8191 + 16384, 15);
        $for C in range(SAMPLE_TILE):
          vout3r${C} = math_asr_s32(vout3r${C} * 8191 + 16384, 15);
          vout3i${C} = math_asr_s32(vout3i${C} * 8191 + 16384, 15);

        $for C in range(SAMPLE_TILE):
          const int32_t vtmp0r${C} = math_asr_s32(vout1r${C} * vtw1r${C} - vout1i${C} * vtw1i${C} + 16384, 15);
          const int32_t vtmp0i${C} = math_asr_s32(vout1r${C} * vtw1i${C} + vout1i${C} * vtw1r${C} + 16384, 15);
        $for C in range(SAMPLE_TILE):
          const int32_t vtmp1r${C} = math_asr_s32(vout2r${C} * vtw2r${C} - vout2i${C} * vtw2i${C} + 16384, 15);
          const int32_t vtmp1i${C} = math_asr_s32(vout2r${C} * vtw2i${C} + vout2i${C} * vtw2r${C} + 16384, 15);
        $for C in range(SAMPLE_TILE):
          const int32_t vtmp2r${C} = math_asr_s32(vout3r${C} * vtw3r${C} - vout3i${C} * vtw3i${C} + 16384, 15);
          const int32_t vtmp2i${C} = math_asr_s32(vout3r${C} * vtw3i${C} + vout3i${C} * vtw3r${C} + 16384, 15);

        $for C in range(SAMPLE_TILE):
          const int32_t vtmp5r${C} = vout0r${C} - vtmp1r${C};
          const int32_t vtmp5i${C} = vout0i${C} - vtmp1i${C};
        $for C in range(SAMPLE_TILE):
          vout0r${C} += vtmp1r${C};
          vout0i${C} += vtmp1i${C};
        $for C in range(SAMPLE_TILE):
          const int32_t vtmp3r${C} = vtmp0r${C} + vtmp2r${C};
          const int32_t vtmp3i${C} = vtmp0i${C} + vtmp2i${C};
        $for C in range(SAMPLE_TILE):
          const int32_t vtmp4r${C} = vtmp0i${C} - vtmp2i${C};
          const int32_t vtmp4i${C} = -(vtmp0r${C} - vtmp2r${C});  // swap r,i and neg i
        $for C in range(SAMPLE_TILE):
          vout2r${C} = vout0r${C} - vtmp3r${C};
          vout2i${C} = vout0i${C} - vtmp3i${C};
        $for C in range(SAMPLE_TILE):
          vout0r${C} += vtmp3r${C};
          vout0i${C} += vtmp3i${C};
        $for C in range(SAMPLE_TILE):
          vout1r${C} = vtmp5r${C} + vtmp4r${C};
          vout1i${C} = vtmp5i${C} + vtmp4i${C};
        $for C in range(SAMPLE_TILE):
          vout3r${C} = vtmp5r${C} - vtmp4r${C};
          vout3i${C} = vtmp5i${C} - vtmp4i${C};

        $for C in range(SAMPLE_TILE):
          data0[${C * 2 + 0}] = (int16_t) vout0r${C};
          data0[${C * 2 + 1}] = (int16_t) vout0i${C};
        data0 += ${SAMPLE_TILE} * 2;
        $for C in range(SAMPLE_TILE):
          data1[${C * 2 + 0}] = (int16_t) vout1r${C};
          data1[${C * 2 + 1}] = (int16_t) vout1i${C};
        data1 += ${SAMPLE_TILE} * 2;
        $for C in range(SAMPLE_TILE):
          data2[${C * 2 + 0}] = (int16_t) vout2r${C};
          data2[${C * 2 + 1}] = (int16_t) vout2i${C};
        data2 += ${SAMPLE_TILE} * 2;
        $for C in range(SAMPLE_TILE):
          data3[${C * 2 + 0}] = (int16_t) vout3r${C};
          data3[${C * 2 + 1}] = (int16_t) vout3i${C};
        data3 += ${SAMPLE_TILE} * 2;
      }
    if XNN_UNLIKELY(s != 0) {
      do {
        int32_t vout0r = (int32_t) data0[0];
        int32_t vout0i = (int32_t) data0[1];
        int32_t vout1r = (int32_t) data1[0];
        int32_t vout1i = (int32_t) data1[1];
        int32_t vout2r = (int32_t) data2[0];
        int32_t vout2i = (int32_t) data2[1];
        int32_t vout3r = (int32_t) data3[0];
        int32_t vout3i = (int32_t) data3[1];

        const int32_t vtw1r = (const int32_t) tw1[0];
        const int32_t vtw1i = (const int32_t) tw1[1];
        const int32_t vtw2r = (const int32_t) tw2[0];
        const int32_t vtw2i = (const int32_t) tw2[1];
        const int32_t vtw3r = (const int32_t) tw3[0];
        const int32_t vtw3i = (const int32_t) tw3[1];
        tw1 = (const int16_t*) ((uintptr_t) tw1 + stride);
        tw2 = (const int16_t*) ((uintptr_t) tw2 + stride * 2);
        tw3 = (const int16_t*) ((uintptr_t) tw3 + stride * 3);

        // Note 32767 / 4 = 8191.  Should be 8192.
        vout0r = math_asr_s32(vout0r * 8191 + 16384, 15);
        vout0i = math_asr_s32(vout0i * 8191 + 16384, 15);
        vout1r = math_asr_s32(vout1r * 8191 + 16384, 15);
        vout1i = math_asr_s32(vout1i * 8191 + 16384, 15);
        vout2r = math_asr_s32(vout2r * 8191 + 16384, 15);
        vout2i = math_asr_s32(vout2i * 8191 + 16384, 15);
        vout3r = math_asr_s32(vout3r * 8191 + 16384, 15);
        vout3i = math_asr_s32(vout3i * 8191 + 16384, 15);

        const int32_t vtmp0r = math_asr_s32(vout1r * vtw1r - vout1i * vtw1i + 16384, 15);
        const int32_t vtmp0i = math_asr_s32(vout1r * vtw1i + vout1i * vtw1r + 16384, 15);
        const int32_t vtmp1r = math_asr_s32(vout2r * vtw2r - vout2i * vtw2i + 16384, 15);
        const int32_t vtmp1i = math_asr_s32(vout2r * vtw2i + vout2i * vtw2r + 16384, 15);
        const int32_t vtmp2r = math_asr_s32(vout3r * vtw3r - vout3i * vtw3i + 16384, 15);
        const int32_t vtmp2i = math_asr_s32(vout3r * vtw3i + vout3i * vtw3r + 16384, 15);

        const int32_t vtmp5r = vout0r - vtmp1r;
        const int32_t vtmp5i = vout0i - vtmp1i;
        vout0r += vtmp1r;
        vout0i += vtmp1i;
        const int32_t vtmp3r = vtmp0r + vtmp2r;
        const int32_t vtmp3i = vtmp0i + vtmp2i;
        const int32_t vtmp4r = vtmp0i - vtmp2i;
        const int32_t vtmp4i = -(vtmp0r - vtmp2r);  // swap r,i and neg i
        vout2r = vout0r - vtmp3r;
        vout2i = vout0i - vtmp3i;

        vout0r += vtmp3r;
        vout0i += vtmp3i;

        vout1r = vtmp5r + vtmp4r;
        vout1i = vtmp5i + vtmp4i;
        vout3r = vtmp5r - vtmp4r;
        vout3i = vtmp5i - vtmp4i;

        data0[0] = (int16_t) vout0r;
        data0[1] = (int16_t) vout0i;
        data1[0] = (int16_t) vout1r;
        data1[1] = (int16_t) vout1i;
        data2[0] = (int16_t) vout2r;
        data2[1] = (int16_t) vout2i;
        data3[0] = (int16_t) vout3r;
        data3[1] = (int16_t) vout3i;
        data0 += 2;
        data1 += 2;
        data2 += 2;
        data3 += 2;

        s -= sizeof(int16_t) * 2;
      } while (s != 0);
    }
  } while (--batch != 0);
}