// Copyright 2009 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once

// ospray stuff
#include "common/Clipping.ih"
#include "common/DGEnum.h"
#include "common/DifferentialGeometry.ih"
#include "common/FeatureFlags.ih"
#include "common/Instance.ih"
#include "common/Ray.ih"
#include "common/RayCone.ih"
#include "common/RayQueryContext.ih"
#include "common/VolumeIntervals.ih"
#include "geometry/GeometryDispatch.ih"
#include "volume/Volume.ih"
// c++ shared
#include "WorldShared.h"

OSPRAY_BEGIN_ISPC_NAMESPACE

#ifdef OSPRAY_TARGET_SYCL
using namespace ospray;
#endif

struct Renderer;

inline void traceGeometryRay(const World *uniform world,
    varying Ray &ray,
    const uniform FeatureFlagsHandler &ffh)
{
  // Skip if no geometries scene
  if (!world->embreeSceneHandleGeometries)
    return;

  const uniform FeatureFlags ff = getFeatureFlags(ffh);

  uniform RayQueryContextDefault context;
  initRayQueryContextDefault(&context, ffh);

  uniform RTCIntersectArguments intersectArgs;
  rtcInitIntersectArguments(&intersectArgs);
  intersectArgs.context = &context.ectx;
  intersectArgs.intersect = (RTCIntersectFunctionN)Geometry_dispatch_intersect;
  intersectArgs.feature_mask = (uniform RTCFeatureFlags)(
      (ff.geometry & ~FFG_OSPRAY_MASK) | RTC_FEATURE_FLAG_INSTANCE);

  rtcIntersectV(world->embreeSceneHandleGeometries,
      (varying RTCRayHit * uniform) & ray,
      &intersectArgs);
}

#ifdef OSPRAY_ENABLE_VOLUMES
inline void traceVolumeRay(
    const World *uniform world, varying Ray &ray, VolumeInterval &interval)
{
  initVolumeInterval(interval);
  if (!world->embreeSceneHandleVolumes)
    return;

  VolumeIntervals intervals;
#ifdef OSPRAY_TARGET_SYCL
  // We don't have access to TLS/dynamic memory on the kernel, so we only
  // support one interval on the GPU (i.e., no overlapping volumes)
  intervals.numVolumeIntervals = 0;
  intervals.numAllocated = 1;
  intervals.intervals = &interval;
#else
  allocVolumeIntervals(intervals);
#endif

  uniform RayQueryContextVolume context;
  InitRayQueryContextVolume(&context, &intervals);

  uniform RTCIntersectArguments intersectArgs;
  rtcInitIntersectArguments(&intersectArgs);
  intersectArgs.context = &context.ectx;
  intersectArgs.intersect = (RTCIntersectFunctionN)Volume_intersect_kernel;
  intersectArgs.feature_mask = RTC_FEATURE_FLAG_INSTANCE
      | RTC_FEATURE_FLAG_USER_GEOMETRY_CALLBACK_IN_ARGUMENTS;

  rtcIntersectV(world->embreeSceneHandleVolumes,
      (varying RTCRayHit * uniform) & ray,
      &intersectArgs);

  // In the SYCL case we write to the interval parameter directly. In the
  // non-sycl case we need to get the first interval to return to the caller
  if (intervals.numVolumeIntervals > 0 && hasInterval(intervals.intervals[0])) {
    interval.instance = *(world->instances + intervals.intervals[0].instID);
    interval.volumetricModel =
        interval.instance->group
            ->volumetricModels[intervals.intervals[0].geomID];
    interval.interval = intervals.intervals[0].interval;
  }
#ifndef OSPRAY_TARGET_SYCL
  freeVolumeIntervals(intervals);
#endif
}

inline void traceVolumeRay(
    const World *uniform world, varying Ray &ray, VolumeIntervals &intervals)
{
  if (!world->embreeSceneHandleVolumes)
    return;

  uniform RayQueryContextVolume context;
  InitRayQueryContextVolume(&context, &intervals);

  uniform RTCIntersectArguments intersectArgs;
  rtcInitIntersectArguments(&intersectArgs);
  intersectArgs.context = &context.ectx;
  intersectArgs.intersect = (RTCIntersectFunctionN)Volume_intersect_kernel;
  intersectArgs.feature_mask = RTC_FEATURE_FLAG_INSTANCE
      | RTC_FEATURE_FLAG_USER_GEOMETRY_CALLBACK_IN_ARGUMENTS;

  rtcIntersectV(world->embreeSceneHandleVolumes,
      (varying RTCRayHit * uniform) & ray,
      &intersectArgs);

  if (intervals.numVolumeIntervals > 0) {
    for (uniform uint32 i = 0; i < intervals.numVolumeIntervals; ++i) {
      Instance *instance = *(world->instances + intervals.intervals[i].instID);
      intervals.intervals[i].instance = instance;
      intervals.intervals[i].volumetricModel =
          instance->group->volumetricModels[intervals.intervals[i].geomID];
    }
  }
}
#endif

// Intersection context structure used for clipping geometries
struct RayQueryContextClipping
{
  uniform RayQueryContextDefault super;
  const World *uniform world;
  varying int32 corrClippingDepth;
  varying uint32 hitsCount;
  varying ClippingHit hits[CLIPPING_HITS_MAX_COUNT];
};

unmasked void clippingIntersectionFilterV(
    const RTCFilterFunctionNArguments *uniform args);

inline void traceClippingRay(const World *uniform world,
    varying Ray &ray,
    varying RayIntervals &rayIntervals,
    const uniform FeatureFlagsHandler &ffh)
{
  // Clipping disabled on GPU for now
#ifdef OSPRAY_TARGET_SYCL
  (void)world;
  (void)ffh;
  rayIntervals.intervals[0] = make_box1f(ray.t0, ray.t);
  rayIntervals.count = 1;
  return;
#else
  // A scene with clipping geometries has to exist
  if (!world->embreeSceneHandleClippers) {
    rayIntervals.intervals[0] = make_box1f(ray.t0, ray.t);
    rayIntervals.count = 1;
    return;
  }

  // Create and initialize intersection context
  RayQueryContextClipping context;
  rtcInitRayQueryContext(&context.super.ectx);
  context.super.type = RQCT_CLIPPING;
  context.super.ffh = &ffh;
  context.world = world;
  context.corrClippingDepth = 0;
  context.hitsCount = 0;

  // Create and initialize intersection arguments
  uniform RTCIntersectArguments intersectArgs;
  rtcInitIntersectArguments(&intersectArgs);
  intersectArgs.context = &context.super.ectx;
  intersectArgs.intersect = (RTCIntersectFunctionN)Geometry_dispatch_intersect;
  intersectArgs.filter = clippingIntersectionFilterV;
  intersectArgs.flags = RTC_RAY_QUERY_FLAG_INVOKE_ARGUMENT_FILTER;

  // Intersect all geometry along given ray,
  // we have to temporarily extend the ray to inf because
  // even distant intersections affect visibility of close objects
  float origT = ray.t;
  ray.t = inf;
  rtcIntersectV(world->embreeSceneHandleClippers,
      (varying RTCRayHit * uniform) & ray,
      &intersectArgs);
  ray.t = origT;

  // Set initial clipping depth,
  // we have to apply correction because ray origin can be inside clipping area
  const uniform int32 voidClippingDepth = world->numInvertedClippers;
  int32 clippingDepth = voidClippingDepth - context.corrClippingDepth;

  // Variables used for intervals construction
  uint32 intervalId = 0;
  bool intervalCompleted = true;

  // Start ray interval if not in clipping area
  if (clippingDepth == 0) {
    rayIntervals.intervals[intervalId].lower = ray.t0;
    intervalCompleted = false;
  }

  // Iterate through collected hits and build ray intervals
  for (uint32 i = 0; i < context.hitsCount; i++) {
    // Do not build ray intervals that are further than ray.t
    float t = context.hits[i].t;
    float absT = abs(t);
    if (absT > ray.t)
      break;

    // Check if coming into or out of clipping area
    if (t < 0.0f) {
      clippingDepth--; // out of clipping area
    } else {
      clippingDepth++; // into clipping area
    }

    // Start interval if previous one is closed and we are not in clipping area
    if (clippingDepth == 0 && intervalCompleted) {
      rayIntervals.intervals[intervalId].lower = absT;
      intervalCompleted = false;
    }

    // End interval if current one is not closed and we enter clipping area
    if (clippingDepth > 0 && !intervalCompleted) {
      rayIntervals.intervals[intervalId].upper = absT;
      intervalCompleted = true;
      intervalId++;
    }
  }

  // Complete ray interval if started
  if (!intervalCompleted) {
    rayIntervals.intervals[intervalId].upper = ray.t;
    intervalId++;
  }

  // Save number of ray intervals
  rayIntervals.count = intervalId;
#endif
}

inline void traceGeometryRayIntervals(const World *uniform world,
    Ray &ray,
    RayIntervals &rayIntervals,
    const uniform FeatureFlagsHandler &ffh)
{
  // Save the ray
  const float t0 = ray.t0;
  const float t = ray.t;

  // Iterate through ray intervals
  for (uint32 i = 0; i < rayIntervals.count; i++) {
    // Set ray interval
    ray.t0 = rayIntervals.intervals[i].lower;
    ray.t = rayIntervals.intervals[i].upper;

    // Skip intervals outside of the ray range
    if (t < ray.t0 || ray.t < t0) {
      ray.t0 = t0;
      ray.t = t;
      continue;
    }

    // Clip interval to the ray range
    ray.t0 = max(ray.t0, t0);
    ray.t = min(ray.t, t);

    // Shoot the ray
    traceGeometryRay(world, ray, ffh);

    // Exit loop if geometry hit
    if (hadHit(ray)) {
      ray.t0 = t0;
      return;
    }
  }

  // Restore the ray
  ray.t0 = t0;
  ray.t = t;
}

inline void traceRay(const World *uniform world,
    varying Ray &ray,
    const uniform FeatureFlagsHandler &ffh)
{
#ifdef OSPRAY_TARGET_SYCL
  // Clipping disabled for now
  traceGeometryRay(world, ray, ffh);
#else
  // Fast path if no clipping geometry
  if (!world->embreeSceneHandleClippers) {
    traceGeometryRay(world, ray, ffh);
    return;
  }

  // Trace ray in clipping geometries scene, fill array with ray intervals
  varying RayIntervals rayIntervals;
  traceClippingRay(world, ray, rayIntervals, ffh);

  // Trace ray intervals
  traceGeometryRayIntervals(world, ray, rayIntervals, ffh);
#endif
}

inline bool isOccludedNoClipping(const World *uniform world,
    varying Ray &ray,
    const uniform FeatureFlagsHandler &ffh)
{
  // Skip if no geometries scene
  if (!world->embreeSceneHandleGeometries)
    return false;
  const uniform FeatureFlags ff = getFeatureFlags(ffh);

  uniform RayQueryContextDefault context;
  initRayQueryContextDefault(&context, ffh);

  uniform RTCOccludedArguments occludedArgs;
  rtcInitOccludedArguments(&occludedArgs);
  occludedArgs.context = &context.ectx;
  occludedArgs.occluded = (RTCOccludedFunctionN)Geometry_dispatch_occluded;
  occludedArgs.feature_mask = (uniform RTCFeatureFlags)(
      (ff.geometry & ~FFG_OSPRAY_MASK) | RTC_FEATURE_FLAG_INSTANCE);

  rtcOccludedV(world->embreeSceneHandleGeometries,
      (varying RTCRay * uniform) & ray,
      &occludedArgs);

  return ray.t < ray.t0;
}

inline bool areIntervalsOccluded(const World *uniform world,
    varying Ray &ray,
    varying RayIntervals &rayIntervals,
    const uniform FeatureFlagsHandler &ffh)
{
  // Iterate through ray intervals
  for (uint32 i = 0; i < rayIntervals.count; i++) {
    // Set ray interval
    ray.t0 = rayIntervals.intervals[i].lower;
    ray.t = rayIntervals.intervals[i].upper;

    // Check for occluders
    if (isOccludedNoClipping(world, ray, ffh))
      return true;
  }

  // No occluder found
  return false;
}

inline bool isOccluded(const World *uniform world,
    varying Ray &ray,
    const uniform FeatureFlagsHandler &ffh)
{
#ifdef OSPRAY_TARGET_SYCL
  // Clipping disabled for now
  return isOccludedNoClipping(world, ray, ffh);
#else
  // Fast path if no clipping geometry
  if (!world->embreeSceneHandleClippers) {
    return isOccludedNoClipping(world, ray, ffh);
  }

  // Allocate array for ray intervals
  varying RayIntervals rayIntervals;
  rayIntervals.count = 0;

  // Trace ray in clipping geometries scene, fill array with ray intervals
  traceClippingRay(world, ray, rayIntervals, ffh);

  // Is there any occluder within given ray intervals
  return areIntervalsOccluded(world, ray, rayIntervals, ffh);
#endif
}

/*! Perform post-intersect computations, i.e. fill the members of
    DifferentialGeometry. Should only get called for rays that actually hit
    that given world. Variables are calculated according to 'flags', a
    bit-combination of DG_PostIntersectFlags.
    The ray, dg.P, dg.Ng, and dg.Ns are in world-coordinates.
    Color defaults to vec4f(1.f) if queried but not present in geometry.
   */
inline void postIntersect(const World *uniform world,
    const Renderer *uniform renderer,
    varying DifferentialGeometry &dg,
    const varying Ray &ray,
    varying RayCone &rayCone,
    uniform int64 flags,
    const uniform FeatureFlagsHandler &ffh)
{
  dg.primID = ray.primID;
  dg.st = make_vec2f(ray.u, ray.v);

  dg.material = NULL;
  dg.renderer = renderer;

  if (flags & DG_COLOR)
    dg.color = make_vec4f(1.f);

  if (flags & DG_TANGENTS) {
    dg.dPds = make_vec3f(1.f, 0.f, 0.f);
    dg.dPdt = make_vec3f(0.f, 1.f, 0.f);
  }

  dg.P = ray.org + ray.t * ray.dir;
  dg.epsilon = 0.f; // per default no geometry-type specific epsilon

  foreach_unique (instID in ray.instID) {
    if (instID != RTC_INVALID_GEOMETRY_ID) {
      Instance *uniform instance = *(world->instances + instID);
      Instance_postIntersect(instance, renderer, dg, ray, flags, false, ffh);
    } else {
      dg.Ns = dg.Ng = ray.Ng;
    }
  }

  // merge geometry-type specific epsilon with general epsilon
  dg.epsilon = max(dg.epsilon, calcEpsilon(dg.P, ray.dir, ray.t));

// some useful combinations; enums unfortunately don't work :-(
#define DG_NG_FACEFORWARD (DG_NG | DG_FACEFORWARD)
#define DG_NS_FACEFORWARD (DG_NS | DG_FACEFORWARD)
#define DG_NG_NORMALIZE (DG_NG | DG_NORMALIZE)
#define DG_NS_NORMALIZE (DG_NS | DG_NORMALIZE)

  vec3f ffnng = normalize(dg.Ng);

  if ((flags & DG_NG_NORMALIZE) == DG_NG_NORMALIZE)
    dg.Ng = ffnng;

  if ((flags & DG_NS_NORMALIZE) == DG_NS_NORMALIZE)
    dg.Ns = normalize(dg.Ns);

  const bool flip = dot(ray.dir, dg.Ng) >= 0.f;
  if (flip)
    ffnng = neg(ffnng);

  if ((flags & DG_NG_FACEFORWARD) == DG_NG_FACEFORWARD && flip)
    dg.Ng = neg(dg.Ng);

  if ((flags & DG_NS_FACEFORWARD) == DG_NS_FACEFORWARD) {
    if (dot(dg.Ng, dg.Ns) < 0.f)
      dg.Ns = neg(dg.Ns);
  }

#undef DG_NG_FACEFORWARD
#undef DG_NS_FACEFORWARD
#undef DG_NG_NORMALIZE
#undef DG_NS_NORMALIZE

  dg.P = dg.P + dg.epsilon * ffnng;

  rayCone.width += rayCone.dwdt * ray.t; // propagate to hit
  const float texArea = length(cross(dg.dPds, dg.dPdt));
  const float world2tex = rsqrtf(texArea);
  const float projConeWidth = rayCone.width / -dot(ray.dir, dg.Ng);
  dg.pixelFootprint = projConeWidth * world2tex;
  dg.pixelFootprint *= renderer->mipBiasFactor;
}

OSPRAY_END_ISPC_NAMESPACE