////////////////////////////////////////////////////////////
//~ Helpers

f32 V_RandFromPos(Vec3 pos)
{
  Texture3D<u32> noise3d = G_Dereference<u32>(V_ShaderConst_NoiseTex);
  // TODO: Compile-time noise dims
  u32 noise = noise3d[(Vec3U32)pos % countof(noise3d)];
  f32 rand = Norm16(noise);
  return rand;
}

Vec4 V_ColorFromParticle(V_ParticleKind particle_kind, u32 particle_idx, u32 density, f32 dryness)
{
  V_ParticleDesc desc = V_DescFromParticleKind(particle_kind);

  Vec4 result = 0;
  u64 seed = MixU64(V_ParticleColorBasis ^ particle_idx);
  f32 rand_color = Norm16(seed >> 0);

  result = desc.color;

  // // FIXME: Base color on particle desc
  // if (particle_kind == V_ParticleKind_Test)
  // {
  //   // result.rgb = Vec3(0, 0, 0);
  //   result = LinearFromSrgb(Vec4(0.5, 0.1, 0.1, 0.5));
  // }
  // else if (particle_kind == V_ParticleKind_Debris)
  // {
  //   result = Color_Orange;
  // }
  // else if (particle_kind == V_ParticleKind_Smoke)
  // {
  //   result = Vec4(0.15, 0.15, 0.15, 1);
  // }

  // Apply density
  {
    if (particle_kind == V_ParticleKind_Smoke)
    {
      // f32 t = saturate(density / 10.0);
      f32 t = smoothstep(-10, 32, density);
      // f32 t = smoothstep(0, 2, (f32)density);
      result.a = lerp(0, 0.85, t);
    }
    else if (particle_kind == V_ParticleKind_BloodTrail || particle_kind == V_ParticleKind_BloodDebris)
    {
      // f32 t = (f32)density / 5;
      // t = pow(t, 2);
      // t = saturate(t);
      // result.rgb *= 1.0 - (t * 0.9);

      f32 t = (f32)density / 5;
      // t = smoothstep(-10, 10, t);
      // t = smoothstep(-5, 5, t);
      t = smoothstep(0, 50, t);
      // result.rgb *= 1.0 - (t * 0.9);

      // result.a = t;
      result.a += (1.0 - result.a) * (t);
    }
  }

  result.rgb = saturate(result.rgb + (rand_color - 0.5) * 0.05);
  // result.a += (rand_alpha - 0.5) * 0.025;
  // result.a *= rand_alpha;

  // Apply dryness
  result.rgb *= 1.0 - (dryness * 0.75);

  return result;
}

////////////////////////////////////////////////////////////
//~ Prepare frame

ComputeShader2D(V_PrepareShadeCS, 8, 8)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  RWTexture2D<Vec4> shade = G_Dereference<Vec4>(frame.shade_rw);
  Vec2 shade_pos = SV_DispatchThreadID + 0.5;
  if (all(shade_pos < countof(shade)))
  {
    // Clear shade
    shade[shade_pos] = 0;
  }
}

//- Prepare cells
ComputeShader2D(V_PrepareCellsCS, 8, 8)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  Texture2D<P_TileKind> tiles = G_Dereference<P_TileKind>(frame.tiles);
  RWTexture2D<u32> stain_cells = G_Dereference<u32>(frame.stain_cells);
  RWTexture2D<u32> ground_cells = G_Dereference<u32>(frame.ground_cells);
  RWTexture2D<u32> air_cells = G_Dereference<u32>(frame.air_cells);
  RWTexture2D<u32> stain_densities = G_Dereference<u32>(frame.stain_densities);
  RWTexture2D<u32> ground_densities = G_Dereference<u32>(frame.ground_densities);
  RWTexture2D<u32> air_densities = G_Dereference<u32>(frame.air_densities);
  RWTexture2D<f32> drynesses = G_Dereference<f32>(frame.drynesses);
  RWTexture2D<u32> occluders = G_Dereference<u32>(frame.occluders);

  Vec2 cell_pos = SV_DispatchThreadID + 0.5;
  if (all(cell_pos < countof(air_cells)))
  {
    Vec2 world_pos = mul(frame.af.cell_to_world, Vec3(cell_pos, 1));
    Vec2 tile_pos = mul(frame.af.world_to_tile, Vec3(world_pos, 1));
    P_TileKind tile = tiles[tile_pos];

    // Update stains
    if (frame.should_clear_particles)
    {
      stain_densities[cell_pos] = 0;
      stain_cells[cell_pos] = 0;
      drynesses[cell_pos] = 0;
    }
    else
    {
      u32 stain = stain_cells[cell_pos];
      {
        stain &= ~(1 << 31);
      }
      f32 dryness = drynesses[cell_pos];
      {
        f32 dry_rate = frame.dt * 0.1;
        dryness = lerp(dryness, 1, dry_rate);
      }
      stain_cells[cell_pos] = stain;
      drynesses[cell_pos] = dryness;
    }

    // Clear cells
    ground_cells[cell_pos] = 0;
    air_cells[cell_pos] = 0;

    // Clear densities
    ground_densities[cell_pos] = 0;
    air_densities[cell_pos] = 0;

    // Reset occluders
    V_OccluderKind occluder = V_OccluderKind_None;
    if (tile == P_TileKind_Wall)
    {
      occluder = V_OccluderKind_Wall;
    }
    occluders[cell_pos] = occluder;
  }
}

//- Clear particles
ComputeShader(V_ClearParticlesCS, 64)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  RWStructuredBuffer<V_Particle> particles = G_Dereference<V_Particle>(frame.particles);
  u32 particle_idx = SV_DispatchThreadID;
  if (particle_idx < V_ParticlesCap)
  {
    particles[particle_idx].kind = V_ParticleKind_None;
  }
}

////////////////////////////////////////////////////////////
//~ Quads

//////////////////////////////
//- Vertex shader

VertexShader(V_QuadVS, V_QuadPSInput)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  StructuredBuffer<V_Quad> quads = G_Dereference<V_Quad>(frame.quads);

  V_Quad quad = quads[SV_InstanceID];

  Vec2 rect_uv = RectUvFromIdx(SV_VertexID);
  Vec2 world_pos = mul(quad.quad_uv_to_world_af, Vec3(rect_uv, 1));
  Vec2 screen_pos = mul(frame.af.world_to_screen, Vec3(world_pos, 1));

  Vec2 samp_uv = lerp(quad.tex_slice_uv.p0, quad.tex_slice_uv.p1, rect_uv);

  V_QuadPSInput result;
  result.sv_position = Vec4(NdcFromPos(screen_pos, frame.screen_dims).xy, 0, 1);
  result.quad_idx = SV_InstanceID;
  result.world_pos = world_pos;
  result.samp_uv = samp_uv;
  return result;
}

//////////////////////////////
//- Pixel shader

PixelShader(V_QuadPS, V_QuadPSOutput, V_QuadPSInput input)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  RWTexture2D<u32> occluders = G_Dereference<u32>(frame.occluders);
  StructuredBuffer<V_Quad> quads = G_Dereference<V_Quad>(frame.quads);
  SamplerState clamp_sampler = G_Dereference(frame.pt_clamp_sampler);

  V_Quad quad = quads[input.quad_idx];
  Texture2D<Vec4> tex = G_Dereference<Vec4>(quad.tex);

  Vec2 world_pos = input.world_pos;
  Vec2 cell_pos = mul(frame.af.world_to_cell, Vec3(world_pos, 1));

  b32 is_in_world = all(cell_pos >= 0) && all(cell_pos < countof(occluders));

  Vec4 albedo = tex.Sample(clamp_sampler, input.samp_uv);

  if (is_in_world)
  {
    // TODO: Don't write occluders using screen space result. Do separate draw pass instead.
    if (albedo.a > 0 && quad.occluder != V_OccluderKind_None && is_in_world)
    {
      InterlockedMax(occluders[cell_pos], quad.occluder);
    }
  }

  V_QuadPSOutput output;
  output.sv_target0 = albedo;
  return output;
}

////////////////////////////////////////////////////////////
//~ Particle simulation

//////////////////////////////
//- Particle emitter shader

ComputeShader(V_EmitParticlesCS, 64)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  StructuredBuffer<V_Emitter> emitters = G_Dereference<V_Emitter>(frame.emitters);
  RWStructuredBuffer<V_Particle> particles = G_Dereference<V_Particle>(frame.particles);

  u32 emitter_idx = SV_DispatchThreadID;
  if (emitter_idx < frame.emitters_count)
  {
    V_Emitter emitter = emitters[emitter_idx];
    i32 semantic_particle_kind = V_ParticleKind_None;
    if (emitter.kind > V_ParticleKind_None)
    {
      semantic_particle_kind = (i32)(emitter_idx + 1) * -1;
    }
    for (u32 emitter_particle_idx = 0; emitter_particle_idx < emitter.count; ++emitter_particle_idx)
    {
      u32 particle_idx = (emitter.first_particle_seq + emitter_particle_idx) % (u32)V_ParticlesCap;

      // InterlockedMin guarantees that the highest emitter index (reflected
      // as negative particle kind) will be used to initialize the particle
      // this frame, in case multiple emitters target the same particle (e.g.
      // more particles pushed this frame than are available in the buffer)
      InterlockedMin(particles[particle_idx].kind, semantic_particle_kind);
    }
  }
}

//////////////////////////////
//- Particle sim shader

ComputeShader(V_SimParticlesCS, 64)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  Texture2D<P_TileKind> tiles = G_Dereference<P_TileKind>(frame.tiles);
  RWStructuredBuffer<V_Particle> particles = G_Dereference<V_Particle>(frame.particles);
  RWTexture2D<u32> stain_cells = G_Dereference<u32>(frame.stain_cells);
  RWTexture2D<u32> ground_cells = G_Dereference<u32>(frame.ground_cells);
  RWTexture2D<u32> air_cells = G_Dereference<u32>(frame.air_cells);
  RWTexture2D<u32> stain_densities = G_Dereference<u32>(frame.stain_densities);
  RWTexture2D<u32> ground_densities = G_Dereference<u32>(frame.ground_densities);
  RWTexture2D<u32> air_densities = G_Dereference<u32>(frame.air_densities);
  RWTexture2D<f32> drynesses = G_Dereference<f32>(frame.drynesses);
  RWTexture2D<u32> occluders = G_Dereference<u32>(frame.occluders);

  u32 particle_idx = SV_DispatchThreadID;
  if (particle_idx < V_ParticlesCap)
  {
    V_Particle particle = particles[particle_idx];
    b32 prune = 0;

    //////////////////////////////
    //- Initialize particle

    if (particle.kind != V_ParticleKind_None)
    {
      u64 seed0 = MixU64(V_ParticleSimBasis ^ particle_idx);
      f32 rand_offset = Norm16(seed0 >> 0);
      f32 rand_angle = Norm16(seed0 >> 16);
      f32 rand_speed = Norm16(seed0 >> 32);
      f32 rand_falloff = Norm16(seed0 >> 48);

      u64 seed1 = MixU64(seed0);
      f32 rand_density = Norm16(seed1 >> 0);

      //////////////////////////////
      //- Init

      if (particle.kind < 0)
      {
        u32 emitter_idx = -particle.kind - 1;
        V_Emitter emitter = G_Dereference<V_Emitter>(frame.emitters)[emitter_idx];

        f32 initial_angle = lerp(emitter.angle.min, emitter.angle.max, rand_angle);
        f32 initial_speed = lerp(emitter.speed.min, emitter.speed.max, rand_speed);

        particle = (V_Particle)0;
        particle.kind = emitter.kind;
        particle.life = 0;
        particle.pos = lerp(emitter.pos.p0, emitter.pos.p1, rand_offset);
        particle.velocity = Vec2(cos(initial_angle), sin(initial_angle)) * initial_speed;
      }

      if (particle.kind > V_ParticleKind_None && particle.kind < V_ParticleKind_COUNT)
      {
        V_ParticleDesc desc = V_DescFromParticleKind((V_ParticleKind)particle.kind);

        u32 packed = 0;
        packed |= (particle_idx & ((1 >> 24) - 1)) << 0;
        packed |= (particle.kind & 0xFF) << 24;
        packed |= 1 << 31;
        StaticAssert(V_ParticlesCap <= (1 << 24));  // particle idx must fit in 24 bits
        StaticAssert(V_ParticleKind_COUNT <= 0x7F); // particle kind must fit in 7 bits

        //////////////////////////////
        //- Move

        b32 collision = 0;

        // TODO: Clip to avoid unnecessary iterations outside of world bounds
        {
          Vec2 p0 = particle.pos;
          Vec2 p1 = particle.pos + particle.velocity * frame.dt;
          f32 t = 1;
          {
            Vec2 occluder_p0 = mul(frame.af.world_to_cell, Vec3(p0, 1));
            Vec2 occluder_p1 = mul(frame.af.world_to_cell, Vec3(p1, 1));
            Vec2I32 cell_p0 = floor(occluder_p0);
            Vec2I32 cell_p1 = floor(occluder_p1);

            Vec2 delta = occluder_p1 - occluder_p0;
            Vec2 inv_delta = 1.0 / delta;
            Vec2 dda_step_dir = Vec2((delta.x > 0) - (delta.x < 0), (delta.y > 0) - (delta.y < 0));
            Vec2 t_delta = abs(inv_delta);
            Vec2 t_max = cell_p0 - occluder_p0;
            t_max.x += dda_step_dir.x > 0;
            t_max.y += dda_step_dir.y > 0;
            t_max *= inv_delta;
            t_max = abs(t_max);

            Vec2 t_hit = 0;

            Vec2I32 cell_pos = cell_p0;

            b32 stepped_x = 0;
            b32 stepped_y = 0;

            // TODO: Tune this
            u32 max_iterations = 128;

            b32 done = 0;
            f32 t_diff = 0;
            for (u32 iteration_idx = 0; iteration_idx < max_iterations && !done; ++iteration_idx)
            {
              if (cell_pos.x == cell_p1.x && cell_pos.y == cell_p1.y)
              {
                done = 1;
              }
              else if (t_max.x < t_max.y)
              {
                cell_pos.x += dda_step_dir.x;
                f32 old = t_hit.x;
                t_hit.x = t_max.x - t_delta.x;
                t_diff = t_hit.x - old;
                t_max.x += t_delta.x;
                stepped_x = 1;
                stepped_y = 0;
              }
              else
              {
                cell_pos.y += dda_step_dir.y;
                f32 old = t_hit.y;
                t_hit.y = t_max.y - t_delta.y;
                t_diff = t_hit.y - old;
                t_max.y += t_delta.y;
                stepped_x = 0;
                stepped_y = 1;
              }

              Vec2 cell_screen_pos_p0 = mul(frame.af.world_to_screen, Vec3(mul(frame.af.cell_to_world, Vec3(floor(cell_pos), 1)), 1));
              Vec2 cell_screen_pos_p1 = mul(frame.af.world_to_screen, Vec3(mul(frame.af.cell_to_world, Vec3(ceil(cell_pos), 1)), 1));
              cell_screen_pos_p1 = max(cell_screen_pos_p1, cell_screen_pos_p0 + 1);

              b32 is_in_world = all(cell_pos >= 0) && all(cell_pos < countof(air_cells));
              b32 is_visible = all(cell_screen_pos_p1 >= 0) && all(cell_screen_pos_p0 < frame.screen_dims);

              if (is_in_world)
              {
                f32 stain_delta = abs(t_diff) * desc.stain_rate * frame.dt;
                particle.stain_accum += stain_delta;

                //- Handle collision
                V_OccluderKind occluder = (V_OccluderKind)occluders[cell_pos];
                if (occluder != V_OccluderKind_None)
                {
                  u64 collision_seed = MixU64(V_ParticleCellBasis ^ seed0 ^ particle.cells_count);
                  f32 rand_collision_angle = Norm16(collision_seed >> 0);
                  f32 rand_collision_velocity = Norm16(collision_seed >> 16);
                  f32 rand_collision_penetration = Norm16(collision_seed >> 32);
                  if (rand_collision_penetration >= desc.pen_rate)
                  {
                    collision = 1;
                    done = 1;
                    {
                      if (stepped_x)
                      {
                        if (!AnyBit(desc.flags, V_ParticleFlag_NoReflect))
                        {
                          particle.velocity.x *= -1;
                        }
                        t = saturate(t_hit.x);
                      }
                      else if (stepped_y)
                      {
                        if (!AnyBit(desc.flags, V_ParticleFlag_NoReflect))
                        {
                          particle.velocity.y *= -1;
                        }
                        t = saturate(t_hit.y);
                      }
                      {
                        f32 collision_angle = lerp(-0.05 * Tau, 0.05 * Tau, rand_collision_angle);

                        f32 collision_velocity_falloff = lerp(50, 100, rand_collision_velocity);
                        // f32 collision_velocity_falloff = lerp(5000, 10000, rand_collision_velocity);
                        // f32 collision_velocity_falloff = lerp(500, 10000, rand_collision_velocity);
                        // f32 collision_velocity_falloff = 0;

                        particle.velocity = RotateVec2Angle(particle.velocity, collision_angle);
                        particle.velocity *= 1.0f - saturate(collision_velocity_falloff * frame.dt);
                      }
                    }
                  }

                }

                if (AnyBit(desc.flags, V_ParticleFlag_PruneWhenStill))
                {
                  if (dot(particle.velocity, particle.velocity) < 0.0001)
                  {
                    prune = 1;
                  }
                }

                if (prune && AnyBit(desc.flags, V_ParticleFlag_StainWhenPruned))
                {
                  particle.stain_accum += 1;
                }

                if (!collision)
                {
                  //- Stain
                  u32 stains_count = floor(particle.stain_accum);
                  if (stains_count > 0)
                  {
                    // TODO: Fixed point
                    u32 density = round(stains_count * rand_density);
                    InterlockedMax(stain_cells[cell_pos], packed);
                    InterlockedAdd(stain_densities[cell_pos], density);
                    drynesses[cell_pos] = 0;
                    particle.stain_accum -= stains_count;
                  }

                  //- Draw
                  {
                    b32 should_draw_ground = is_visible && AnyBit(desc.flags, V_ParticleFlag_Ground);
                    b32 should_draw_air = is_visible && AnyBit(desc.flags, V_ParticleFlag_Air);

                    if (should_draw_ground)
                    {
                      InterlockedMax(ground_cells[cell_pos], packed);
                      InterlockedAdd(ground_densities[cell_pos], 1);
                    }

                    if (should_draw_air)
                    {
                      InterlockedMax(air_cells[cell_pos], packed);
                      InterlockedAdd(air_densities[cell_pos], 1);
                    }
                  }
                }
              }
              else
              {
                done = 1;
                prune = 1;
              }

              particle.cells_count += 1;
              iteration_idx += 1;
            }
          }

          f32 falloff = saturate(lerp(10, 20, rand_falloff) * frame.dt);
          // f32 falloff = saturate(lerp(1, 2, rand_falloff) * frame.dt);
          particle.velocity *= 1.0f - falloff;

          particle.pos = p0 + (p1 - p0) * t;
        }

          // Increment life
          particle.life += frame.dt;
      }

      if (prune)
      {
        particle.kind = V_ParticleKind_None;
      }

      particles[particle_idx] = particle;
    }
  }
}

////////////////////////////////////////////////////////////
//~ Shade

// TODO: Remove this

ComputeShader2D(V_ShadeCS, 8, 8)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  RWTexture2D<Vec4> shade_tex = G_Dereference<Vec4>(frame.shade_rw);
  Texture2D<Vec4> albedo_tex = G_Dereference<Vec4>(frame.albedo_ro);
  Texture2D<P_TileKind> tiles = G_Dereference<P_TileKind>(frame.tiles);
  RWTexture2D<u32> stain_cells = G_Dereference<u32>(frame.stain_cells);
  RWTexture2D<f32> drynesses = G_Dereference<f32>(frame.drynesses);
  SamplerState clamp_sampler = G_Dereference(frame.pt_clamp_sampler);

  Vec2 shade_pos = SV_DispatchThreadID + 0.5;
  Vec2 world_pos = mul(frame.af.shade_to_world, Vec3(shade_pos, 1));
  Vec2 cell_pos = mul(frame.af.world_to_cell, Vec3(world_pos, 1));
  Vec2 tile_pos = mul(frame.af.world_to_tile, Vec3(world_pos, 1));

  P_TileKind tile = tiles[tile_pos];

  Vec2 half_world_dims = Vec2(P_WorldPitch, P_WorldPitch) * 0.5;
  b32 is_in_world = all(cell_pos >= 0) && all(cell_pos < countof(stain_cells));

  //////////////////////////////
  //- Compute result

  Vec4 result = 0;

  //////////////////////////////
  //- Write result

  if (all(shade_pos < countof(shade_tex)))
  {
    shade_tex[shade_pos] = result;
  }
}

////////////////////////////////////////////////////////////
//~ Composite

ComputeShader2D(V_CompositeCS, 8, 8)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  // Texture2D<Vec4> shade_tex = G_Dereference<Vec4>(frame.shade_ro);
  Texture2D<Vec4> albedo_tex = G_Dereference<Vec4>(frame.albedo_ro);
  RWTexture2D<Vec4> screen_tex = G_Dereference<Vec4>(frame.screen_rw);
  RWTexture2D<u32> stain_cells = G_Dereference<u32>(frame.stain_cells);
  RWTexture2D<u32> ground_cells = G_Dereference<u32>(frame.ground_cells);
  RWTexture2D<u32> stain_densities = G_Dereference<u32>(frame.stain_densities);
  RWTexture2D<u32> air_cells = G_Dereference<u32>(frame.air_cells);
  RWTexture2D<u32> ground_densities = G_Dereference<u32>(frame.ground_densities);
  RWTexture2D<u32> air_densities = G_Dereference<u32>(frame.air_densities);
  RWTexture2D<f32> drynesses = G_Dereference<f32>(frame.drynesses);
  Texture2D<P_TileKind> tiles = G_Dereference<P_TileKind>(frame.tiles);
  SamplerState clamp_sampler = G_Dereference(frame.pt_clamp_sampler);
  RWStructuredBuffer<V_Particle> particles = G_Dereference<V_Particle>(frame.particles);

  Vec2 screen_pos = SV_DispatchThreadID.xy + 0.5;
  Vec2 world_pos = mul(frame.af.screen_to_world, Vec3(screen_pos, 1));
  Vec2 tile_pos = mul(frame.af.world_to_tile, Vec3(world_pos, 1));
  Vec2 cell_pos = mul(frame.af.world_to_cell, Vec3(world_pos, 1));
  Vec2 shade_pos = mul(frame.af.screen_to_shade, Vec3(screen_pos.xy, 1));

  Vec2 half_world_dims = Vec2(P_WorldPitch, P_WorldPitch) * 0.5;
  Vec2 world_bounds_screen_p0 = mul(frame.af.world_to_screen, Vec3(-half_world_dims.xy, 1));
  Vec2 world_bounds_screen_p1 = mul(frame.af.world_to_screen, Vec3(half_world_dims.xy, 1));
  b32 is_in_world = all(cell_pos >= 0) && all(cell_pos < countof(ground_cells));
  b32 is_in_screen = all(screen_pos >= 0) && all(screen_pos < countof(screen_tex));

  P_TileKind tile = tiles[tile_pos];
  P_TileKind equipped_tile = frame.equipped_tile;

  //////////////////////////////
  //- World color

  Vec4 world_color = Vec4(0.025, 0.025, 0.025, 1);
  if (is_in_world)
  {
    //////////////////////////////
    //- Shade color

    Vec4 shade_color = 0;
    // if (all(shade_pos >= Vec2(0, 0)) && all(shade_pos < countof(shade_tex)))
    // {
    //   Vec2 shade_uv = shade_pos / countof(shade_tex);
    //   shade_color = shade_tex.SampleLevel(clamp_sampler, shade_uv, 0);
    // }

    //////////////////////////////
    //- Tile

    // TODO: Remove this

    b32 tile_is_wall = 0;
    Vec4 tile_color = 0;
    {
      P_TileKind tile_tl = tiles[Vec2(tile_pos.x - 0.99, tile_pos.y - 0.99)];
      P_TileKind tile_tr = tiles[Vec2(tile_pos.x + 0.99, tile_pos.y - 0.99)];
      P_TileKind tile_br = tiles[Vec2(tile_pos.x + 0.99, tile_pos.y + 0.99)];
      P_TileKind tile_bl = tiles[Vec2(tile_pos.x - 0.99, tile_pos.y + 0.99)];
      P_TileKind tile_t = tiles[Vec2(tile_pos.x, tile_pos.y - 0.99)];
      P_TileKind tile_r = tiles[Vec2(tile_pos.x + 0.99, tile_pos.y)];
      P_TileKind tile_b = tiles[Vec2(tile_pos.x, tile_pos.y + 0.99)];
      P_TileKind tile_l = tiles[Vec2(tile_pos.x - 0.99, tile_pos.y)];

      f32 tile_edge_dist = Inf;
      P_TileKind edge_tile = tile;
      if (tile_tl != tile) { edge_tile = tile_tl;  tile_edge_dist = min(tile_edge_dist, length(tile_pos - Vec2(floor(tile_pos.x), floor(tile_pos.y)))); }
      if (tile_tr != tile) { edge_tile = tile_tr;  tile_edge_dist = min(tile_edge_dist, length(tile_pos - Vec2(ceil(tile_pos.x), floor(tile_pos.y)))); }
      if (tile_br != tile) { edge_tile = tile_br;  tile_edge_dist = min(tile_edge_dist, length(tile_pos - Vec2(ceil(tile_pos.x), ceil(tile_pos.y)))); }
      if (tile_bl != tile) { edge_tile = tile_bl;  tile_edge_dist = min(tile_edge_dist, length(tile_pos - Vec2(floor(tile_pos.x), ceil(tile_pos.y)))); }
      if (tile_l != tile) { edge_tile = tile_l;  tile_edge_dist = min(tile_edge_dist, frac(tile_pos.x)); }
      if (tile_r != tile) { edge_tile = tile_r;  tile_edge_dist = min(tile_edge_dist, 1.0 - frac(tile_pos.x)); }
      if (tile_t != tile) { edge_tile = tile_t;  tile_edge_dist = min(tile_edge_dist, frac(tile_pos.y)); }
      if (tile_b != tile) { edge_tile = tile_b;  tile_edge_dist = min(tile_edge_dist, 1.0 - frac(tile_pos.y)); }

      if (tile == P_TileKind_Wall)
      {
        Vec4 outer = LinearFromSrgb(Vec4(0.05, 0.05, 0.05, 1));
        Vec4 inner = LinearFromSrgb(Vec4(0.15, 0.15, 0.15, 1));
        tile_color = lerp(outer, inner, smoothstep(0, 1, tile_edge_dist / 0.375));
        tile_is_wall = 1;
      }
      else if (tile != P_TileKind_Empty)
      {
        V_TileDesc tile_desc = frame.tile_descs[tile];
        Texture2D<Vec4> tile_tex = G_Dereference<Vec4>(tile_desc.tex);
        Vec2 tile_samp_uv = lerp(tile_desc.tex_slice_uv.p0, tile_desc.tex_slice_uv.p1, frac(world_pos));
        tile_color = tile_tex.SampleLevel(clamp_sampler, tile_samp_uv, 0);
      }
      // Checkered grid
      else if (tile == P_TileKind_Empty)
      {
        i32 color_idx = 0;
        Vec4 colors[2] = {
          LinearFromSrgb(Vec4(0.30, 0.30, 0.30, 1)),
          LinearFromSrgb(Vec4(0.15, 0.15, 0.15, 1))
        };
        const f32 checker_size = 0.5;
        Vec2 world_pos_modded = fmod(abs(world_pos), Vec2(checker_size * 2, checker_size * 2));
        if (world_pos_modded.x < checker_size)
        {
          color_idx = !color_idx;
        }
        if (world_pos_modded.y < checker_size)
        {
          color_idx = !color_idx;
        }
        if (world_pos.x < 0)
        {
          color_idx = !color_idx;
        }
        if (world_pos.y < 0)
        {
          color_idx = !color_idx;
        }
        tile_color = colors[color_idx];
      }
    }

    //////////////////////////////
    //- Albedo tex

    Vec4 albedo_tex_color = albedo_tex[screen_pos];

    //////////////////////////////
    //- Particles

    Vec4 stain_particle_color = 0;
    Vec4 ground_particle_color = 0;
    Vec4 air_particle_color = 0;
    {
      //- Stain
      {
        {
          u32 packed = stain_cells[cell_pos];
          V_ParticleKind particle_kind = (V_ParticleKind)((packed >> 24) & 0x7F);
          if (particle_kind != V_ParticleKind_None)
          {
            u32 particle_idx = packed & ((1 << 24) - 1);
            u32 density = stain_densities[cell_pos];
            f32 dryness = drynesses[cell_pos];
            stain_particle_color = V_ColorFromParticle(particle_kind, particle_idx, density, dryness);
          }
        }
        stain_particle_color.rgb *= 1.0 - (0.30 * tile_is_wall); // Darken wall stains
        stain_particle_color.rgb *= stain_particle_color.a;
      }
      //- Ground
      {
        {
          u32 packed = ground_cells[cell_pos];
          V_ParticleKind particle_kind = (V_ParticleKind)((packed >> 24) & 0x7F);
          if (particle_kind != V_ParticleKind_None)
          {
            u32 particle_idx = packed & ((1 << 24) - 1);
            u32 density = ground_densities[cell_pos];
            ground_particle_color = V_ColorFromParticle(particle_kind, particle_idx, density, 0);
          }
        }
        ground_particle_color.rgb *= ground_particle_color.a;
      }
      //- Air
      {
        {
          u32 packed = air_cells[cell_pos];
          V_ParticleKind particle_kind = (V_ParticleKind)((packed >> 24) & 0x7F);
          if (particle_kind != V_ParticleKind_None)
          {
            u32 particle_idx = packed & ((1 << 24) - 1);
            u32 density = air_densities[cell_pos];
            air_particle_color = V_ColorFromParticle(particle_kind, particle_idx, density, 0);
          }
        }
        air_particle_color.rgb *= air_particle_color.a;
      }
    }

    //////////////////////////////
    //- Compose world

    // world_color = BlendPremul(shade_color, world_color);
    world_color = BlendPremul(stain_particle_color, world_color);
    world_color = BlendPremul(ground_particle_color, world_color);
    if (!tile_is_wall)
    {
      world_color = BlendPremul(tile_color, world_color);  // Blend ground tile
      world_color = BlendPremul(stain_particle_color, world_color);  // Blend ground stain
      world_color = BlendPremul(ground_particle_color, world_color);  // Blend ground particle
    }
    world_color = BlendPremul(albedo_tex_color, world_color);
    if (tile_is_wall)
    {
      world_color = BlendPremul(tile_color, world_color);  // Blend wall tile
      world_color = BlendPremul(stain_particle_color, world_color);  // Blend wall stain
      world_color = BlendPremul(ground_particle_color, world_color);  // Blend wall particle
    }
    world_color = BlendPremul(air_particle_color, world_color);
  }

  //////////////////////////////
  //- Overlay color

  Vec4 overlay_color = 0;
  {
    f32 half_thickness = 1;

    //////////////////////////////
    //- Tile selection overlay

    Vec4 selection_color = 0;
    if (
      frame.is_editing &&
      frame.edit_mode == V_EditMode_Tile &&
      frame.has_mouse_focus &&
      is_in_world
    )
    {
      Vec4 border_color = LinearFromSrgb(Vec4(1, 1, 1, 1));
      // Vec4 inner_color = LinearFromSrgb(Vec4(0.4, 0.4, 0.4, 0.25));
      Vec4 inner_color = LinearFromSrgb(Vec4(0.4, 0.8, 0.4, 0.6));

      Rng2 screen_selection = frame.screen_selection;
      Rng2 world_selection = frame.world_selection;

      Rng2 tile_selection;
      tile_selection.p0 = floor(mul(frame.af.world_to_tile, Vec3(world_selection.p0, 1)));
      tile_selection.p1 = ceil(mul(frame.af.world_to_tile, Vec3(world_selection.p1, 1)));
      tile_selection.p1 = max(tile_selection.p1, tile_selection.p0 + 1);

      f32 dist = 100000000;
      dist = min(dist, screen_pos.x - screen_selection.p0.x);
      dist = min(dist, screen_pos.y - screen_selection.p0.y);
      dist = min(dist, screen_selection.p1.x - screen_pos.x);
      dist = min(dist, screen_selection.p1.y - screen_pos.y);
      dist = -dist;

      // if (dist >= -half_thickness && dist <= half_thickness)
      // {
      //     selection_color = border_color;
      // }
      // else
      {
        if (
          tile_pos.x >= tile_selection.p0.x &&
          tile_pos.x <= tile_selection.p1.x &&
          tile_pos.y >= tile_selection.p0.y &&
          tile_pos.y <= tile_selection.p1.y
        )
        {
          selection_color = inner_color;
        }
      }

      // Premultiply
      selection_color.rgb *= selection_color.a;
    }

    //////////////////////////////
    //- Grid

    Vec4 grid_color = 0;
    if (is_in_world)
    {
      b32 debug_draw = !!frame.show_console;

      // Grid outline
      if (frame.show_console)
      {
        const Vec4 line_color = LinearFromSrgb(Vec4(1, 1, 1, 0.1));
        Vec2 line_screen_p0 = mul(frame.af.world_to_screen, Vec3(floor(world_pos), 1));
        Vec2 line_screen_p1 = mul(frame.af.world_to_screen, Vec3(ceil(world_pos), 1));
        f32 line_dist = 100000;
        line_dist = min(line_dist, abs(screen_pos.x - line_screen_p0.x));
        line_dist = min(line_dist, abs(screen_pos.x - line_screen_p1.x));
        line_dist = min(line_dist, abs(screen_pos.y - line_screen_p0.y));
        line_dist = min(line_dist, abs(screen_pos.y - line_screen_p1.y));
        if (line_dist <= half_thickness * 0.5)
        {
          grid_color = line_color;
        }
      }

      // Axis
      if (frame.show_console)
      {
        const Vec4 x_axis_color = LinearFromSrgb(Vec4(0.75, 0, 0, 1));
        const Vec4 y_axis_color = LinearFromSrgb(Vec4(0, 0.75, 0, 1));

        Vec2 zero_screen = mul(frame.af.world_to_screen, Vec3(0, 0, 1));
        f32 x_dist = abs(screen_pos.x - zero_screen.x);
        f32 y_dist = abs(screen_pos.y - zero_screen.y);
        if (y_dist <= half_thickness)
        {
          grid_color = x_axis_color;
        }
        else if (x_dist <= half_thickness)
        {
          grid_color = y_axis_color;
        }
      }

      // World bounds
      {
        const Vec4 bounds_color = LinearFromSrgb(Vec4(0.75, 0.75, 0, 1));
        f32 bounds_dist = 100000;
        bounds_dist = min(bounds_dist, abs(screen_pos.x - world_bounds_screen_p0.x));
        bounds_dist = min(bounds_dist, abs(screen_pos.x - world_bounds_screen_p1.x));
        bounds_dist = min(bounds_dist, abs(screen_pos.y - world_bounds_screen_p0.y));
        bounds_dist = min(bounds_dist, abs(screen_pos.y - world_bounds_screen_p1.y));
        if (bounds_dist <= half_thickness)
        {
          grid_color = bounds_color;
        }
      }

      // Premultiply
      grid_color.rgb *= grid_color.a;
    }

    //////////////////////////////
    //- Crosshair

    // TODO: Remove this
    // TODO: Move to final step after post-processing pass

    Vec4 crosshair_color = 0;
    if (!frame.is_editing)
    {
      f32 dist = length(frame.screen_crosshair - screen_pos);
      if (dist < 4)
      {
        // Adaptive crosshair color based on underlying luminance
        f32 world_luminance = LuminanceFromColor(world_color);
        f32 adaptive_threshold = 0.5;
        Vec4 adapted_crosshair_color = crosshair_color;
        if (world_luminance <= adaptive_threshold)
        {
          crosshair_color = Color_White;
        }
        else
        {
          crosshair_color = InvertColor(Color_White);
        }
        crosshair_color = Premul(crosshair_color);
      }
    }

    //////////////////////////////
    //- Compose overlay

    overlay_color = BlendPremul(selection_color, overlay_color);
    overlay_color = BlendPremul(grid_color, overlay_color);
    overlay_color = BlendPremul(crosshair_color, overlay_color);
  }

  //////////////////////////////
  //- Compose result

  Vec4 result = Vec4(0, 0, 0, 1);
  result = BlendPremul(world_color, result);
  result = BlendPremul(overlay_color, result);

  result = Unpremul(result);

  if (is_in_screen)
  {
    screen_tex[screen_pos] = result;
  }
}

////////////////////////////////////////////////////////////
//~ Debug shapes

//////////////////////////////
//- Vertex shader

VertexShader(V_DVertVS, V_DVertPSInput)
{
  V_SharedFrame frame = G_Dereference<V_SharedFrame>(V_ShaderConst_Frame)[0];
  StructuredBuffer<V_DVert> verts = G_Dereference<V_DVert>(frame.dverts);

  V_DVert vert = verts[SV_VertexID];

  Vec2 screen_pos = vert.pos;

  V_DVertPSInput result;
  result.sv_position = Vec4(NdcFromPos(screen_pos, frame.screen_dims).xy, 0, 1);
  result.color_lin = vert.color_lin;
  return result;
}

//////////////////////////////
//- Pixel shader

PixelShader(V_DVertPS, V_DVertPSOutput, V_DVertPSInput input)
{
  V_DVertPSOutput output;
  output.sv_target0 = input.color_lin;
  return output;
}