File angle_corrector_calibration_based.hpp

// Copyright 2024 TIER IV, Inc.
//
// 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.

#pragma once

#include "nebula_core_decoders/angles.hpp"
#include "nebula_hesai_common/hesai_common.hpp"
#include "nebula_hesai_decoders/decoders/angle_corrector.hpp"

#include <nebula_core_common/nebula_common.hpp>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <memory>

namespace nebula::drivers
{

template <size_t ChannelN, size_t AngleUnit>
class AngleCorrectorCalibrationBased
: public AngleCorrector<HesaiCalibrationConfiguration, ChannelN>
{
private:
  static constexpr size_t max_azimuth = 360 * AngleUnit;

  std::array<float, ChannelN> elevation_angle_rad_{};
  std::array<float, ChannelN> azimuth_offset_rad_{};
  std::array<float, max_azimuth> block_azimuth_rad_{};

  std::array<float, ChannelN> elevation_cos_{};
  std::array<float, ChannelN> elevation_sin_{};
  std::array<std::array<float, ChannelN>, max_azimuth> azimuth_cos_{};
  std::array<std::array<float, ChannelN>, max_azimuth> azimuth_sin_{};

  size_t min_correction_index_{};
  size_t max_correction_index_{};

  [[nodiscard]] int32_t to_exact_angle(double angle_deg) const
  {
    return std::round(angle_deg * AngleUnit);
  }

  [[nodiscard]] float to_radians(int32_t angle_exact) const
  {
    return deg2rad(angle_exact / static_cast<double>(AngleUnit));
  }

public:
  explicit AngleCorrectorCalibrationBased(
    const std::shared_ptr<const HesaiCalibrationConfiguration> & sensor_calibration)
  {
    if (sensor_calibration == nullptr) {
      throw std::runtime_error(
        "Cannot instantiate AngleCorrectorCalibrationBased without calibration data");
    }

    // ////////////////////////////////////////
    // Elevation lookup tables
    // ////////////////////////////////////////

    for (size_t channel_id = 0; channel_id < ChannelN; ++channel_id) {
      float elevation_angle_deg = sensor_calibration->elev_angle_map.at(channel_id);
      float azimuth_offset_deg = sensor_calibration->azimuth_offset_map.at(channel_id);

      elevation_angle_rad_[channel_id] = deg2rad(elevation_angle_deg);
      azimuth_offset_rad_[channel_id] = deg2rad(azimuth_offset_deg);

      elevation_cos_[channel_id] = cosf(elevation_angle_rad_[channel_id]);
      elevation_sin_[channel_id] = sinf(elevation_angle_rad_[channel_id]);
    }

    // ////////////////////////////////////////
    // Azimuth lookup tables
    // ////////////////////////////////////////

    for (size_t block_azimuth = 0; block_azimuth < max_azimuth; block_azimuth++) {
      block_azimuth_rad_[block_azimuth] = to_radians(block_azimuth);
      for (size_t channel_id = 0; channel_id < ChannelN; ++channel_id) {
        float spatial_azimuth_rad =
          block_azimuth_rad_[block_azimuth] + azimuth_offset_rad_[channel_id];
        azimuth_cos_[block_azimuth][channel_id] = cosf(spatial_azimuth_rad);
        azimuth_sin_[block_azimuth][channel_id] = sinf(spatial_azimuth_rad);
      }
    }

    const auto & az = azimuth_offset_rad_;
    min_correction_index_ = std::min_element(az.begin(), az.end()) - az.begin();
    max_correction_index_ = std::max_element(az.begin(), az.end()) - az.begin();
  }

  [[nodiscard]] CorrectedAngleData get_corrected_angle_data(
    uint32_t block_azimuth, uint32_t channel_id) const override
  {
    float azimuth_rad = block_azimuth_rad_[block_azimuth] + azimuth_offset_rad_[channel_id];
    azimuth_rad = normalize_angle(azimuth_rad, M_PIf * 2);

    float elevation_rad = elevation_angle_rad_[channel_id];

    return {
      azimuth_rad,
      elevation_rad,
      azimuth_sin_[block_azimuth][channel_id],
      azimuth_cos_[block_azimuth][channel_id],
      elevation_sin_[channel_id],
      elevation_cos_[channel_id]};
  }

  [[nodiscard]] CorrectedAzimuths<ChannelN, float> get_corrected_azimuths(
    uint32_t block_azimuth) const override
  {
    CorrectedAzimuths<ChannelN, float> corrected_azimuths;
    float block_azimuth_rad = block_azimuth_rad_[block_azimuth];

    for (size_t channel_id = 0; channel_id < ChannelN; ++channel_id) {
      float exact_azimuth = block_azimuth_rad + azimuth_offset_rad_[channel_id];
      exact_azimuth = normalize_angle(exact_azimuth, 2 * M_PIf);
      corrected_azimuths.azimuths[channel_id] = exact_azimuth;
    }

    // Use precomputed min/max indices based on correction terms (offsets).
    // min_correction_index = channel that lags behind (smallest offset)
    // max_correction_index = channel that races ahead (largest offset)
    // These are invariant across block_azimuth values.
    corrected_azimuths.min_correction_index = min_correction_index_;
    corrected_azimuths.max_correction_index = max_correction_index_;

    return corrected_azimuths;
  }
};

}  // namespace nebula::drivers