#ifndef __PBR_INC
#define __PBR_INC

#include "data.cginc"
#include "decal.cginc"
#include "filamented.cginc"
#include "instancing.cginc"
#include "interpolators.cginc"
#include "letter_grid.cginc"
#include "texture_utils.cginc"

#if defined(_PARALLAX_HEIGHTMAP)
float2 parallax_offset(float2 uv, float3 view_dir_world, float3x3 tbn) {
  float3 view_dir_tangent = mul(tbn, view_dir_world);
  float view_z = max(view_dir_tangent.z, 1e-3f);
  float2 uv_step = view_dir_tangent.xy / view_z * _Parallax_Heightmap_Scale;

#if defined(_PARALLAX_HEIGHTMAP_RAY_MARCHING)
  // Adapt steps by angle to keep cost down while preserving glancing detail.
  float angle = saturate(view_z);
  float base_steps = _Parallax_Heightmap_Ray_Marching_Steps;
  float step_count = lerp(base_steps * 1.5, base_steps * 0.75, angle);
  step_count = clamp(step_count, 2.0, max(base_steps, 2.0));

  float2 delta_uv = uv_step / step_count;
  float delta_depth = 1.0 / step_count;

  float2 cur_uv = uv;
  float2 prev_uv = uv;
  float cur_depth = 0.0;
  float cur_height = 1.0 - _Parallax_Heightmap.Sample(linear_repeat_s,
      cur_uv * _Parallax_Heightmap_ST.xy + _Parallax_Heightmap_ST.zw).r;
  cur_height = (cur_height - (1.0 - _Parallax_Heightmap_Bias));

  // If starting inside geometry, march backwards
  bool inside = cur_depth < cur_height;
  if (!inside) {
    delta_uv = -delta_uv;
    delta_depth = -delta_depth;
  }

  float prev_depth = cur_depth;
  float prev_height = cur_height;

  for (int i = 0; i < (int)step_count; i++) {
    bool was_inside = cur_depth < cur_height;
    if (was_inside != inside) break;

    prev_depth = cur_depth;
    prev_height = cur_height;
    prev_uv = cur_uv;

    cur_uv += delta_uv;
    cur_depth += delta_depth;
    cur_height = 1.0 - _Parallax_Heightmap.Sample(linear_repeat_s,
        cur_uv * _Parallax_Heightmap_ST.xy + _Parallax_Heightmap_ST.zw).r;
    cur_height = (cur_height - (1.0 - _Parallax_Heightmap_Bias));
  }

  // Short binary refine between last two samples to tighten the hit
  float before = prev_height - prev_depth;
  float2 low_uv = prev_uv;
  float low_depth = prev_depth;
  float low_sign = before;
  float2 high_uv = cur_uv;
  float high_depth = cur_depth;

  [unroll(2)]
  for (int j = 0; j < 2; j++) {
    float mid_height = 1.0 - _Parallax_Heightmap.Sample(linear_repeat_s,
        0.5 * (low_uv + high_uv) * _Parallax_Heightmap_ST.xy + _Parallax_Heightmap_ST.zw).r;
    mid_height = (mid_height - (1.0 - _Parallax_Heightmap_Bias));
    float mid_depth = 0.5 * (low_depth + high_depth);
    float mid_sign = mid_height - mid_depth;
    if (mid_sign == 0.0 || sign(mid_sign) == sign(low_sign)) {
      low_uv = 0.5 * (low_uv + high_uv);
      low_depth = mid_depth;
      low_sign = mid_sign;
    } else {
      high_uv = 0.5 * (low_uv + high_uv);
      high_depth = mid_depth;
    }
  }

  float2 refine_uv = 0.5 * (low_uv + high_uv);
  return refine_uv - uv;
#else
  float2 heightmap_uv = uv * _Parallax_Heightmap_ST.xy + _Parallax_Heightmap_ST.zw;
  float height = _Parallax_Heightmap.Sample(linear_repeat_s, heightmap_uv).r;
  height = saturate(height - _Parallax_Heightmap_Bias);

  return uv_step * height;
#endif
}
#endif  // _PARALLAX_HEIGHTMAP

// Tokuyoshi and Kaplanyan 2019 "Improved Geometric Specular Antialiasing"
float normalFiltering(float3 geometric_normal, float perceptual_roughness) {
  float3 du = ddx(geometric_normal);
  float3 dv = ddy(geometric_normal);

  float variance = _Specular_AA_Variance * (dot(du, du) + dot(dv, dv));
  float roughness = perceptual_roughness * perceptual_roughness;
  float kernel_roughness = min(2.0f * variance, _Specular_AA_Threshold);
  float square_roughness = saturate(roughness * roughness + kernel_roughness);

  return saturate(sqrt(sqrt(square_roughness)));
}

void propagateSmoothness(inout Pbr pbr) {
  pbr.smoothness = 1.0f - normalFiltering(pbr.geometric_normal, 1.0f - pbr.smoothness);

  pbr.roughness_perceptual = clamp(1.0f - pbr.smoothness, MIN_PERCEPTUAL_ROUGHNESS, 1);
  pbr.roughness = clamp(pbr.roughness_perceptual * pbr.roughness_perceptual, MIN_ROUGHNESS, 1);
#if defined(_CLEARCOAT)
  pbr.cc_roughness_perceptual = clamp(pbr.cc_roughness_perceptual, MIN_PERCEPTUAL_ROUGHNESS, 1);
  pbr.cc_roughness_perceptual = normalFiltering(pbr.geometric_normal, pbr.cc_roughness_perceptual);
  pbr.cc_roughness = max(MIN_ROUGHNESS, pbr.cc_roughness_perceptual * pbr.cc_roughness_perceptual);
#endif
}

void apply_marble(float3 world_pos, inout float3 albedo) {
#if defined(_MARBLE)
  float3 uvw = world_pos * _Marble_Scale;

  float3 offset = 0;
#if defined(_MARBLE_TIME)
  offset += _Time[0] * _Marble_Speed * _Marble_Direction;
#endif
#if defined(_MARBLE_OFFSET)
  offset += _Marble_Offset;
#endif

  uvw += offset;

  float3 noise = domain_warp_procedural(uvw, _Marble_Strength,
      _Marble_Octaves, _Marble_Lacunarity, _Marble_Gain);

  noise = _Marble_Post_Ramp.Sample(linear_clamp_s, float2(noise.x, 0));

  albedo = noise;
#endif
}

void apply_kintsugi(float3 world_pos, inout float3 albedo, inout float smoothness, inout float metallic) {
#if defined(_KINTSUGI)
  float3 uvw = world_pos * _Kintsugi_Scale;

#if defined(_KINTSUGI_DOMAIN_WARPING)
  float3 warp = domain_warp_procedural(
      world_pos * _Kintsugi_Domain_Warping_Scale,
      _Kintsugi_Domain_Warping_Strength, _Kintsugi_Domain_Warping_Octaves,
      _Kintsugi_Domain_Warping_Lacunarity, _Kintsugi_Domain_Warping_Gain);
  uvw += warp;
#endif

  float mask = voronoi_edge_distance(uvw) + 0.5f;
  float width = max(fwidth(mask) * 0.5f, _Kintsugi_Width);
  float threshold = _Kintsugi_Threshold;
  mask = smoothstep(threshold - width, threshold + width, mask);

#if defined(_KINTSUGI_NOISE_INVERT)
  mask = 1.0f - mask;
#endif  // _KINTSUGI_NOISE_INVERT

  albedo = lerp(albedo, _Kintsugi_Color, mask);
  smoothness = lerp(smoothness, _Kintsugi_Smoothness, mask);
  metallic = lerp(metallic, _Kintsugi_Metallic, mask);
#endif  // _KINTSUGI
}

void apply_letter_grid(v2f i, inout Pbr pbr) {
#if defined(_LETTER_GRID)
  LetterGridOutput lg = LetterGrid(i);
  pbr.albedo.rgb = lerp(pbr.albedo.rgb, lg.albedo, lg.albedo.a);
  pbr.metallic = lerp(pbr.metallic, lg.metallic, lg.albedo.a);
  pbr.roughness = lerp(pbr.roughness, lg.roughness, lg.albedo.a);
#if defined(FORWARD_BASE_PASS)
  pbr.emission += lg.emission * lg.albedo.a;
#endif
#endif
}

#if defined(_BURLEY_TILING)
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;
};

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) * 0.5f;
  // 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;
  float2x2 rot = float2x2(cos(theta), -sin(theta), sin(theta), cos(theta));
  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;
  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

void apply_burley_tiling(v2f i, inout Pbr pbr, inout float3 normal_tangent) {
#if defined(_BURLEY_TILING)
  // Center at 0.
  float2 uv = i.uv01.xy - 0.5;
  // Scale so that any rotation remains within [0, 1] bounds.
  uv *= TWO_OVER_SQRT_3;
  uv /= _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;
  float3 weights = burley_apply_blend_gamma(baryc, _Burley_Tiling_Blend_Gamma);
  float2 uv_dx = ddx(uv);
  float2 uv_dy = ddy(uv);
  BurleyPatchTransform patch_0_transform = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_0, input_scale);
  BurleyPatchTransform patch_1_transform = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_1, input_scale);
  BurleyPatchTransform patch_2_transform = burley_make_patch_transform(uv, uv_dx, uv_dy, vertex_2, input_scale);

  float4 patch_0 = burley_sample_patch(_Burley_Tiling_Maintex, patch_0_transform);
  float4 patch_1 = burley_sample_patch(_Burley_Tiling_Maintex, patch_1_transform);
  float4 patch_2 = burley_sample_patch(_Burley_Tiling_Maintex, patch_2_transform);
  float4 gaussian_blend = patch_0 * weights.x + patch_1 * weights.y + patch_2 * weights.z;
  pbr.albedo.xyz = burley_degaussianize(_Burley_Tiling_Maintex_LUT, burley_apply_soft_clipping(gaussian_blend.rgb, weights));

#if defined(_BURLEY_TILING_SMOOTHNESS)
  float4 patch_0_smoothness = burley_sample_patch(_Burley_Tiling_Smoothness_Map, patch_0_transform);
  float4 patch_1_smoothness = burley_sample_patch(_Burley_Tiling_Smoothness_Map, patch_1_transform);
  float4 patch_2_smoothness = burley_sample_patch(_Burley_Tiling_Smoothness_Map, patch_2_transform);
  float4 smoothness_blend = patch_0_smoothness * weights.x + patch_1_smoothness * weights.y + patch_2_smoothness * weights.z;
  pbr.smoothness = burley_degaussianize(_Burley_Tiling_Smoothness_Map_LUT, burley_apply_soft_clipping(smoothness_blend, weights)).r;
#if defined(_BURLEY_TILING_SMOOTHNESS_INVERT)
  pbr.smoothness = 1.0f - pbr.smoothness;
#endif  // _BURLEY_TILING_SMOOTHNESS_INVERT
#endif  // _BURLEY_TILING_SMOOTHNESS

#if defined(_BURLEY_TILING_NORMAL)
  // TODO whiteout blending?
  float4 patch_0_normal = burley_sample_patch(_Burley_Tiling_Normal_Map, patch_0_transform);
  float4 patch_1_normal = burley_sample_patch(_Burley_Tiling_Normal_Map, patch_1_transform);
  float4 patch_2_normal = burley_sample_patch(_Burley_Tiling_Normal_Map, patch_2_transform);
  float4 normal_blend = patch_0_normal * weights.x + patch_1_normal * weights.y + patch_2_normal * weights.z;
  normal_tangent = burley_degaussianize(_Burley_Tiling_Normal_Map_LUT, burley_apply_soft_clipping(normal_blend.rgb, weights), false);
  normal_tangent.xy = normal_tangent.xy * 2 - 1;
  normal_tangent.xy *= _Burley_Tiling_Normal_Strength;
  normal_tangent = normalize(normal_tangent);
#endif  // _BURLEY_TILING_NORMAL

#endif  // _BURLEY_TILING
}

Pbr getPbr(v2f i) {
  Pbr pbr = (Pbr) 0;

  float3 n = i.normal;
  pbr.geometric_normal = n;
  float3 t = i.tangent.xyz;
  t = normalize(t - n * dot(n, t));  // Gram-Schmidt to avoid skew
  float3 b = normalize(cross(n, t)) * i.tangent.w;
  pbr.tbn = float3x3(t, b, n);

#if defined(_UV_SCROLL)
  i.uv01.xy += getTime() * _UV_Scroll_Speed;
#endif  // _UV_SCROLL

#if defined(_PARALLAX_HEIGHTMAP)
  float2 uv_parallax = i.uv01.xy + parallax_offset(i.uv01.xy, normalize(i.eyeVec.xyz), pbr.tbn);
#else
  float2 uv_parallax = i.uv01.xy;
#endif  // _PARALLAX_HEIGHTMAP

#if defined(OUTLINES_PASS) && defined(_OUTLINES)
  pbr.albedo = _Outlines_Color;
#else
  pbr.albedo = _MainTex.Sample(aniso4_trilinear_repeat_s, uv_parallax * _MainTex_ST.xy + _MainTex_ST.zw);
  pbr.albedo *= _Color;
#endif

  float3 normal_tangent = UnpackNormal(_BumpMap.Sample(aniso4_trilinear_repeat_s, uv_parallax * _BumpMap_ST.xy));
  normal_tangent.xy *= _BumpScale;

#if defined(_DETAILS)
  float2 detail_uv = get_uv_by_channel(i, _Details_UV_Channel);
  float3 detail_normal = UnpackNormal(_DetailNormalMap.Sample(aniso4_trilinear_repeat_s, detail_uv * _DetailNormalMap_ST.xy));
  detail_normal.xy *= _DetailNormalMapScale;
  float detail_mask = _DetailMask.Sample(aniso4_trilinear_repeat_s, detail_uv * _DetailMask_ST.xy).r;
  detail_normal.xy *= detail_mask;
  normal_tangent = blendNormalsHill12(normal_tangent, detail_normal);
#endif

  float4 metallic_gloss = _MetallicGlossMap.Sample(aniso4_trilinear_repeat_s, uv_parallax * _MetallicGlossMap_ST.xy);
  pbr.smoothness = metallic_gloss.a * _Glossiness;
  pbr.metallic = metallic_gloss.r * _Metallic;

  apply_marble(i.worldPos, pbr.albedo.xyz);
  apply_kintsugi(i.worldPos, pbr.albedo.xyz, pbr.smoothness, pbr.metallic);
  apply_letter_grid(i, pbr);
  apply_burley_tiling(i, pbr, normal_tangent);

  applyDecals(i, pbr, normal_tangent);

  pbr.normal = normalize(mul(normal_tangent, pbr.tbn));

#if defined(_BENT_NORMALS)
  float3 bent_ts = UnpackNormal(_Bent_Normals_Map.Sample(
        aniso4_trilinear_repeat_s, uv_parallax * _Bent_Normals_Map_ST.xy +
        _Bent_Normals_Map_ST.zw));
  pbr.bent_normal = normalize(mul(bent_ts, pbr.tbn));
#endif

#if defined(_CLEARCOAT)
  pbr.cc_roughness = _Clearcoat_Roughness;
  pbr.cc_roughness_perceptual = sqrt(pbr.cc_roughness);
  pbr.cc_strength = _Clearcoat_Strength;
#if defined(_CLEARCOAT_NORMALS)
  float3 cc_normal_ts = UnpackNormal(_Clearcoat_Normals.Sample(
        aniso4_trilinear_repeat_s, uv_parallax * _Clearcoat_Normals_ST.xy +
        _Clearcoat_Normals_ST.zw));
  cc_normal_ts.xy *= _Clearcoat_Normals_Strength;
  pbr.cc_normal = normalize(mul(cc_normal_ts, pbr.tbn));
#else
  pbr.cc_normal = i.normal;
#endif  // _CLEARCOAT_NORMALS
#endif  // _CLEARCOAT
  propagateSmoothness(pbr);

#if defined(_EMISSIONS) && defined(FORWARD_BASE_PASS)
  float3 emission_tint = _EmissionColor;
  float3 emission_color = _EmissionMap.Sample(trilinear_repeat_s,
      i.uv01.xy * _EmissionMap_ST.xy + _EmissionMap_ST.zw);
  float emission_mask = _EmissionMask.Sample(trilinear_repeat_s,
      i.uv01.xy * _EmissionMask_ST.xy + _EmissionMask_ST.zw).r;
  pbr.emission = emission_tint * emission_color * emission_mask;
#endif

  return pbr;
}

#endif  // __PBR_INC