/* SPDX-License-Identifier: Apache-2.0
 * Copyright 2011-2022 Blender Foundation */

#include "scene/light_tree.h"
#include "scene/mesh.h"
#include "scene/object.h"

#include "util/progress.h"

CCL_NAMESPACE_BEGIN

float OrientationBounds::calculate_measure() const
{
  if (this->is_empty()) {
    return 0.0f;
  }

  float theta_w = fminf(M_PI_F, theta_o + theta_e);
  float cos_theta_o = cosf(theta_o);
  float sin_theta_o = sinf(theta_o);

  return M_2PI_F * (1 - cos_theta_o) +
         M_PI_2_F * (2 * theta_w * sin_theta_o - cosf(theta_o - 2 * theta_w) -
                     2 * theta_o * sin_theta_o + cos_theta_o);
}

OrientationBounds merge(const OrientationBounds &cone_a, const OrientationBounds &cone_b)
{
  if (cone_a.is_empty()) {
    return cone_b;
  }
  if (cone_b.is_empty()) {
    return cone_a;
  }

  /* Set cone a to always have the greater theta_o. */
  const OrientationBounds *a = &cone_a;
  const OrientationBounds *b = &cone_b;
  if (cone_b.theta_o > cone_a.theta_o) {
    a = &cone_b;
    b = &cone_a;
  }

  float cos_a_b = dot(a->axis, b->axis);
  float theta_d = safe_acosf(cos_a_b);
  float theta_e = fmaxf(a->theta_e, b->theta_e);

  /* Return axis and theta_o of a if it already contains b. */
  /* This should also be called when b is empty. */
  if (a->theta_o + 5e-4f >= fminf(M_PI_F, theta_d + b->theta_o)) {
    return OrientationBounds({a->axis, a->theta_o, theta_e});
  }

  /* Compute new theta_o that contains both a and b. */
  float theta_o = (theta_d + a->theta_o + b->theta_o) * 0.5f;

  if (theta_o >= M_PI_F) {
    return OrientationBounds({a->axis, M_PI_F, theta_e});
  }

  /* Slerp between a and b. */
  float3 new_axis;
  if (cos_a_b < -0.9995f) {
    /* Opposite direction, any orthogonal vector is fine. */
    float3 unused;
    make_orthonormals(a->axis, &new_axis, &unused);
  }
  else {
    float theta_r = theta_o - a->theta_o;
    float3 ortho = safe_normalize(b->axis - a->axis * cos_a_b);
    new_axis = a->axis * cosf(theta_r) + ortho * sinf(theta_r);
  }

  return OrientationBounds({new_axis, theta_o, theta_e});
}

LightTreeEmitter::LightTreeEmitter(Scene *scene, int prim_id, int object_id)
    : prim_id(prim_id), object_id(object_id)
{
  if (is_triangle()) {
    float3 vertices[3];
    Object *object = scene->objects[object_id];
    Mesh *mesh = static_cast<Mesh *>(object->get_geometry());
    Mesh::Triangle triangle = mesh->get_triangle(prim_id);
    Shader *shader = static_cast<Shader *>(mesh->get_used_shaders()[mesh->get_shader()[prim_id]]);

    for (int i = 0; i < 3; i++) {
      vertices[i] = mesh->get_verts()[triangle.v[i]];
    }

    /* instanced mesh lights have not applied their transform at this point.
     * in this case, these points have to be transformed to get the proper
     * spatial bound. */
    if (!mesh->transform_applied) {
      const Transform &tfm = object->get_tfm();
      for (int i = 0; i < 3; i++) {
        vertices[i] = transform_point(&tfm, vertices[i]);
      }
    }

    /* TODO: need a better way to handle this when textures are used. */
    float area = triangle_area(vertices[0], vertices[1], vertices[2]);
    measure.energy = area * average(shader->emission_estimate);

    /* NOTE: the original implementation used the bounding box centroid, but triangle centroid
     * seems to work fine */
    centroid = (vertices[0] + vertices[1] + vertices[2]) / 3.0f;

    const bool is_front_only = (shader->emission_sampling == EMISSION_SAMPLING_FRONT);
    const bool is_back_only = (shader->emission_sampling == EMISSION_SAMPLING_BACK);
    if (is_front_only || is_back_only) {
      /* One-sided. */
      measure.bcone.axis = safe_normalize(
          cross(vertices[1] - vertices[0], vertices[2] - vertices[0]));
      if (is_back_only) {
        measure.bcone.axis = -measure.bcone.axis;
      }
      if (transform_negative_scale(object->get_tfm())) {
        measure.bcone.axis = -measure.bcone.axis;
      }
      measure.bcone.theta_o = 0;
    }
    else {
      /* Double sided: any vector in the plane. */
      measure.bcone.axis = safe_normalize(vertices[0] - vertices[1]);
      measure.bcone.theta_o = M_PI_2_F;
    }
    measure.bcone.theta_e = M_PI_2_F;

    for (int i = 0; i < 3; i++) {
      measure.bbox.grow(vertices[i]);
    }
  }
  else {
    Light *lamp = scene->lights[object_id];
    LightType type = lamp->get_light_type();
    const float size = lamp->get_size();
    float3 strength = lamp->get_strength();

    centroid = scene->lights[object_id]->get_co();
    measure.bcone.axis = normalize(lamp->get_dir());

    if (type == LIGHT_AREA) {
      measure.bcone.theta_o = 0;
      measure.bcone.theta_e = lamp->get_spread() * 0.5f;

      /* For an area light, sizeu and sizev determine the 2 dimensions of the area light,
       * while axisu and axisv determine the orientation of the 2 dimensions.
       * We want to add all 4 corners to our bounding box. */
      const float3 half_extentu = 0.5f * lamp->get_sizeu() * lamp->get_axisu() * size;
      const float3 half_extentv = 0.5f * lamp->get_sizev() * lamp->get_axisv() * size;
      measure.bbox.grow(centroid + half_extentu + half_extentv);
      measure.bbox.grow(centroid + half_extentu - half_extentv);
      measure.bbox.grow(centroid - half_extentu + half_extentv);
      measure.bbox.grow(centroid - half_extentu - half_extentv);

      strength *= 0.25f; /* eval_fac scaling in `area.h` */
    }
    else if (type == LIGHT_POINT) {
      measure.bcone.theta_o = M_PI_F;
      measure.bcone.theta_e = M_PI_2_F;

      /* Point and spot lights can emit light from any point within its radius. */
      const float3 radius = make_float3(size);
      measure.bbox.grow(centroid - radius);
      measure.bbox.grow(centroid + radius);

      strength *= 0.25f * M_1_PI_F; /* eval_fac scaling in `spot.h` and `point.h` */
    }
    else if (type == LIGHT_SPOT) {
      measure.bcone.theta_o = 0;

      const float unscaled_theta_e = lamp->get_spot_angle() * 0.5f;
      const float len_u = len(lamp->get_axisu());
      const float len_v = len(lamp->get_axisv());
      const float len_w = len(lamp->get_dir());

      measure.bcone.theta_e = fast_atanf(fast_tanf(unscaled_theta_e) * fmaxf(len_u, len_v) /
                                         len_w);

      /* Point and spot lights can emit light from any point within its radius. */
      const float3 radius = make_float3(size);
      measure.bbox.grow(centroid - radius);
      measure.bbox.grow(centroid + radius);

      strength *= 0.25f * M_1_PI_F; /* eval_fac scaling in `spot.h` and `point.h` */
    }
    else if (type == LIGHT_BACKGROUND) {
      /* Set an arbitrary direction for the background light. */
      measure.bcone.axis = make_float3(0.0f, 0.0f, 1.0f);
      /* TODO: this may depend on portal lights as well. */
      measure.bcone.theta_o = M_PI_F;
      measure.bcone.theta_e = 0;

      /* integrate over cosine-weighted hemisphere */
      strength *= lamp->get_average_radiance() * M_PI_F;
    }
    else if (type == LIGHT_DISTANT) {
      measure.bcone.theta_o = 0;
      measure.bcone.theta_e = 0.5f * lamp->get_angle();
    }

    if (lamp->get_shader()) {
      strength *= lamp->get_shader()->emission_estimate;
    }

    /* Use absolute value of energy so lights with negative strength are properly supported in the
     * light tree. */
    measure.energy = fabsf(average(strength));
  }
}

LightTree::LightTree(Scene *scene,
                     DeviceScene *dscene,
                     Progress &progress,
                     uint max_lights_in_leaf)
    : progress_(progress), max_lights_in_leaf_(max_lights_in_leaf)
{
  KernelIntegrator *kintegrator = &dscene->data.integrator;

  /* Add both lights and emissive triangles to this vector for light tree construction. */
  emitters_.reserve(kintegrator->num_distribution);
  distant_lights_.reserve(kintegrator->num_distant_lights);
  uint *object_offsets = dscene->object_lookup_offset.alloc(scene->objects.size());

  /* When we keep track of the light index, only contributing lights will be added to the device.
   * Therefore, we want to keep track of the light's index on the device.
   * However, we also need the light's index in the scene when we're constructing the tree. */
  int device_light_index = 0;
  int scene_light_index = 0;
  for (Light *light : scene->lights) {
    if (light->is_enabled) {
      if (light->light_type == LIGHT_BACKGROUND || light->light_type == LIGHT_DISTANT) {
        distant_lights_.emplace_back(scene, ~device_light_index, scene_light_index);
      }
      else {
        emitters_.emplace_back(scene, ~device_light_index, scene_light_index);
      }

      device_light_index++;
    }

    scene_light_index++;
  }

  /* Similarly, we also want to keep track of the index of triangles that are emissive. */
  int object_id = 0;
  for (Object *object : scene->objects) {
    if (progress_.get_cancel()) {
      return;
    }

    if (!object->usable_as_light()) {
      object_id++;
      continue;
    }

    object_offsets[object_id] = num_triangles;

    /* Count emissive triangles. */
    Mesh *mesh = static_cast<Mesh *>(object->get_geometry());
    size_t mesh_num_triangles = mesh->num_triangles();
    for (size_t i = 0; i < mesh_num_triangles; i++) {
      int shader_index = mesh->get_shader()[i];
      Shader *shader = (shader_index < mesh->get_used_shaders().size()) ?
                           static_cast<Shader *>(mesh->get_used_shaders()[shader_index]) :
                           scene->default_surface;

      if (shader->emission_sampling != EMISSION_SAMPLING_NONE) {
        emitters_.emplace_back(scene, i, object_id);
      }
    }
    num_triangles += mesh_num_triangles;
    object_id++;
  }

  /* Copy array to device. */
  dscene->object_lookup_offset.copy_to_device();
}

LightTreeNode *LightTree::build(Scene * /* scene */, DeviceScene * /* dscene */)
{
  if (emitters_.empty() && distant_lights_.empty()) {
    return nullptr;
  }

  /* At this stage `emitters_` only contains local lights, the distant lights will be merged
   * into `emitters_` when Light Tree building is finished. */
  const int num_local_lights = emitters_.size();
  const int num_distant_lights = distant_lights_.size();

  root_ = create_node(LightTreeMeasure::empty, 0);

  /* All local lights are grouped to the left child as an inner node. */
  recursive_build(left, root_.get(), 0, num_local_lights, emitters_.data(), 0, 1);
  task_pool.wait_work();

  /* All distant lights are grouped to the right child as a leaf node. */
  root_->children[right] = create_node(LightTreeMeasure::empty, 1);
  for (int i = 0; i < num_distant_lights; i++) {
    root_->children[right]->add(distant_lights_[i]);
  }
  root_->children[right]->make_leaf(num_local_lights, num_distant_lights);

  /* Append distant lights to the end of `light_prims` */
  std::move(distant_lights_.begin(), distant_lights_.end(), std::back_inserter(emitters_));

  return root_.get();
}

void LightTree::recursive_build(const Child child,
                                LightTreeNode *parent,
                                const int start,
                                const int end,
                                LightTreeEmitter *emitters,
                                const uint bit_trail,
                                const int depth)
{
  if (progress_.get_cancel()) {
    return;
  }

  parent->children[child] = create_node(LightTreeMeasure::empty, bit_trail);
  LightTreeNode *node = parent->children[child].get();

  /* Find the best place to split the emitters into 2 nodes.
   * If the best split cost is no better than making a leaf node, make a leaf instead. */
  int split_dim = -1, middle;
  if (should_split(emitters, start, middle, end, node->measure, split_dim)) {

    if (split_dim != -1) {
      /* Partition the emitters between start and end based on the centroids.  */
      std::nth_element(emitters + start,
                       emitters + middle,
                       emitters + end,
                       [split_dim](const LightTreeEmitter &l, const LightTreeEmitter &r) {
                         return l.centroid[split_dim] < r.centroid[split_dim];
                       });
    }

    /* Recursively build the left branch. */
    if (middle - start > MIN_EMITTERS_PER_THREAD) {
      task_pool.push(
          [=] { recursive_build(left, node, start, middle, emitters, bit_trail, depth + 1); });
    }
    else {
      recursive_build(left, node, start, middle, emitters, bit_trail, depth + 1);
    }

    /* Recursively build the right branch. */
    if (end - middle > MIN_EMITTERS_PER_THREAD) {
      task_pool.push([=] {
        recursive_build(right, node, middle, end, emitters, bit_trail | (1u << depth), depth + 1);
      });
    }
    else {
      recursive_build(right, node, middle, end, emitters, bit_trail | (1u << depth), depth + 1);
    }
  }
  else {
    node->make_leaf(start, end - start);
  }
}

bool LightTree::should_split(LightTreeEmitter *emitters,
                             const int start,
                             int &middle,
                             const int end,
                             LightTreeMeasure &measure,
                             int &split_dim)
{
  const int num_emitters = end - start;
  if (num_emitters < 2) {
    if (num_emitters) {
      /* Do not try to split if there is only one emitter. */
      measure = (emitters + start)->measure;
    }
    return false;
  }

  middle = (start + end) / 2;

  BoundBox centroid_bbox = BoundBox::empty;
  for (int i = start; i < end; i++) {
    centroid_bbox.grow((emitters + i)->centroid);
  }

  const float3 extent = centroid_bbox.size();
  const float max_extent = max4(extent.x, extent.y, extent.z, 0.0f);

  /* Check each dimension to find the minimum splitting cost. */
  float total_cost = 0.0f;
  float min_cost = FLT_MAX;
  for (int dim = 0; dim < 3; dim++) {
    /* If the centroid bounding box is 0 along a given dimension, skip it. */
    if (centroid_bbox.size()[dim] == 0.0f && dim != 0) {
      continue;
    }

    const float inv_extent = 1 / (centroid_bbox.size()[dim]);

    /* Fill in buckets with emitters. */
    std::array<LightTreeBucket, LightTreeBucket::num_buckets> buckets;
    for (int i = start; i < end; i++) {
      const LightTreeEmitter *emitter = emitters + i;

      /* Place emitter into the appropriate bucket, where the centroid box is split into equal
       * partitions. */
      int bucket_idx = LightTreeBucket::num_buckets *
                       (emitter->centroid[dim] - centroid_bbox.min[dim]) * inv_extent;
      bucket_idx = clamp(bucket_idx, 0, LightTreeBucket::num_buckets - 1);

      buckets[bucket_idx].add(*emitter);
    }

    /* Precompute the left bucket measure cumulatively. */
    std::array<LightTreeBucket, LightTreeBucket::num_buckets - 1> left_buckets;
    left_buckets.front() = buckets.front();
    for (int i = 1; i < LightTreeBucket::num_buckets - 1; i++) {
      left_buckets[i] = left_buckets[i - 1] + buckets[i];
    }

    if (dim == 0) {
      /* Calculate node measure by summing up the bucket measure. */
      measure = left_buckets.back().measure + buckets.back().measure;

      /* Degenerate case with co-located emitters. */
      if (is_zero(centroid_bbox.size())) {
        break;
      }

      total_cost = measure.calculate();
      if (total_cost == 0.0f) {
        break;
      }
    }

    /* Precompute the right bucket measure cumulatively. */
    std::array<LightTreeBucket, LightTreeBucket::num_buckets - 1> right_buckets;
    right_buckets.back() = buckets.back();
    for (int i = LightTreeBucket::num_buckets - 3; i >= 0; i--) {
      right_buckets[i] = right_buckets[i + 1] + buckets[i + 1];
    }

    /* Calculate the cost of splitting at each point between partitions. */
    const float regularization = max_extent * inv_extent;
    for (int split = 0; split < LightTreeBucket::num_buckets - 1; split++) {
      const float left_cost = left_buckets[split].measure.calculate();
      const float right_cost = right_buckets[split].measure.calculate();
      const float cost = regularization * (left_cost + right_cost);

      if (cost < total_cost && cost < min_cost) {
        min_cost = cost;
        split_dim = dim;
        middle = start + left_buckets[split].count;
      }
    }
  }
  return min_cost < total_cost || num_emitters > max_lights_in_leaf_;
}

__forceinline LightTreeMeasure operator+(const LightTreeMeasure &a, const LightTreeMeasure &b)
{
  LightTreeMeasure c(a);
  c.add(b);
  return c;
}

LightTreeBucket operator+(const LightTreeBucket &a, const LightTreeBucket &b)
{
  return LightTreeBucket(a.measure + b.measure, a.count + b.count);
}

CCL_NAMESPACE_END