#ifndef __FOURIER_WATER_COMPUTE_INC
#define __FOURIER_WATER_COMPUTE_INC

#include "UnityCG.cginc"
#include "UnityCustomRenderTexture.cginc"
#include "fourier_water_twiddles.cginc"

#define DIR_HORIZONTAL 0
#define DIR_VERTICAL   1

uniform uint _Stage;
uniform uint _Big_N;
uniform float _Size_M;
uniform uint _Stride;
uniform uint _Direction;
uniform Texture2D _Signal;
uniform float4 _Signal_TexelSize;

#if defined(_RADIX_8)
#define FFT_RADIX 8
#elif defined(_RADIX_4)
#define FFT_RADIX 4
#else
#define FFT_RADIX 2
#endif

float2 cmul(float2 x, float2 y) {
  return float2(x[0] * y[0] - x[1] * y[1], x[0] * y[1] + x[1] * y[0]);
}

float4 mul_ni(float4 x) {
  return x.yxwz * float4(1, -1, 1, -1);
}

float4 rot_45(float4 x) {
  return (x - mul_ni(x)) * 0.7071067811865475;
}

// Factored butterflies from https://yummers.dev/a-fast-cpu-fft.html.
float4 fft_radix_2(float4 a, float4 b, uint lane) {
  return lane == 0 ? a + b : a - b;
}

float4 fft_radix_4(float4 a, float4 b, float4 c, float4 d, uint lane) {
  float4 even = (lane & 1) == 0 ? a + c : a - c;
  float4 odd = (lane & 1) == 0 ? b + d : mul_ni(b - d);
  return lane < 2 ? even + odd : even - odd;
}

float4 fft_radix_8(float4 a, float4 b, float4 c, float4 d,
    float4 e, float4 f, float4 g, float4 h, uint lane) {
  float4 even, odd;
  switch (lane & 3) {
    case 0:
      even = (a + e) + (c + g);
      odd = (b + f) + (d + h);
      break;
    case 1:
      even = (a - e) + mul_ni(c - g);
      odd = rot_45(mul_ni(b - f) - (d - h));
      break;
    case 2:
      even = (a + e) - (c + g);
      odd = mul_ni((b + f) - (d + h));
      break;
    default:
      even = (a - e) - mul_ni(c - g);
      odd = rot_45(-(b - f) + mul_ni(d - h));
      break;
  }
  return lane < 4 ? even + odd : even - odd;
}

float4 fft_stage(uint x, uint y) {
  uint main_coord = _Direction == DIR_HORIZONTAL ? x : y;
  uint alt_coord = _Direction == DIR_HORIZONTAL ? y : x;
  uint start_idx = main_coord % _Stride;
  uint group_size = _Big_N / FFT_RADIX;
  uint k = (main_coord / _Stride) % group_size;
  uint lane = (main_coord / _Stride) / group_size;
  float4 v[FFT_RADIX];
  [unroll]
  for (uint j = 0; j < FFT_RADIX; j++) {
    uint2 coords = uint2(start_idx + (FFT_RADIX * k + j) * _Stride, alt_coord);
    if (_Direction == DIR_VERTICAL) {
      coords = coords.yx;
    }
    v[j] = _Signal.Load(uint3(coords, 0));
    if (j != 0 && k != 0) {
      float2 factor = kTwiddles[j * k * _Stride];
      v[j] = float4(cmul(v[j].rg, factor), cmul(v[j].ba, factor));
    }
  }
#if defined(_RADIX_8)
  return fft_radix_8(v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], lane);
#elif defined(_RADIX_4)
  return fft_radix_4(v[0], v[1], v[2], v[3], lane);
#else
  return fft_radix_2(v[0], v[1], lane);
#endif
}

float4 frag(v2f_customrendertexture i) : SV_Target {
  float2 px_sz = float2(_CustomRenderTextureWidth, _CustomRenderTextureHeight);
  uint2 px_pos = i.globalTexcoord.xy * px_sz;
#if defined(_TIME_EVOLVE) || defined(_DERIVATIVE) || defined(_DISPLACEMENT) || defined(_DISPLACEMENT_DERIVATIVE)
  float2 input = _Signal.Load(uint3(px_pos, 0));
  float2 k2 = px_pos;
  k2.x = k2.x < px_sz.x / 2 ? k2.x : k2.x - px_sz.x;
  k2.y = k2.y < px_sz.y / 2 ? k2.y : k2.y - px_sz.y;
  k2 *= 2.0f * 3.14159265f / _Size_M;
  float k = length(k2);
  float omega = sqrt(9.81f * k + 0.076f * k * k * k / 1000.0f);
  float omega_s, omega_c;
  sincos(omega * _Time[1], omega_s, omega_c);
  float2 H = cmul(input, float2(omega_c, omega_s));
  // Negative-exponent FFT: -i k/|k| moves points toward height crests.
  float inv_k = k > 0.0 ? rcp(k) : 0.0;
#if defined(_DERIVATIVE)
  float2 slopeX = cmul(H, float2(0, -k2.x));
  float2 slopeY = cmul(H, float2(0, -k2.y));
  return float4(slopeX, slopeY);
#elif defined(_DISPLACEMENT)
  return float4(cmul(H, float2(0, -k2.x * inv_k)),
      cmul(H, float2(0, -k2.y * inv_k)));
#elif defined(_DISPLACEMENT_DERIVATIVE)
  // RG: dDx/dy = dDy/dx; BA: dDy/dy.
  return float4(-H * (k2.x * k2.y * inv_k),
      -H * (k2.y * k2.y * inv_k));
#else
  // RG: height; BA: dDx/dx.
  return float4(H, -H * (k2.x * k2.x * inv_k));
#endif

#else
  return fft_stage(px_pos.x, px_pos.y);
#endif
}

#endif  // __FOURIER_WATER_COMPUTE_INC