// SPDX-FileCopyrightText: Copyright (c) Copyright 2021 NVIDIA Corporation
// SPDX-License-Identifier: BSD-3-Clause

// # RTC Kernel:
// ** Volume Edge Enhancement **

// # Summary:
// Compute the colormap alpha gradient in the current volume and darken samples with high gradient
// magnitude a a user-defined rate.

// Define the user-defined data structure
struct Edge_enhancement_params
{
  float sample_range; // [10, 10] 1.0 sample range
  float rh;           // 1.0f ray sampling difference
  int stp_num;        // 6 [GUI] additional samples along ray
  float min_alpha;    // 0.2f minimum alpha for sampling (improves performance)
};

using namespace nv::index::xac;
using namespace nv::index::xaclib;

class Volume_sample_program
{
  NV_IDX_VOLUME_SAMPLE_PROGRAM

public:
  const Edge_enhancement_params*
    m_edge_enhancement_params; // define variables to bind user-defined buffer to

public:
  NV_IDX_DEVICE_INLINE_MEMBER
  void initialize()
  {
    // Bind the contents of the buffer slot 0 to the variable
    m_edge_enhancement_params = state.bind_parameter_buffer<Edge_enhancement_params>(0);
  }

  NV_IDX_DEVICE_INLINE_MEMBER
  int execute(const Sample_info_self&   sample_info,
                    Sample_output&      sample_output)
  {
    const auto& volume = state.self;
    const float3& ps = sample_info.sample_position_object_space;
    const Colormap& colormap = volume.get_colormap();

    const auto  svol_sampler =
      volume.generate_sampler<float>(0u, sample_info.sample_context);


    // sample volume
    const float3 ray_dir = sample_info.ray_direction;
    const float3 p0 = ps - (ray_dir * (m_edge_enhancement_params->rh *
                                       m_edge_enhancement_params->sample_range));
    const float3 p2 = ps + (ray_dir * (m_edge_enhancement_params->rh *
                                       m_edge_enhancement_params->sample_range));

    const float vs_p0 = svol_sampler.fetch_sample(p0);
    const float vs_ps = svol_sampler.fetch_sample(ps);
    const float vs_p2 = svol_sampler.fetch_sample(p2);

    const float4 ps_color = colormap.lookup(vs_ps);
    const float ps_alpha = ps_color.w;
    const int min_hits = m_edge_enhancement_params->stp_num;
    const int num_steps = min_hits * 2;

    if (ps_alpha < m_edge_enhancement_params->min_alpha)
    {
      // sample below given alpha threshold
      // return NV_IDX_PROG_DISCARD_SAMPLE;
    }

    if ((num_steps > 0) && (ps_alpha > m_edge_enhancement_params->min_alpha))
    {
      // init result color
      int sum_over_b = 0;
      int sum_under_b = 0;
      int sum_over_a = 0;
      int sum_under_a = 0;

      // iterate steps
      for (int sc = 0; sc <= num_steps; sc++)
      {
        // get step size
        const float rt = float(sc) / float(num_steps);
        const float3 pi_b = (1.0f - rt) * p0 + rt * ps;
        const float3 pi_a = (1.0f - rt) * ps + rt * p2;

        const float vs_pi_b = svol_sampler.fetch_sample(pi_b);
        const float vs_pi_a = svol_sampler.fetch_sample(pi_a);

        const float4 pia_color = colormap.lookup(vs_pi_b);
        const float4 pib_color = colormap.lookup(vs_pi_a);

        // check by alpha
        if (pib_color.w > ps_alpha)
          sum_over_b += 1;
        if (pib_color.w < ps_alpha)
          sum_under_b += 1;
        if (pia_color.w > ps_alpha)
          sum_over_a += 1;
        if (pia_color.w < ps_alpha)
          sum_under_a += 1;
      }

      // check if all are larger
      // if (sum_over > 1 && sum_under == 0)
      //    sample_output.color = make_float4(ps_color.x, ps_color.y, ps_color.z, ps_color.w);

      // check numer of samples below / above reference
      if (((sum_under_b > min_hits) && (sum_under_a > min_hits)) ||
        ((sum_over_b > min_hits) && (sum_over_a > min_hits)))
      {
        sample_output.color = ps_color / float(sum_under_b + sum_under_a - 2 * min_hits + 1);
        sample_output.color.w = ps_color.w;

        return NV_IDX_PROG_OK;
      }
      else
      {
        sample_output.color = ps_color;

        // return NV_IDX_PROG_DISCARD_SAMPLE;
      }
    }
    else
    {
      sample_output.color = ps_color;
    }

    return NV_IDX_PROG_OK;
  }
}; // class Volume_sample_program
