mujoco

// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "render/filament/support/model_lights.h"

#include <algorithm>
#include <memory>
#include <numbers>
#include <string>
#include <utility>

#include <math/TVecHelpers.h>
#include <math/mat4.h>
#include <math/mathfwd.h>
#include <math/vec3.h>
#include <math/vec4.h>
#include <mujoco/mjrfilament.h>
#include <mujoco/mujoco.h>
#include "render/filament/mjrfilament_cpp.h"
#include "render/filament/support/filament_util.h"
#include "render/filament/support/model_objects.h"

namespace mujoco {

using filament::math::float3;
using filament::math::float4;
using filament::math::mat3;
using filament::math::mat4;

static UniquePtr<mjrfTexture> CreateFallbackIndirectLightTexture(
    mjrfContext* ctx) {
  const std::string filename = ResolveFilamentAssetPath("ibl.ktx");
  mjResource* resource =
      mju_openResource("", filename.c_str(), nullptr, nullptr, 0);
  if (!resource) {
    mju_error("Failed to open resource: %s", filename.c_str());
  }
  const void* bytes = nullptr;
  const int nbytes = mju_readResource(resource, &bytes);
  if (bytes == nullptr || nbytes <= 0) {
    mju_error("Failed to read resource: %s", filename.c_str());
  }

  mjrfTextureConfig config;
  mjrf_defaultTextureConfig(&config);
  config.width = 1;
  config.height = 1;
  config.sampler_type = mjTEXTURE_CUBE;
  config.format = mjPIXEL_FORMAT_KTX;
  config.color_space = mjCOLORSPACE_AUTO;

  auto texture = CreateTexture(ctx, config);

  mjrfTextureData payload;
  mjrf_defaultTextureData(&payload);
  payload.bytes = bytes;
  payload.num_bytes = nbytes;
  payload.release =
      +[](void* user_data) { mju_closeResource((mjResource*)user_data); };
  payload.user_data = resource;

  mjrf_setTextureData(texture.get(), &payload);
  return texture;
}

ModelLights::ModelLights(mjrfScene* scene, ModelObjects* model_objects)
    : scene_(scene), model_objects_(model_objects) {
  const mjModel* model = model_objects->GetModel();
  shadowsize_ = model->vis.quality.shadowsize;
  fallback_head_light_intensity_ =
      ReadElement(model, "filament.fallback.head_light_intensity",
                  fallback_head_light_intensity_);
  fallback_scene_light_intensity_ =
      ReadElement(model, "filament.fallback.scene_light_intensity",
                  fallback_scene_light_intensity_);
  fallback_environment_light_intensity_ =
      ReadElement(model, "filament.fallback.environment_light_intensity",
                  fallback_environment_light_intensity_);
  Prepare();
}

ModelLights::~ModelLights() {
  for (auto& iter : lights_) {
    mjrf_removeLightFromScene(scene_, iter.get());
  }
  lights_.clear();
  if (fallback_ibl_) {
    mjrf_removeLightFromScene(scene_, fallback_ibl_.get());
  }
  fallback_ibl_.reset();
  if (fallback_directional_) {
    mjrf_removeLightFromScene(scene_, fallback_directional_.get());
  }
  fallback_directional_.reset();
}

void ModelLights::Prepare() {
  mjrfContext* ctx = model_objects_->GetContext();
  const mjModel* model = model_objects_->GetModel();

  // Default to the shadow map resolution in mjVisual, like the classic
  // renderer. Filament clamps shadow map sizes to 2048.
  int default_shadow_map_size = std::min(model->vis.quality.shadowsize, 2048);
  default_shadow_map_size =
      ReadElement(model, "filament.shadows.map_size", default_shadow_map_size);

  bool has_image_based_light = false;
  bool has_directional_light = false;
  float total_light_intensity = 0.0f;
  for (int i = 0; i < model->nlight; ++i) {
    total_light_intensity += model->light_intensity[i];
    has_directional_light |= model->light_type[i] == mjLIGHT_DIRECTIONAL;

    if (model->light_type[i] == mjLIGHT_IMAGE) {
      mjrfLightParams params;
      mjrf_defaultLightParams(&params);
      params.type = mjLIGHT_IMAGE;
      params.texture = model_objects_->GetTexture(model->light_texid[i]);
      params.intensity = model->light_intensity[i];
      auto light_obj = CreateLight(ctx, params);
      mjrf_addLightToScene(scene_, light_obj.get());
      lights_.emplace_back(std::move(light_obj));
      has_image_based_light = true;
    } else {
      mjrfLightParams params;
      mjrf_defaultLightParams(&params);
      params.color[0] = model->light_diffuse[0];
      params.color[1] = model->light_diffuse[1];
      params.color[2] = model->light_diffuse[2];
      params.type = (mjtLightType)model->light_type[i];
      params.cast_shadows = model->light_castshadow[i];
      // light_bulbradius is the radius of the emitting surface in meters.
      // Filament's DPCF/PCSS internally scale shadowBulbRadius from meters
      // to light-space texels, so we pass it through in meters for all types.
      // VSM requires a single uniform blur width for the map; we approximate
      // this using the angular size of the bulb from the scene center.
      const float bulb_radius = model->light_bulbradius[i];
      const float map_size = default_shadow_map_size;
      const float3 to_center = ReadFloat3(model->stat.center) -
                               ReadFloat3(model->light_pos0, i);
      const float distance = std::max(length(to_center), 1e-6f);
      const float bulb_angle = bulb_radius / distance;
      params.bulb_radius = bulb_radius;
      switch (params.type) {
        case mjLIGHT_SPOT: {
          const float fov =
              2.0f * model->light_cutoff[i] * std::numbers::pi / 180.0f;
          params.vsm_blur_width = bulb_angle * map_size / fov;
          break;
        }
        case mjLIGHT_POINT:
          params.vsm_blur_width =
              bulb_angle * map_size / (0.5f * std::numbers::pi);
          break;
        case mjLIGHT_DIRECTIONAL: {
          const float coverage =
              2.0f * model->vis.map.shadowclip * model->stat.extent;
          params.vsm_blur_width = bulb_radius * map_size / coverage;
          break;
        }
        default:
          break;
      }
      params.vsm_blur_width = std::min(params.vsm_blur_width, 125.0f);
      params.range = model->light_range[i];
      params.intensity = model->light_intensity[i];
      params.shadow_map_size = default_shadow_map_size;
      if (params.type == mjLIGHT_SPOT) {
        params.spot_cone_angle = model->light_cutoff[i];
        params.spot_softness = model->light_softness[i];
      }

      auto light_obj = CreateLight(ctx, params);
      mjrf_addLightToScene(scene_, light_obj.get());
      lights_.emplace_back(std::move(light_obj));
    }
  }

  // Add a placeholder (black) headlight as our last light. Going forward, we'll
  // assume lights_.back() is always the headlight.
  {
    mjrfLightParams params;
    mjrf_defaultLightParams(&params);
    // We break with the spec here slightly and use a spot light for the head
    // light instead of a directional params. This is because filament only
    // supports a single directional light, and we'd rather allow a scene
    // light to be that directional params. It's also a bit odd for a
    // directional light to move with the camera.
    params.type = mjLIGHT_SPOT;
    params.cast_shadows = 0;
    params.intensity = 0.0f;
    params.spot_cone_angle = 90.0f;
    auto light_obj = CreateLight(ctx, params);
    mjrf_addLightToScene(scene_, light_obj.get());
    lights_.emplace_back(std::move(light_obj));
  }

  // Workaround for an upstream filament bug, present since 1.74.0
  // (https://github.com/google/filament/issues/10249): under the non-PCF
  // shadow types (VSM/DPCF/PCSS), a scene where a punctual (spot/point) light
  // casts shadows renders fully black unless a directional light is also
  // present. Mere presence suffices: zero intensity, shadows off. Keep
  // filament's directional slot occupied with an invisible light whenever the
  // model has no directional light.
  if (!has_directional_light) {
    mjrfLightParams params;
    mjrf_defaultLightParams(&params);
    params.type = mjLIGHT_DIRECTIONAL;
    params.cast_shadows = 0;
    params.intensity = 0.0f;
    fallback_directional_ = CreateLight(ctx, params);
    mjrf_addLightToScene(scene_, fallback_directional_.get());
  }

  if (!has_image_based_light && total_light_intensity > 0.0f) {
    // Create a black indirect light to ensure that the skybox is
    // oriented to respect mujoco's Z-up convention.
    mjrfLightParams params;
    mjrf_defaultLightParams(&params);
    params.type = mjLIGHT_IMAGE;
    params.intensity = 10.0f;
    fallback_ibl_ = CreateLight(ctx, params);
    mjrf_addLightToScene(scene_, fallback_ibl_.get());
  }

  // There are no "physical" lights in the scene which means we're likely
  // dealing with a "classic renderer" scene. In this case, let's add a
  // default environment light and set the light intensity ourselves.
  if (total_light_intensity == 0.0f) {
    // Create a fallback environment light.
    fallback_ibl_texture_ = CreateFallbackIndirectLightTexture(ctx);

    mjrfLightParams params;
    mjrf_defaultLightParams(&params);
    params.type = mjLIGHT_IMAGE;
    params.texture = fallback_ibl_texture_.get();
    params.intensity = fallback_environment_light_intensity_;
    fallback_ibl_ = CreateLight(ctx, params);
    mjrf_addLightToScene(scene_, fallback_ibl_.get());

    // Distribute the fallback scene light intensity among the lights.
    const float intensity = fallback_scene_light_intensity_ / lights_.size();
    for (auto& light : lights_) {
      if (light) {
        const bool is_headlight = (light == lights_.back());
        mjrf_setLightIntensity(light.get(), is_headlight
                                                ? fallback_head_light_intensity_
                                                : intensity);
      }
    }
  }

  mjrf_setSceneSkybox(scene_, model_objects_->GetSkyboxTexture());
}

void ModelLights::UpdateShadowMapSize() {
  const mjModel* model = model_objects_->GetModel();
  if (model->vis.quality.shadowsize != shadowsize_) {
    shadowsize_ = model->vis.quality.shadowsize;
    const int map_size = std::min(shadowsize_, 2048);
    for (auto& light : lights_) {
      mjrf_setLightShadowMapSize(light.get(), map_size);
    }
  }
}

void ModelLights::Update(const mjData* data) {
  UpdateShadowMapSize();
  if (data == nullptr) {
    return;
  }
  const mjModel* model = model_objects_->GetModel();
  for (int i = 0; i <= model->nlight; ++i) {
    // Light with index nlight is the headlight.
    mjrfLight* light = lights_[i].get();
    if (i == model->nlight) {
      const float3 color = ReadFloat3(model->vis.headlight.diffuse);
      mjrf_setLightColor(light, color.v);
      mjrf_setLightEnabled(light, model->vis.headlight.active);
    } else {
      const float3 pos = ReadFloat3(data->light_xpos, i);
      const float3 dir = ReadFloat3(data->light_xdir, i);
      mjrf_setLightTransform(light, pos.v, dir.v);

      const float3 color = ReadFloat3(model->light_diffuse, i);
      mjrf_setLightColor(light, color.v);

      mjrf_setLightEnabled(light, model->light_active[i]);
    }
  }
}

mjrfLight* ModelLights::GetLight(int index) {
  if (index < 0 || index >= lights_.size()) {
    return nullptr;
  }
  return lights_[index].get();
}
}  // namespace mujoco