/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

#pragma once

#include "draw_shader_shared.hh"
#include "draw_view.bsl.hh"
#include "eevee_lightprobe_sphere.bsl.hh"
#include "eevee_surfel.bsl.hh"
#include "gpu_shader_utildefines_lib.glsl"

namespace eevee::surfel {

float avg_albedo(float3 albedo)
{
  return saturate(dot(albedo, float3(1.0f / 3.0f)));
}

struct SurfelRay {
  [[resource_table]] srt_t<LightprobeSphereRenderData> lightprobe_spheres;
  [[resource_table]] srt_t<SurfelData> surfels_data;

  [[push_constant]] const int radiance_src;
  [[push_constant]] const int radiance_dst;

  void radiance_transfer(Surfel &surfel, float3 in_radiance, float in_visibility, float3 L)
  {
    [[resource_table]] SurfelData &surfels = surfels_data;

    /* Clamped brightness. */
    float luma = max(1e-8f, reduce_max(in_radiance));
    in_radiance *= 1.0f - max(0.0f, luma - surfels.capture_info_buf.clamp_indirect) / luma;

    float NL = dot(surfel.normal, L);
    /* Lambertian BSDF. Albedo applied later depending on which side of the surfel was hit. */
    constexpr float bsdf = M_1_PI;
    /* From "Global Illumination using Parallel Global Ray-Bundles"
     * Eq. 3: Outgoing light */
    float transfert_fn = (M_TAU / surfels.capture_info_buf.sample_count) * bsdf * abs(NL);

    SurfelRadiance radiance = surfel.radiance_indirect[radiance_dst];

    float sample_weight = 1.0f / surfels.capture_info_buf.sample_count;
    bool front_facing = (NL > 0.0f);
    if (front_facing) {
      /* Store radiance normalized for spherical harmonic accumulation and for visualization. */
      radiance.front *= radiance.front_weight;
      radiance.front += float4(in_radiance, in_visibility) * transfert_fn *
                        float4(surfel.albedo_front, avg_albedo(surfel.albedo_front));
      radiance.front_weight += sample_weight;
      radiance.front /= radiance.front_weight;
    }
    else {
      /* Store radiance normalized for spherical harmonic accumulation and for visualization. */
      radiance.back *= radiance.back_weight;
      radiance.back += float4(in_radiance, in_visibility) * transfert_fn *
                       float4(surfel.albedo_back, avg_albedo(surfel.albedo_back));
      radiance.back_weight += sample_weight;
      radiance.back /= radiance.back_weight;
    }

    surfel.radiance_indirect[radiance_dst] = radiance;
  }

  void radiance_transfer_surfel(Surfel &receiver, Surfel sender)
  {
    [[resource_table]] SurfelData &surfels = surfels_data;

    float3 L = safe_normalize(sender.position - receiver.position);
    bool front_facing = dot(-L, sender.normal) > 0.0f;

    float4 radiance_vis;
    SurfelRadiance sender_radiance_indirect = sender.radiance_indirect[radiance_src];
    if (front_facing) {
      radiance_vis = sender.radiance_direct.front;
      radiance_vis += sender_radiance_indirect.front * sender_radiance_indirect.front_weight;
    }
    else {
      radiance_vis = sender.radiance_direct.back;
      radiance_vis += sender_radiance_indirect.back * sender_radiance_indirect.back_weight;
    }

    if (!surfels.capture_info_buf.capture_indirect) {
      radiance_vis.rgb = float3(0.0f);
    }

    radiance_transfer(receiver, radiance_vis.rgb, radiance_vis.a, L);
  }

  void radiance_transfer_world(Surfel &receiver, float3 L)
  {
    [[resource_table]] const LightprobeSphereRenderData &lp_spheres = lightprobe_spheres;
    [[resource_table]] SurfelData &surfels = surfels_data;

    float3 radiance = float3(0.0f);
    float visibility = 0.0f;

    if (surfels.capture_info_buf.capture_world_indirect) {
      SphereProbeUvArea atlas_coord = surfels.capture_info_buf.world_atlas_coord;
      radiance = lp_spheres.sample_probe(L, 0.0f, atlas_coord).rgb;
    }

    if (surfels.capture_info_buf.capture_visibility_indirect) {
      visibility = 1.0f;
    }

    radiance_transfer(receiver, radiance, visibility, L);
  }
};

/**
 * For every surfel, compute the incoming radiance from both side.
 * For that, walk the ray surfel linked-list and gather the light from the neighbor surfels.
 * This shader is dispatched for a random ray in a uniform hemisphere as we evaluate the
 * radiance in both directions.
 *
 * Dispatched as 1 thread per surfel.
 */
[[compute, local_size(SURFEL_GROUP_SIZE)]]
void ray_main([[resource_table]] SurfelRay &srt,
              [[resource_table]] const draw::View &views,
              [[global_invocation_id]] const uint3 global_id)
{
  [[resource_table]] SurfelData &surfels = srt.surfels_data;

  int surfel_index = int(global_id.x);
  if (surfel_index >= int(surfels.capture_info_buf.surfel_len)) {
    return;
  }

  const ViewMatrices view = views.get(0);
  Surfel surfel = surfels.surfel_buf[surfel_index];

  float3 sky_L = view.world_incident_vector(surfel.position);

  if (surfel.next > -1) {
    Surfel surfel_next = surfels.surfel_buf[surfel.next];
    srt.radiance_transfer_surfel(surfel, surfel_next);
  }
  else {
    srt.radiance_transfer_world(surfel, -sky_L);
  }

  if (surfel.prev > -1) {
    Surfel surfel_prev = surfels.surfel_buf[surfel.prev];
    srt.radiance_transfer_surfel(surfel, surfel_prev);
  }
  else {
    srt.radiance_transfer_world(surfel, sky_L);
  }

  surfels.surfel_buf[surfel_index] = surfel;
}

}  // namespace eevee::surfel

PipelineCompute eevee_surfel_ray(eevee::surfel::ray_main);
