#ifndef __RAY_MARCHING_INC
#define __RAY_MARCHING_INC

#if defined(_RAY_MARCHING)

#include "math.cginc"
#include "ray_marching_maps.hlsl"
#include "texture_utils.cginc"
#include "vertex.cginc"

struct RayMarchResult {
  float3 objPos;
  float3 objNorm;
  float3 objTan;
};

float3 getObjPos(v2f i)
{
#if defined(_VERTEX_DEFORMATION_FRAGMENT_NORMALS) || defined(_VERTEX_DEFORMATION_TESSELLATION)
  return i.objPos_orig;
#else
  return i.objPos;
#endif
}

void GetRoRd(v2f i, out float3 ro, out float3 rd) {
  ro = getObjPos(i);
  float3 rd_world = i.worldPos.xyz - _WorldSpaceCameraPos;
  rd = normalize(mul(unity_WorldToObject, rd_world));
#if defined(_RAY_MARCHING_SCALING)
  float3 scale = float3(
    _Ray_Marching_Scaling_Factor_X,
    _Ray_Marching_Scaling_Factor_Y,
    _Ray_Marching_Scaling_Factor_Z);
  ro *= scale;
#endif  // _RAY_MARCHING_SCALING
}

float map(float3 p) {
    float d = 1e9;
#if defined(_RAY_MARCHING_BALL)
    {
      float r = _Ray_Marching_Ball_Radius;
      d = min(d, map_ball(p, _Ray_Marching_Ball_Radius));
    }
#endif
#if defined(_RAY_MARCHING_HEXAGON)
    {
      float r = _Ray_Marching_Hexagon_Radius;
      float h = _Ray_Marching_Hexagon_Height;
      d = min(d, map_hexagon(p, float2(r, h)));
    }
#endif
  return d;
}

#if defined(_RAY_MARCHING_HEX_GRID)
float domain_repeat_hex_grid(inout float3 p) {
  const float gridCount = max(_Ray_Marching_Hex_Grid_Count, 1.0f);
  const float invGridCount = 1.0f / gridCount;
  const float3 hex = cart_to_hex(p.xy);
  const float3 scaledHex = hex * gridCount;
  // Cell ID.
  const float3 id = round_hex(scaledHex);
  // Coordinates within the current cell.
  const float3 local = (scaledHex - id) * invGridCount;
  const float3 cubeId = float3(id.y, id.z, -id.y - id.z);
  const float limit = max((gridCount - 1.0f) * 0.5f, 0.0f);
  const float3 cubeLimit = float3(limit, limit, limit);

  float3 bestPos = p;
  float d = 1e9;
  if (!any(abs(cubeId) > cubeLimit)) {
    float3 pp = p;
    pp.xy = hex_to_cart(local);
    const float d_cur = map(pp);
    if (d_cur < d) {
      d = d_cur;
      bestPos = pp;
    }
  }

#if defined(_RAY_MARCHING_CORRECT_REPETITION)
  const float3 cubeOffsets[6] = {
      float3( 1.0f, -1.0f,  0.0f),
      float3( 1.0f,  0.0f, -1.0f),
      float3( 0.0f,  1.0f, -1.0f),
      float3(-1.0f,  1.0f,  0.0f),
      float3(-1.0f,  0.0f,  1.0f),
      float3( 0.0f, -1.0f,  1.0f),
  };

  for (int idx = 0; idx < 6; ++idx) {
    const float3 cube = cubeId + cubeOffsets[idx];

    if (any(abs(cube) > cubeLimit)) {
      continue;
    }

    const float3 rid = float3(cube.x + cube.y, cube.x, cube.y);
    const float3 neighbourHex = (scaledHex - rid) * invGridCount;
    float3 neighbour = p;
    neighbour.xy = hex_to_cart(neighbourHex);
    const float d_cur = map(neighbour);
    if (d_cur < d) {
      d = d_cur;
      bestPos = neighbour;
    }
  }
#endif  // _RAY_MARCHING_CORRECT_REPETITION

  p = bestPos;
  return d;
}
#endif  // _RAY_MARCHING_HEX_GRID

float domain_repeat(inout float3 p) {
  float d;
#if defined(_RAY_MARCHING_CART_GRID)
  {
    const float3 count = float3(
        _Ray_Marching_Cart_Grid_Count_X,
        _Ray_Marching_Cart_Grid_Count_Y,
        _Ray_Marching_Cart_Grid_Count_Z);
    const float3 span = float3(
        _Ray_Marching_Cart_Grid_Span_X,
        _Ray_Marching_Cart_Grid_Span_Y,
        _Ray_Marching_Cart_Grid_Span_Z) * 2.0f;
    const float3 period = span / count;
    const float3 half_period = period * 0.5f;
    // count = 2
    // which = 0 -> ok
    // which = 1 -> ok
    // which = -1 -> not ok
    const float3 max_allowed = count / 2;
    const float3 min_allowed = -count / 2 + 0.5;
    const float3 offset = float3(
        count.x % 2 == 0 ? half_period.x : 0,
        count.y % 2 == 0 ? half_period.y : 0,
        count.z % 2 == 0 ? half_period.z : 0);
    const float3 shifted = p + offset;
    float3 which = floor((shifted + half_period) / period);
    if (any(which > max_allowed || which < min_allowed)) {
      p = 1e6;
    } else {
      p = glsl_mod(shifted + half_period, period) - half_period;
    }
    d = map(p);
  }
#elif defined(_RAY_MARCHING_HEX_GRID)
  d = domain_repeat_hex_grid(p);
#else
  d = map(p);
#endif
  return d;
}

bool get_one_hit(float3 ro, float3 rd, out float3 hitPos, out float hitD) {
  const float kMinDist = _Ray_Marching_Min_Dist;
  const float kMaxDist = _Ray_Marching_Max_Dist;
  const uint kMaxIter = _Ray_Marching_Max_Iter;
  float d_cur;
  float d_acc = 0;
  for (uint ii = 0; ii < kMaxIter; ++ii) {
    float3 p = ro + rd * d_acc;

    d_cur = domain_repeat(p);

#if defined(_RAY_MARCHING_OVERSTEP)
    d_cur *= (d_cur > 0 ? _Ray_Marching_Overstepping_Factor : 1.0f);
#endif

    d_acc += d_cur;
    if (abs(d_cur) < kMinDist) {
      hitPos = p;
      hitD = d_acc;
      return true;
    }
    if (d_acc > kMaxDist) {
      break;
    }
  }

  return false;
}

bool get_hit(float3 ro, float3 rd, out float3 hitPos, out float hitD) {
#if defined(_RAY_MARCHING_CART_INSTANCING)
  const float3 span = float3(
      _Ray_Marching_Cart_Instancing_Span_X,
      _Ray_Marching_Cart_Instancing_Span_Y,
      _Ray_Marching_Cart_Instancing_Span_Z) * 2.0f;
  const float3 count = float3(
      _Ray_Marching_Cart_Instancing_Count_X,
      _Ray_Marching_Cart_Instancing_Count_Y,
      _Ray_Marching_Cart_Instancing_Count_Z);
  const float3 count_minus_1 = count - 1.0f;
  const float3 offset = span / count;
  const float3 midpoint = offset * count_minus_1 * 0.5f;
#if defined(_RAY_MARCHING_CART_INSTANCING_OFFSETS)
  const float x_every_y = _Ray_Marching_Cart_Instancing_Offsets_X_Every_Y;
  const float x_every_z = _Ray_Marching_Cart_Instancing_Offsets_X_Every_Z;
  const float y_every_x = _Ray_Marching_Cart_Instancing_Offsets_Y_Every_X;
  const float y_every_z = _Ray_Marching_Cart_Instancing_Offsets_Y_Every_Z;
  const float z_every_x = _Ray_Marching_Cart_Instancing_Offsets_Z_Every_X;
  const float z_every_y = _Ray_Marching_Cart_Instancing_Offsets_Z_Every_Y;
  const float3 max_offset = float3(
    count_minus_1.y * x_every_y + count_minus_1.z * x_every_z,
    count_minus_1.x * y_every_x + count_minus_1.z * y_every_z,
    count_minus_1.x * z_every_x + count_minus_1.y * z_every_y
  );
  const float3 offset_midpoint = max_offset * 0.5f;
#endif  // _RAY_MARCHING_CART_INSTANCING_OFFSETS

  hitD = 1e9;
  for (uint xi = 0; xi < count.x; ++xi)
  for (uint yi = 0; yi < count.y; ++yi)
  for (uint zi = 0; zi < count.z; ++zi) {
    float3 hitPos_tmp;
    float hitD_tmp;
    float3 cur_pos = ro + offset * float3(xi, yi, zi) - midpoint;
#if defined(_RAY_MARCHING_CART_INSTANCING_OFFSETS)
    const float3 cur_offset = float3(
        x_every_y * yi + x_every_z * zi,
        y_every_x * xi + y_every_z * zi,
        z_every_x * xi + z_every_y * yi);
    cur_pos += cur_offset - offset_midpoint;
#endif  // _RAY_MARCHING_CART_INSTANCING_OFFSETS
    if (!get_one_hit(cur_pos, rd, hitPos_tmp, hitD_tmp)) {
      continue;
    }
    if (hitD_tmp < hitD) {
      hitPos = hitPos_tmp;
      hitD = hitD_tmp;
    }
  }
  return hitD != 1e9;
#else  // !_RAY_MARCHING_CART_INSTANCING
  return get_one_hit(ro, rd, hitPos, hitD);
#endif  // _RAY_MARCHING_CART_INSTANCING
}
#endif  // _RAY_MARCHING

void ray_march(inout v2f i, bool is_fragment) {
#if defined(_RAY_MARCHING)
  float3 ro, rd;
  GetRoRd(i, ro, rd);

#if defined(_VERTEX_DEFORMATION)
  // TODO optimize, we don't need to pass in `dummy`.
  {
    float3 dummy = 0;
    float3 tmp_pos = ro;
    undeform_normal(tmp_pos, dummy, rd);
    rd = normalize(rd);
  }
#endif

  float3 hit_pos;
  float hit_d;
  // Fragment shader can discard pixels, and here it's desired.
  // Vertex shader very well may not have hit anything, but that doesn't mean
  // that the fragment shader can't finish the job.
  if (!get_hit(ro, rd, hit_pos, hit_d) && is_fragment) {
    discard;
  }

  float3 lclPos = hit_pos;
  float3 lclNorm;
  float3 lclTan;
  // TODO loop should say which primitive was hit, and be dispatched here.
#if defined(_RAY_MARCHING_BALL)
  {
    float r = _Ray_Marching_Ball_Radius;
    map_ball_normal(r, lclPos, lclNorm, lclTan);
  }
#elif defined(_RAY_MARCHING_HEXAGON)
  {
    float r = _Ray_Marching_Hexagon_Radius;
    float h = _Ray_Marching_Hexagon_Height;
    map_hexagon_normal(float2(r, h), lclPos, lclNorm, lclTan);
  }
#endif

#if defined(_VERTEX_DEFORMATION)
  if (is_fragment) {
    float3 tmp_pos = ro;
    deform_normal(tmp_pos, lclNorm, lclTan);
  }
#endif

  i.objPos = lclPos;
  if (is_fragment) {
    i.normal = lclNorm;
    i.tangent.xyz = lclTan;
  }
#endif  // _RAY_MARCHING
}

#endif  // __RAY_MARCHING_INC