#ifndef __CUSTOM31_INC
#define __CUSTOM31_INC

#define PI      3.14159265f
#define PI_RCP  0.31830988f
#define TAU     (2.0f * PI)
#define HALF_PI (0.5f * PI)

#define glsl_mod(x,y) (((x)-(y)*floor((x)/(y))))

// Differentiable versions of common operators.
#define dabs(x) sqrt((x) * (x) + 1e-6)
#define dlerp(x, y, t) ((x) * (1-t) + (y) * t)
// This was derived using fourier analysis. See Scripts/approximate.py.
#define dfrac(x) \
  4.997559e-01 - \
  3.183096e-01 * sin(6.283185e+00*(x)) - \
  1.591544e-01 * sin(1.256637e+01*(x)) - \
  1.061025e-01 * sin(1.884956e+01*(x)) - \
  7.957647e-02 * sin(2.513274e+01*(x))

#define cubic(t) ((t) * (t) * (3.0 - 2.0 * (t)))
#define d_cubic(t) (6.0f * (t) * (1.0f-(t)))

// Macros for transforming normal and tangent using autodiff.
// r3r3 refers to "r3 to r3 transform", aka a mapping between 3d real-valued
// spaces.
#define R3R3_DECLARE_BASIS_VECTORS(xyz) \
  DifferentialPair<float3> dp_x = diffPair(xyz, float3(1, 0, 0)); \
  DifferentialPair<float3> dp_y = diffPair(xyz, float3(0, 1, 0)); \
  DifferentialPair<float3> dp_z = diffPair(xyz, float3(0, 0, 1))

#define R3R3_AUTODIFF_BASIS_VECTORS(fun, ...) \
  DifferentialPair<float3> dp_x_out = fwd_diff(fun)(dp_x, __VA_ARGS__); \
  DifferentialPair<float3> dp_y_out = fwd_diff(fun)(dp_y, __VA_ARGS__); \
  DifferentialPair<float3> dp_z_out = fwd_diff(fun)(dp_z, __VA_ARGS__)

#define R3R3_DEFORM_NORMAL_AND_TANGENT(normal, tangent) \
  float3x3 jacobian = float3x3(                           \
    float3(dp_x_out.d.x, dp_y_out.d.x, dp_z_out.d.x),     \
    float3(dp_x_out.d.y, dp_y_out.d.y, dp_z_out.d.y),     \
    float3(dp_x_out.d.z, dp_y_out.d.z, dp_z_out.d.z)      \
  );                                                      \
  float jac_det = determinant(jacobian);                  \
  float3x3 itjac = inverse(transpose(jacobian), jac_det); \
  normal = mul(itjac, normal) * sign(jac_det);            \
  tangent = mul(jacobian, tangent)

// Syntactic sugar - wraps the previous three macros.
#define R3R3_NORMALS(xyz, normal, tangent, fun, ...)  \
  R3R3_DECLARE_BASIS_VECTORS(xyz);                    \
  R3R3_AUTODIFF_BASIS_VECTORS(fun, __VA_ARGS__);      \
  R3R3_DEFORM_NORMAL_AND_TANGENT(normal, tangent);    \
  xyz = fun(xyz, __VA_ARGS__)

float3x3 inverse(float3x3 m, float det) {
  det = max(1e-6 * abs(det), abs(det)) * sign(det);

  float invDet = 1.0f / det;
  float3x3 inv;

  inv._11 = (m._22 * m._33 - m._23 * m._32) * invDet;
  inv._12 = (m._13 * m._32 - m._12 * m._33) * invDet;
  inv._13 = (m._12 * m._23 - m._13 * m._22) * invDet;
  inv._21 = (m._23 * m._31 - m._21 * m._33) * invDet;
  inv._22 = (m._11 * m._33 - m._13 * m._31) * invDet;
  inv._23 = (m._13 * m._21 - m._11 * m._23) * invDet;
  inv._31 = (m._21 * m._32 - m._22 * m._31) * invDet;
  inv._32 = (m._31 * m._12 - m._11 * m._32) * invDet;
  inv._33 = (m._11 * m._22 - m._12 * m._21) * invDet;

  return inv;
}

// Takes map a 2x2 quad on the xy plane to a tube. The circular cross section
// is on the xz plane.
[Differentiable]
public float3 plane_to_tube(float3 xyz, no_diff float t) {
  float x0 = xyz.x;
  float y0 = xyz.y;
  float z0 = xyz.z;

  float theta = x0 * PI;
  float radius = ((1.0f - z0) / (dabs(t) + 1e-4f)) * sign(t);

  float x = sin(theta / radius) * radius * PI_RCP;
  // The z0 term here is required to make the jacobian invertible.
  float z = z0 + (1.0f - cos(theta / radius)) * radius * PI_RCP;

  return float3(x, y0, z);
}

public void plane_to_tube_normal(inout float3 xyz, inout float3 normal,
    inout float3 tangent, float t) {
  R3R3_NORMALS(xyz, normal, tangent, plane_to_tube, t);
}

[Differentiable]
public float3 seal(float3 xyz, no_diff float A, no_diff float k, no_diff float t) {
  float x = xyz.x;
  float y = xyz.y;
  float z = xyz.z;

  float x0 = x + sin(y * k) * 0.1;
  float y0 = y + sin(x * k * 2) * 0.5 + cos(z);
  float z0 = (z + sin(x * y * k * PI + t)) * A * (1.0 + sin(y * k) * sin(x * k));

  x0 += z0 * 0.1 * sin(z0 * PI + 1.5);
  y0 += z0 * 0.1 * sin(z0 * PI + 1.5);

  x0 -= 0.0;
  y0 -= 1.2;

  return float3(
      x0, y0, z0
  );
}

public void seal_normal(inout float3 xyz, inout float3 normal,
    inout float3 tangent, float A, float k, float t) {
  R3R3_NORMALS(xyz, normal, tangent, seal, A, k, t);
}

[Differentiable]
public float3 sine_wave(float3 xyz,
    no_diff float3 amplitude,
    no_diff float3 direction,
    no_diff float3 k,
    no_diff float3 omega,
    no_diff float t) {
  xyz += amplitude * sin(k * dot(xyz, direction) + omega * t);
  return xyz;
}

public void sine_wave_normal(inout float3 xyz, inout float3 normal, inout float3 tangent,
    float3 amplitude, float3 direction, float3 k, float3 omega, float t) {
  R3R3_NORMALS(xyz, normal, tangent, sine_wave, amplitude, direction, k, omega, t);
}

[Differentiable]
float3 rand3_hash3(float3 p)
{
    // Improved Murmurhash3 by Squirrel Eiserloh (GDC 2017)
    p = float3(dot(p, float3(127.1, 311.7, 74.7)),
               dot(p, float3(269.5, 183.3, 246.1)),
               dot(p, float3(113.5, 271.9, 124.6)));
    return frac(sin(p) * 43758.5453123);
}

[Differentiable]
float rand3_hash1(float3 p)
{
    // Improved Murmurhash3 by Squirrel Eiserloh (GDC 2017)
    float p0 = dot(p, float3(127.1, 311.7, 74.7));
    return frac(sin(p0) * 43758.5453123);
}

// Calculate value noise and jacobian in one shot.
// Based on https://iquilezles.org/articles/morenoise/
[Differentiable]
float3 value_noise(float3 xyz, out float3 dx, out float3 dy, out float3 dz) {
  float3 cell = floor(xyz);
  float3 f = xyz - cell;

  float ux = cubic(f.x);
  float uy = cubic(f.y);
  float uz = cubic(f.z);

  float dux = d_cubic(f.x);
  float duy = d_cubic(f.y);
  float duz = d_cubic(f.z);

  float3 n000 = rand3_hash3(cell + float3(0.0f, 0.0f, 0.0f));
  float3 n001 = rand3_hash3(cell + float3(0.0f, 0.0f, 1.0f));
  float3 n010 = rand3_hash3(cell + float3(0.0f, 1.0f, 0.0f));
  float3 n011 = rand3_hash3(cell + float3(0.0f, 1.0f, 1.0f));
  float3 n100 = rand3_hash3(cell + float3(1.0f, 0.0f, 0.0f));
  float3 n101 = rand3_hash3(cell + float3(1.0f, 0.0f, 1.0f));
  float3 n110 = rand3_hash3(cell + float3(1.0f, 1.0f, 0.0f));
  float3 n111 = rand3_hash3(cell + float3(1.0f, 1.0f, 1.0f));

  float3 n00 = lerp(n000, n001, uz);
  float3 n01 = lerp(n010, n011, uz);
  float3 n10 = lerp(n100, n101, uz);
  float3 n11 = lerp(n110, n111, uz);

  float3 n0 = lerp(n00, n01, uy);
  float3 n1 = lerp(n10, n11, uy);

  float oneMinusUx = 1.0f - ux;
  float oneMinusUy = 1.0f - uy;

  float3 noise_half = lerp(n0, n1, ux);

  float3 dnoise_half_dx = (n1 - n0) * dux;

  float3 dn0_dy = (n01 - n00) * duy;
  float3 dn1_dy = (n11 - n10) * duy;
  float3 dnoise_half_dy = dn0_dy * oneMinusUx + dn1_dy * ux;

  float3 dn00_dz = (n001 - n000) * duz;
  float3 dn01_dz = (n011 - n010) * duz;
  float3 dn10_dz = (n101 - n100) * duz;
  float3 dn11_dz = (n111 - n110) * duz;
  float3 dn0_dz = dn00_dz * oneMinusUy + dn01_dz * uy;
  float3 dn1_dz = dn10_dz * oneMinusUy + dn11_dz * uy;
  float3 dnoise_half_dz = dn0_dz * oneMinusUx + dn1_dz * ux;

  dx = 2.0f * dnoise_half_dx;
  dy = 2.0f * dnoise_half_dy;
  dz = 2.0f * dnoise_half_dz;

  return -1.0f + 2.0f * noise_half;
}

float3 fbm_with_derivatives(float3 xyz,
    no_diff float t,
    no_diff float amplitude,
    no_diff float gain,
    no_diff float lacunarity,
    no_diff float scale,
    no_diff float octaves,
    no_diff float3 velocity,
    out float3 dx,
    out float3 dy,
    out float3 dz) {
  float3 noise = float3(0.0f, 0.0f, 0.0f);
  float gain_i = amplitude;
  float freq_i = scale;

  dx = float3(0.0f, 0.0f, 0.0f);
  dy = float3(0.0f, 0.0f, 0.0f);
  dz = float3(0.0f, 0.0f, 0.0f);

  for (uint i = 0; i < octaves; ++i) {
    float3 dx_i, dy_i, dz_i;
    float3 octave_noise = value_noise((xyz - velocity * t) * freq_i, dx_i, dy_i, dz_i);
    noise += gain_i * octave_noise;
    dx += gain_i * freq_i * dx_i;
    dy += gain_i * freq_i * dy_i;
    dz += gain_i * freq_i * dz_i;
    freq_i *= lacunarity;
    gain_i *= gain;
  }

  dx += float3(1.0f, 0.0f, 0.0f);
  dy += float3(0.0f, 1.0f, 0.0f);
  dz += float3(0.0f, 0.0f, 1.0f);

  return xyz + noise;
}

public float3 fbm(float3 xyz,
    no_diff float t,
    no_diff float amplitude,
    no_diff float gain,
    no_diff float lacunarity,
    no_diff float scale,
    no_diff float octaves,
    no_diff float3 velocity) {
  float3 dx_unused, dy_unused, dz_unused;
  return fbm_with_derivatives(xyz, t, amplitude, gain, lacunarity, scale, octaves,
      velocity, dx_unused, dy_unused, dz_unused);
}

public void fbm_normal(inout float3 xyz, inout float3 normal, inout float3 tangent,
    no_diff float t,
    no_diff float amplitude, no_diff float gain, no_diff float lacunarity,
    no_diff float scale, no_diff float octaves, no_diff float3 velocity) {
  float3 dx, dy, dz;
  xyz = fbm_with_derivatives(xyz, t, amplitude, gain, lacunarity, scale, octaves,
      velocity, dx, dy, dz);
  float3x3 jac = float3x3(
      dx.x, dy.x, dz.x,
      dx.y, dy.y, dz.y,
      dx.z, dy.z, dz.z);
  float jac_det = determinant(jac);
  float3x3 itjac = inverse(transpose(jac), jac_det);
  normal = mul(itjac, normal) * sign(jac_det);
  tangent = mul(jac, tangent);
}

#endif  // __CUSTOM31_INC