/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#include "obs-precomp.h"  // Precompiled headers
//
#include <mrpt/containers/stl_containers_utils.h>
#include <mrpt/core/round.h>
#include <mrpt/obs/CObservationVelodyneScan.h>
#include <mrpt/poses/CPose3DInterpolator.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/serialization/stl_serialization.h>

#include <iostream>

using namespace std;
using namespace mrpt::obs;

// shortcut:
using Velo = mrpt::obs::CObservationVelodyneScan;

// This must be added to any CSerializable class implementation file.
IMPLEMENTS_SERIALIZABLE(CObservationVelodyneScan, CObservation, mrpt::obs)

static CSinCosLookUpTableFor2DScans velodyne_sincos_tables;

const float VLP16_BLOCK_TDURATION = 110.592f;  // [us]
const float VLP16_DSR_TOFFSET = 2.304f;        // [us]
const float VLP16_FIRING_TOFFSET = 55.296f;    // [us]

const double HDR32_DSR_TOFFSET = 1.152;     // [us]
const double HDR32_FIRING_TOFFSET = 46.08;  // [us]

mrpt::system::TTimeStamp Velo::getOriginalReceivedTimeStamp() const
{
  return originalReceivedTimestamp;
}

uint8_t Velo::serializeGetVersion() const { return 2; }
void Velo::serializeTo(mrpt::serialization::CArchive& out) const
{
  out << timestamp << sensorLabel;
  out << minRange << maxRange << sensorPose;
  out.WriteAs<uint32_t>(scan_packets.size());
  if (!scan_packets.empty())
    out.WriteBuffer(&scan_packets[0], sizeof(scan_packets[0]) * scan_packets.size());
  out.WriteAs<uint32_t>(calibration.laser_corrections.size());
  if (!calibration.laser_corrections.empty())
    out.WriteBuffer(
        &calibration.laser_corrections[0],
        sizeof(calibration.laser_corrections[0]) * calibration.laser_corrections.size());
  out << point_cloud.x << point_cloud.y << point_cloud.z << point_cloud.intensity;
  out << has_satellite_timestamp;  // v1
  // v2:
  out << point_cloud.timestamp << point_cloud.azimuth << point_cloud.laser_id
      << point_cloud.pointsForLaserID;
}

void Velo::serializeFrom(mrpt::serialization::CArchive& in, uint8_t version)
{
  switch (version)
  {
    case 0:
    case 1:
    case 2:
    {
      in >> timestamp >> sensorLabel;

      in >> minRange >> maxRange >> sensorPose;
      {
        uint32_t N;
        in >> N;
        scan_packets.resize(N);
        if (N) in.ReadBuffer(&scan_packets[0], sizeof(scan_packets[0]) * N);
      }
      {
        uint32_t N;
        in >> N;
        calibration.laser_corrections.resize(N);
        if (N)
          in.ReadBuffer(
              &calibration.laser_corrections[0], sizeof(calibration.laser_corrections[0]) * N);
      }
      point_cloud.clear();
      in >> point_cloud.x >> point_cloud.y >> point_cloud.z >> point_cloud.intensity;
      if (version >= 1)
        in >> has_satellite_timestamp;
      else
        has_satellite_timestamp = (this->timestamp != this->originalReceivedTimestamp);
      if (version >= 2)
        in >> point_cloud.timestamp >> point_cloud.azimuth >> point_cloud.laser_id >>
            point_cloud.pointsForLaserID;
    }
    break;
    default:
      MRPT_THROW_UNKNOWN_SERIALIZATION_VERSION(version);
  };

  // m_cachedMap.clear();
}

void Velo::getDescriptionAsText(std::ostream& o) const
{
  CObservation::getDescriptionAsText(o);
  o << "Homogeneous matrix for the sensor 3D pose, relative to robot base:\n";
  o << sensorPose.getHomogeneousMatrixVal<mrpt::math::CMatrixDouble44>() << "\n"
    << sensorPose << "\n";
  o << format("Sensor min/max range: %.02f / %.02f m\n", minRange, maxRange);
  o << "Raw packet count: " << scan_packets.size() << "\n";
}

/** [us] */
namespace
{
double HDL32AdjustTimeStamp(int firingblock, int dsr)
{
  return (firingblock * HDR32_FIRING_TOFFSET) + (dsr * HDR32_DSR_TOFFSET);
}
/** [us] */
double VLP16AdjustTimeStamp(int firingblock, int dsr, int firingwithinblock)
{
  return (firingblock * VLP16_BLOCK_TDURATION) + (dsr * VLP16_DSR_TOFFSET) +
         (firingwithinblock * VLP16_FIRING_TOFFSET);
}

void velodyne_scan_to_pointcloud(
    const Velo& scan,
    const Velo::TGeneratePointCloudParameters& params,
    Velo::PointCloudStorageWrapper& out_pc)
{
  // Initially based on code from ROS velodyne & from
  // vtkVelodyneHDLReader::vtkInternal::ProcessHDLPacket().
  using mrpt::round;

  // Access to sin/cos table:
  mrpt::obs::T2DScanProperties scan_props;
  scan_props.aperture = 2 * M_PI;
  scan_props.nRays = Velo::ROTATION_MAX_UNITS;
  scan_props.rightToLeft = true;
  // The LUT contains sin/cos values for angles in this order: [180deg ... 0
  // deg ... -180 deg]
  const CSinCosLookUpTableFor2DScans::TSinCosValues& lut_sincos =
      velodyne_sincos_tables.getSinCosForScan(scan_props);

  const int minAzimuth_int = round(params.minAzimuth_deg * 100);
  const int maxAzimuth_int = round(params.maxAzimuth_deg * 100);
  const float realMinDist = std::max(mrpt::d2f(scan.minRange), params.minDistance);
  const float realMaxDist = std::min(params.maxDistance, mrpt::d2f(scan.maxRange));
  const auto isolatedPointsFilterDistance_units =
      mrpt::round(params.isolatedPointsFilterDistance / Velo::DISTANCE_RESOLUTION);

  // This is: 16,32,64 depending on the LIDAR model
  const size_t num_lasers = scan.calibration.laser_corrections.size();

  out_pc.resizeLaserCount(num_lasers);
  out_pc.reserve(Velo::SCANS_PER_BLOCK * scan.scan_packets.size() * Velo::BLOCKS_PER_PACKET + 16);

  for (size_t iPkt = 0; iPkt < scan.scan_packets.size(); iPkt++)
  {
    const Velo::TVelodyneRawPacket* raw = &scan.scan_packets[iPkt];

    mrpt::system::TTimeStamp pkt_tim;  // Find out timestamp of this pkt
    {
      const uint32_t us_pkt0 = scan.scan_packets[0].gps_timestamp();
      const uint32_t us_pkt_this = raw->gps_timestamp();
      // Handle the case of time counter reset by new hour 00:00:00
      const uint32_t us_ellapsed = (us_pkt_this >= us_pkt0)
                                       ? (us_pkt_this - us_pkt0)
                                       : (1000000UL * 3600UL + us_pkt_this - us_pkt0);
      pkt_tim = mrpt::system::timestampAdd(scan.timestamp, us_ellapsed * 1e-6);
    }

    // Take the median rotational speed as a good value for interpolating
    // the missing azimuths:
    int median_azimuth_diff;
    {
      // In dual return, the azimuth rate is actually twice this
      // estimation:
      const unsigned int nBlocksPerAzimuth = (raw->laser_return_mode == Velo::RETMODE_DUAL) ? 2 : 1;
      const size_t nDiffs = Velo::BLOCKS_PER_PACKET - nBlocksPerAzimuth;
      std::vector<int> diffs(nDiffs);
      for (size_t i = 0; i < nDiffs; ++i)
      {
        int localDiff = (Velo::ROTATION_MAX_UNITS + raw->blocks[i + nBlocksPerAzimuth].rotation() -
                         raw->blocks[i].rotation()) %
                        Velo::ROTATION_MAX_UNITS;
        diffs[i] = localDiff;
      }
      std::nth_element(
          diffs.begin(), diffs.begin() + Velo::BLOCKS_PER_PACKET / 2,
          diffs.end());  // Calc median
      median_azimuth_diff = diffs[Velo::BLOCKS_PER_PACKET / 2];
    }

    // Firings per packet
    for (int block = 0; block < Velo::BLOCKS_PER_PACKET; block++)
    {
      // ignore packets with mangled or otherwise different contents
      if ((num_lasers != 64 && Velo::UPPER_BANK != raw->blocks[block].header()) ||
          (raw->blocks[block].header() != Velo::UPPER_BANK &&
           raw->blocks[block].header() != Velo::LOWER_BANK))
      {
        cerr << "[Velo] skipping invalid packet: block " << block << " header value is "
             << raw->blocks[block].header();
        continue;
      }

      const int dsr_offset = (raw->blocks[block].header() == Velo::LOWER_BANK) ? 32 : 0;
      const auto azimuth_raw_f = mrpt::d2f(raw->blocks[block].rotation());
      const bool block_is_dual_2nd_ranges =
          (raw->laser_return_mode == Velo::RETMODE_DUAL && ((block & 0x01) != 0));
      const bool block_is_dual_last_ranges =
          (raw->laser_return_mode == Velo::RETMODE_DUAL && ((block & 0x01) == 0));

      for (int dsr = 0, k = 0; dsr < Velo::SCANS_PER_FIRING; dsr++, k++)
      {
        if (!raw->blocks[block].laser_returns[k].distance())  // Invalid return?
          continue;

        const auto rawLaserId = static_cast<uint8_t>(dsr + dsr_offset);
        uint8_t laserId = rawLaserId;

        // Detect VLP-16 data and adjust laser id if necessary
        bool firingWithinBlock = false;
        if (num_lasers == 16)
        {
          if (laserId >= 16)
          {
            laserId -= 16;
            firingWithinBlock = true;
          }
        }

        ASSERT_LT_(laserId, num_lasers);
        const mrpt::obs::VelodyneCalibration::PerLaserCalib& calib =
            scan.calibration.laser_corrections[laserId];

        // In dual return, if the distance is equal in both ranges,
        // ignore one of them:
        if (block_is_dual_2nd_ranges)
        {
          if (raw->blocks[block].laser_returns[k].distance() ==
              raw->blocks[block - 1].laser_returns[k].distance())
            continue;  // duplicated point
          if (!params.dualKeepStrongest) continue;
        }
        if (block_is_dual_last_ranges && !params.dualKeepLast) continue;

        // Return distance:
        const float distance = mrpt::d2f(
            raw->blocks[block].laser_returns[k].distance() * Velo::DISTANCE_RESOLUTION +
            calib.distanceCorrection);
        if (distance < realMinDist || distance > realMaxDist) continue;

        // Isolated points filtering:
        if (params.filterOutIsolatedPoints)
        {
          bool pass_filter = true;
          const int16_t dist_this = raw->blocks[block].laser_returns[k].distance();
          if (k > 0)
          {
            const int16_t dist_prev = raw->blocks[block].laser_returns[k - 1].distance();
            if (!dist_prev || std::abs(dist_this - dist_prev) > isolatedPointsFilterDistance_units)
              pass_filter = false;
          }
          if (k < (Velo::SCANS_PER_FIRING - 1))
          {
            const int16_t dist_next = raw->blocks[block].laser_returns[k + 1].distance();
            if (!dist_next || std::abs(dist_this - dist_next) > isolatedPointsFilterDistance_units)
              pass_filter = false;
          }
          if (!pass_filter) continue;  // Filter out this point
        }

        // Azimuth correction: correct for the laser rotation as a
        // function of timing during the firings
        double timestampadjustment = 0.0;  // [us] since beginning of scan
        double blockdsr0 = 0.0;
        double nextblockdsr0 = 1.0;
        switch (num_lasers)
        {
          // VLP-16
          case 16:
          {
            if (raw->laser_return_mode == Velo::RETMODE_DUAL)
            {
              timestampadjustment = VLP16AdjustTimeStamp(block / 2, laserId, firingWithinBlock);
              nextblockdsr0 = VLP16AdjustTimeStamp(block / 2 + 1, 0, 0);
              blockdsr0 = VLP16AdjustTimeStamp(block / 2, 0, 0);
            }
            else
            {
              timestampadjustment = VLP16AdjustTimeStamp(block, laserId, firingWithinBlock);
              nextblockdsr0 = VLP16AdjustTimeStamp(block + 1, 0, 0);
              blockdsr0 = VLP16AdjustTimeStamp(block, 0, 0);
            }
          }
          break;
          // HDL-32:
          case 32:
            timestampadjustment = HDL32AdjustTimeStamp(block, dsr);
            nextblockdsr0 = HDL32AdjustTimeStamp(block + 1, 0);
            blockdsr0 = HDL32AdjustTimeStamp(block, 0);
            break;
          case 64:
            break;
          default:
          {
            THROW_EXCEPTION("Error: unhandled LIDAR model!");
          }
        };

        const int azimuthadjustment = mrpt::round(
            median_azimuth_diff *
            ((timestampadjustment - blockdsr0) / (nextblockdsr0 - blockdsr0)));

        const float azimuth_corrected_f = azimuth_raw_f + azimuthadjustment;
        const int azimuth_corrected = ((int)round(azimuth_corrected_f)) % Velo::ROTATION_MAX_UNITS;

        // Filter by azimuth:
        if (!((minAzimuth_int < maxAzimuth_int && azimuth_corrected >= minAzimuth_int &&
               azimuth_corrected <= maxAzimuth_int) ||
              (minAzimuth_int > maxAzimuth_int &&
               (azimuth_corrected <= maxAzimuth_int || azimuth_corrected >= minAzimuth_int))))
          continue;

        // Vertical axis mis-alignment calibration:
        const float cos_vert_angle = mrpt::d2f(calib.cosVertCorrection);
        const float sin_vert_angle = mrpt::d2f(calib.sinVertCorrection);
        const float horz_offset = mrpt::d2f(calib.horizontalOffsetCorrection);
        const float vert_offset = mrpt::d2f(calib.verticalOffsetCorrection);

        float xy_distance = distance * cos_vert_angle;
        if (vert_offset != .0f) xy_distance += vert_offset * sin_vert_angle;

        const int azimuth_corrected_for_lut =
            (azimuth_corrected + (Velo::ROTATION_MAX_UNITS / 2)) % Velo::ROTATION_MAX_UNITS;
        const float cos_azimuth = lut_sincos.ccos[azimuth_corrected_for_lut];
        const float sin_azimuth = lut_sincos.csin[azimuth_corrected_for_lut];

        // Compute raw position
        const mrpt::math::TPoint3Df pt(
            xy_distance * cos_azimuth + horz_offset * sin_azimuth,     // MRPT +X = Velodyne +Y
            -(xy_distance * sin_azimuth - horz_offset * cos_azimuth),  // MRPT +Y = Velodyne -X
            distance * sin_vert_angle + vert_offset);

        bool add_point = true;
        if (params.filterByROI &&
            (pt.x > params.ROI_x_max || pt.x < params.ROI_x_min || pt.y > params.ROI_y_max ||
             pt.y < params.ROI_y_min || pt.z > params.ROI_z_max || pt.z < params.ROI_z_min))
          add_point = false;

        if (params.filterBynROI &&
            (pt.x <= params.nROI_x_max && pt.x >= params.nROI_x_min && pt.y <= params.nROI_y_max &&
             pt.y >= params.nROI_y_min && pt.z <= params.nROI_z_max && pt.z >= params.nROI_z_min))
          add_point = false;

        if (!add_point) continue;

        // Insert point:
        out_pc.add_point(
            pt.x, pt.y, pt.z, raw->blocks[block].laser_returns[k].intensity(), pkt_tim,
            azimuth_corrected_f, laserId);

      }  // end for k,dsr=[0,31]
    }    // end for each block [0,11]
  }      // end for each data packet
}
}  // namespace

void Velo::generatePointCloud(
    PointCloudStorageWrapper& dest, const TGeneratePointCloudParameters& params)
{
  velodyne_scan_to_pointcloud(*this, params, dest);
}

void Velo::generatePointCloud(const TGeneratePointCloudParameters& params)
{
  struct PointCloudStorageWrapper_Inner : public PointCloudStorageWrapper
  {
    Velo& me_;
    const TGeneratePointCloudParameters& params_;
    PointCloudStorageWrapper_Inner(Velo& me, const TGeneratePointCloudParameters& p) :
        me_(me), params_(p)
    {
      // Reset point cloud:
      me_.point_cloud.clear();
    }

    void resizeLaserCount(std::size_t n) override { me_.point_cloud.pointsForLaserID.resize(n); }

    void reserve(std::size_t n) override
    {
      me_.point_cloud.reserve(n);
      if (!me_.point_cloud.pointsForLaserID.empty())
      {
        const std::size_t n_per_ring = 100 + n / (me_.point_cloud.pointsForLaserID.size());
        for (auto& v : me_.point_cloud.pointsForLaserID) v.reserve(n_per_ring);
      }
    }

    void add_point(
        float pt_x,
        float pt_y,
        float pt_z,
        uint8_t pt_intensity,
        const mrpt::system::TTimeStamp& tim,
        const float azimuth,
        uint16_t laser_id) override
    {
      const auto idx = me_.point_cloud.x.size();
      me_.point_cloud.x.push_back(pt_x);
      me_.point_cloud.y.push_back(pt_y);
      me_.point_cloud.z.push_back(pt_z);
      me_.point_cloud.intensity.push_back(pt_intensity);
      if (params_.generatePerPointTimestamp)
      {
        me_.point_cloud.timestamp.push_back(tim);
      }
      if (params_.generatePerPointAzimuth)
      {
        const int azimuth_corrected = round(azimuth) % Velo::ROTATION_MAX_UNITS;
        me_.point_cloud.azimuth.push_back(azimuth_corrected * ROTATION_RESOLUTION);
      }
      me_.point_cloud.laser_id.push_back(laser_id);
      if (params_.generatePointsForLaserID)
        me_.point_cloud.pointsForLaserID[laser_id].push_back(idx);
    }
  };

  PointCloudStorageWrapper_Inner my_pc_wrap(*this, params);
  point_cloud.clear();

  generatePointCloud(my_pc_wrap, params);
}

void Velo::generatePointCloudAlongSE3Trajectory(
    const mrpt::poses::CPose3DInterpolator& vehicle_path,
    std::vector<mrpt::math::TPointXYZIu8>& out_points,
    TGeneratePointCloudSE3Results& results_stats,
    const TGeneratePointCloudParameters& params)
{
  struct PointCloudStorageWrapper_SE3_Interp : public PointCloudStorageWrapper
  {
    Velo& me_;
    const mrpt::poses::CPose3DInterpolator& vehicle_path_;
    std::vector<mrpt::math::TPointXYZIu8>& out_points_;
    TGeneratePointCloudSE3Results& results_stats_;
    mrpt::system::TTimeStamp last_query_tim_;
    mrpt::poses::CPose3D last_query_;
    bool last_query_valid_;

    void reserve(std::size_t n) override { out_points_.reserve(out_points_.size() + n); }

    PointCloudStorageWrapper_SE3_Interp(
        Velo& me,
        const mrpt::poses::CPose3DInterpolator& vehicle_path,
        std::vector<mrpt::math::TPointXYZIu8>& out_points,
        TGeneratePointCloudSE3Results& results_stats) :
        me_(me),
        vehicle_path_(vehicle_path),
        out_points_(out_points),
        results_stats_(results_stats),
        last_query_tim_(INVALID_TIMESTAMP),
        last_query_valid_(false)
    {
    }
    void add_point(
        float pt_x,
        float pt_y,
        float pt_z,
        uint8_t pt_intensity,
        const mrpt::system::TTimeStamp& tim,
        [[maybe_unused]] const float azimuth,
        [[maybe_unused]] uint16_t laser_id) override
    {
      // Use a cache since it's expected that the same timestamp is
      // queried several times in a row:
      if (last_query_tim_ != tim)
      {
        last_query_tim_ = tim;
        vehicle_path_.interpolate(tim, last_query_, last_query_valid_);
      }

      if (last_query_valid_)
      {
        mrpt::poses::CPose3D global_sensor_pose(mrpt::poses::UNINITIALIZED_POSE);
        global_sensor_pose.composeFrom(last_query_, me_.sensorPose);
        double gx, gy, gz;
        global_sensor_pose.composePoint(pt_x, pt_y, pt_z, gx, gy, gz);
        out_points_.emplace_back(gx, gy, gz, pt_intensity);
        ++results_stats_.num_correctly_inserted_points;
      }
      ++results_stats_.num_points;
    }
  };

  PointCloudStorageWrapper_SE3_Interp my_pc_wrap(*this, vehicle_path, out_points, results_stats);
  velodyne_scan_to_pointcloud(*this, params, my_pc_wrap);
}

void Velo::TPointCloud::clear()
{
  x.clear();
  y.clear();
  z.clear();
  intensity.clear();
  timestamp.clear();
  azimuth.clear();
  laser_id.clear();
  pointsForLaserID.clear();
}
void Velo::TPointCloud::clear_deep()
{
  using mrpt::containers::deep_clear;

  deep_clear(x);
  deep_clear(y);
  deep_clear(z);
  deep_clear(intensity);
  deep_clear(timestamp);
  deep_clear(azimuth);
  deep_clear(laser_id);
  deep_clear(pointsForLaserID);
}

void Velo::TPointCloud::reserve(std::size_t n)
{
  x.reserve(n);
  y.reserve(n);
  z.reserve(n);
  intensity.reserve(n);
  timestamp.reserve(n);
  azimuth.reserve(n);
  laser_id.reserve(n);
  pointsForLaserID.reserve(64);  // ring number
}

void Velo::unload() const { const_cast<Velo&>(*this).point_cloud.clear_deep(); }

// Default ctor. Do NOT move to the .h, that causes build errors.
Velo::TGeneratePointCloudParameters::TGeneratePointCloudParameters() = default;
