#ifndef __BURLEY_INC
#define __BURLEY_INC

#include "data.cginc"
#include "math.cginc"

#if defined(_BURLEY_TILING) || defined(TRIPLANAR_BURLEY)
float2 burley_tri_to_cart(float2 tri_coord) {
  return float2(
      tri_coord.x + tri_coord.y * 0.5f,
      tri_coord.y * SQRT_3_OVER_2);
}

float3 burley_apply_blend_gamma(float3 weights, float gamma) {
  weights = pow(weights, gamma);
  return weights / (weights.x + weights.y + weights.z);
}

// Equation 4 (first half).
float3 burley_soft_clipping_lower_half(float3 x_hat, float w_hat) {
  float linear_start = 0.25f * (2.0f - w_hat);
  float3 linear_value = (x_hat - 0.5f) / w_hat + 0.5f;
  float3 linear_mask = step(float3(linear_start, linear_start, linear_start), x_hat);

  if (w_hat >= TWO_OVER_THREE) {
    float3 t = x_hat / (2.0f - w_hat);
    float3 quadratic = 8.0f * (1.0f / w_hat - 1.0f) * t * t + (3.0f - 2.0f / w_hat) * t;
    return lerp(quadratic, linear_value, linear_mask);
  }

  float quadratic_start = 0.25f * (2.0f - 3.0f * w_hat);
  float3 d = (x_hat - quadratic_start) / w_hat;
  float3 quadratic = d * d;
  float3 quadratic_mask = step(float3(quadratic_start, quadratic_start, quadratic_start), x_hat);
  float3 result = quadratic * quadratic_mask;
  return lerp(result, linear_value, linear_mask);
}

// Equation 4.
float3 burley_soft_clipping_contrast(float3 x_hat, float w_hat) {
  float3 upper_mask = step(0.5f, x_hat);
  float3 lower_x = min(x_hat, 1.0f - x_hat);
  float3 lower_y = burley_soft_clipping_lower_half(lower_x, w_hat);
  return lerp(lower_y, 1.0f - lower_y, upper_mask);
}

float3 burley_apply_soft_clipping(float3 gaussian_color, float3 weights) {
  float w_hat = sqrt(dot(weights, weights));
  return burley_soft_clipping_contrast(gaussian_color, w_hat);
}

float3 burley_degaussianize(texture2D lut, float3 gaussian_color, bool decode_srgb = true) {
  float2 uv_r = float2(gaussian_color.r, 0.5f);
  float2 uv_g = float2(gaussian_color.g, 0.5f);
  float2 uv_b = float2(gaussian_color.b, 0.5f);
  float lut_r = lut.Sample(linear_clamp_s, uv_r).r;
  float lut_g = lut.Sample(linear_clamp_s, uv_g).g;
  float lut_b = lut.Sample(linear_clamp_s, uv_b).b;
  float3 restored = float3(lut_r, lut_g, lut_b);
  return decode_srgb ? srgb_to_linear(restored) : restored;
}

struct BurleyPatchTransform {
  float2 uv;
  float2 dx;
  float2 dy;
  float2x2 uv_to_patch;
};

BurleyPatchTransform burley_make_patch_transform(float2 uv, float2 uv_dx, float2 uv_dy,
                                                 float2 tri_vertex, float input_scale) {
  float3 cube_id = float3(tri_vertex.x, tri_vertex.y, -tri_vertex.x - tri_vertex.y);
  float3 tile_rand3 = hash33_fast(cube_id);
  float2 vertex_uv = burley_tri_to_cart(tri_vertex);
  // Map the unit-radius hex support to the unit square so arbitrary rotation
  // stays within bounds.
  float2 local_uv = (uv - vertex_uv);
  // Apply input scaling.
  local_uv *= input_scale;
  float2 sample_dx = uv_dx * (0.5f * input_scale);
  float2 sample_dy = uv_dy * (0.5f * input_scale);
  // Rotate.
  float theta = hash31_ff(tile_rand3) * TAU - PI;
#if defined(_BURLEY_TILING_ROTATION_CONSTRAINT)
  theta *= _Burley_Tiling_Rotation_Constraint;
#endif  // _BURLEY_TILING_ROTATION_CONSTRAINT
  float s;
  float c;
  sincos(theta, s, c);
  float2x2 rot = float2x2(c, -s, s, c);
  local_uv = mul(rot, local_uv);
  sample_dx = mul(rot, sample_dx);
  sample_dy = mul(rot, sample_dy);
  // Apply randomized offset, staying within bounds.
  // The scaled-and-rotated footprint is bounded by [-Input_Scale / 2, Input_Scale / 2],
  // so we can offset by [(1 - Input_Scale) / 2].
  float2 random_offset = (tile_rand3.yz * 2.0f - 1.0f) * (0.5f * (1.0f - input_scale));
  local_uv += random_offset;
  // Finally, remap onto [0, 1].
  local_uv += 0.5f;

  BurleyPatchTransform patch;
  patch.uv = local_uv;
  patch.dx = sample_dx;
  patch.dy = sample_dy;
  patch.uv_to_patch = rot * (0.5f * input_scale);
  return patch;
}

float4 burley_sample_patch(texture2D tex, BurleyPatchTransform patch) {
  return tex.SampleGrad(
      aniso4_trilinear_repeat_s, patch.uv, patch.dx, patch.dy);
}
#endif  // _BURLEY_TILING || TRIPLANAR_BURLEY

#if defined(_BURLEY_TILING)
struct BurleyTilingContext {
  BurleyPatchTransform patch_0;
  BurleyPatchTransform patch_1;
  BurleyPatchTransform patch_2;
  float3 weights;
  float2 base_uv;
  float uv_scale;
};

static BurleyTilingContext _burley_ctx;

void burley_tiling_setup(float2 base_uv) {
  _burley_ctx.base_uv = base_uv;
  float2 uv = base_uv - 0.5;
  // Scale so that any rotation remains within [0, 1] bounds.
  uv *= TWO_OVER_SQRT_3;
  uv /= _Burley_Tiling_Output_Scale;
  _burley_ctx.uv_scale = TWO_OVER_SQRT_3 / _Burley_Tiling_Output_Scale;
  float3 hex_coord = cart_to_hex(uv);
  float2 tri_coord = hex_coord.yz;
  float2 tri_cell = floor(tri_coord);
  float2 tri_frac = tri_coord - tri_cell;
  float2 vertex_0;
  float2 vertex_1;
  float2 vertex_2;
  float3 baryc;
  if (tri_frac.x + tri_frac.y < 1.0f) {
    vertex_0 = tri_cell;
    vertex_1 = tri_cell + float2(1.0f, 0.0f);
    vertex_2 = tri_cell + float2(0.0f, 1.0f);
    baryc = float3(1.0f - (tri_frac.x + tri_frac.y), tri_frac.x, tri_frac.y);
  } else {
    vertex_0 = tri_cell + 1.0f;
    vertex_1 = tri_cell + float2(0.0f, 1.0f);
    vertex_2 = tri_cell + float2(1.0f, 0.0f);
    baryc = float3(tri_frac.x + tri_frac.y - 1.0f, 1.0f - tri_frac.x, 1.0f - tri_frac.y);
  }

  float input_scale = _Burley_Tiling_Input_Scale;
  _burley_ctx.weights = burley_apply_blend_gamma(baryc, _Burley_Tiling_Blend_Gamma);
  float2 uv_dx = ddx(uv);
  float2 uv_dy = ddy(uv);
  _burley_ctx.patch_0 = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_0, input_scale);
  _burley_ctx.patch_1 = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_1, input_scale);
  _burley_ctx.patch_2 = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_2, input_scale);
}
#endif  // _BURLEY_TILING

#if defined(_BURLEY_TILING_HEIGHTMAP) || defined(_BURLEY_TILING_AMBIENT_OCCLUSION)
float burley_sample_scalar(texture2D tex, texture2D lut) {
  float4 patch_0 = burley_sample_patch(tex, _burley_ctx.patch_0);
  float4 patch_1 = burley_sample_patch(tex, _burley_ctx.patch_1);
  float4 patch_2 = burley_sample_patch(tex, _burley_ctx.patch_2);
  float4 gaussian_blend = patch_0 * _burley_ctx.weights.x +
      patch_1 * _burley_ctx.weights.y +
      patch_2 * _burley_ctx.weights.z;
  return burley_degaussianize(
      lut,
      burley_apply_soft_clipping(gaussian_blend.rgb, _burley_ctx.weights),
      false).r;
}
#endif  // _BURLEY_TILING_HEIGHTMAP || _BURLEY_TILING_AMBIENT_OCCLUSION

#endif  // __BURLEY_INC