__host__ function BVH
bvh_build()
{
  U64 max_bvh_nodes = 2 * MAX_NUM_ENTITIES - 1;
  BVH bvh = {0};
  bvh.nodes = (BVHNode *)_aligned_malloc(sizeof(BVHNode)*max_bvh_nodes, 64);
  bvh.max_num_nodes = max_bvh_nodes;
  bvh.used_nodes = 2; // Skip by two, TODO(anton): Comment this.
  bvh.minimum_entities_in_leaf = 2;
  U32 root_index = 0;
  BVHNode *root = &bvh.nodes[root_index];
  root->left_first = 0; 
  root->tri_count = MAX_NUM_ENTITIES;

  bvh_update_bounds(&bvh, 0);
  
  bvh_subdivide(&bvh, 0);

  return bvh;
}

__host__ function void
bvh_subdivide(BVH *bvh, U32 node_idx)
{
  BVHNode *node = &bvh->nodes[node_idx];
  if(node->tri_count <= bvh->minimum_entities_in_leaf) 
  {
    return;
  }

  // Split box
  Vec3F32 extent = sub_V3F32(node->aabb_max, node->aabb_min);
  U32 axis = 0;
  if(extent.y > extent.x) axis = 1;
  if(extent.z > extent.v[axis]) axis = 2;
  F32 split_pos = node->aabb_min.v[axis] + extent.v[axis] * 0.5f;

  // Sorting into left and right partitions
  U32 i = node->left_first;
  U32 j = node->tri_count + i - 1;
  while (i <= j)
  {
    U32 tri_idx = h_tri_indices[i];
    if(h_entities[tri_idx].center.v[axis] < split_pos)
    {
      i += 1;
    }
    else
    {
      h_tri_indices[i] = h_tri_indices[j];
      h_tri_indices[j] = tri_idx;
      j -= 1;
    }
  }

  U32 left_count = i - node->left_first;
  if(left_count == 0 || left_count == node->tri_count)
  {
    // One of the partitions are empty, don't subdivide further.
    return;
  }

  // Create child nodes and subdivide
  U32 left_child_index  = bvh->used_nodes++;
  U32 right_child_index = bvh->used_nodes++;

  bvh->nodes[left_child_index].left_first = node->left_first;
  bvh->nodes[left_child_index].tri_count = left_count;
  bvh->nodes[right_child_index].left_first = i;
  bvh->nodes[right_child_index].tri_count = node->tri_count - left_count;
  node->left_first = left_child_index;
  node->tri_count = 0;
  bvh_update_bounds(bvh, left_child_index);
  bvh_update_bounds(bvh, right_child_index);
  bvh_subdivide(bvh, left_child_index);
  bvh_subdivide(bvh, right_child_index);
}


__host__ function void 
bvh_update_bounds(BVH *bvh, U32 node_idx)
{
  BVHNode *node = &bvh->nodes[node_idx];

  node->aabb_min = vec3F32(F32_MAX, F32_MAX, F32_MAX);
  node->aabb_max = vec3F32(F32_MIN, F32_MIN, F32_MIN);
  
  U32 first_tri_idx = node->left_first;
  for(U32 i = 0; i < node->tri_count; i += 1)
  {
    U32 leaf_tri_idx = h_tri_indices[first_tri_idx + i];
    Entity *tri = &h_entities[leaf_tri_idx];
    node->aabb_min = h_min_V3F32(node->aabb_min, tri->vertex0);
    node->aabb_min = h_min_V3F32(node->aabb_min, tri->vertex1);
    node->aabb_min = h_min_V3F32(node->aabb_min, tri->vertex2);
    node->aabb_max = h_max_V3F32(node->aabb_max, tri->vertex0);
    node->aabb_max = h_max_V3F32(node->aabb_max, tri->vertex1);
    node->aabb_max = h_max_V3F32(node->aabb_max, tri->vertex2);
  } 
}

__host__ function void
bvh_host_intersect(BVH *bvh, RayF32 *ray, HitRecord *rec, U32 node_idx)
{
  BVHNode *node = &bvh->nodes[node_idx];
  U32 any_hit = 0;
  if(h_intersect_aabb(ray, node->aabb_min, node->aabb_max, rec->t))
  {
    if(node->tri_count > 0)
    {
      //LOG("Hit a leaf node %i with tri count %i \n", node_idx, node->tri_count);
      for(U32 i = 0; i < node->tri_count; i+=1)
      {
        U32 tri_index = h_tri_indices[node->left_first + i];
        Entity *tri = &h_entities[tri_index];
        hit_triangle_host(ray, rec, tri);
        if(rec->hit)
        { 
          any_hit = 1;
          //LOG("got hit in bvh_host_intersect loop \n");
        }
      }
    }
    else
    {
      bvh_host_intersect(bvh, ray, rec, node->left_first);
      bvh_host_intersect(bvh, ray, rec, node->left_first + 1);
    }
  }

  if(!rec->hit) 
  {
     rec->hit = any_hit;
  }
}