/*
This file is part of FAST-LIVO2: Fast, Direct LiDAR-Inertial-Visual Odometry.

Developer: Chunran Zheng <zhengcr@connect.hku.hk>

For commercial use, please contact me at <zhengcr@connect.hku.hk> or
Prof. Fu Zhang at <fuzhang@hku.hk>.

This file is subject to the terms and conditions outlined in the 'LICENSE' file,
which is included as part of this source code package.

Multi-camera vectorization (N-camera, joint ESIKF, cross-camera photometric
consistency) ported from Omni-LIVO — see
docs/OMNI_LIVO_DUAL_CAMERA_PORTING_PLAN.md Phase 4.
*/

#ifndef VIO_H_
#define VIO_H_

#include "voxel_map.h"
#include "feature.h"
#include <opencv2/imgproc/imgproc_c.h>
#include <pcl/filters/voxel_grid.h>
#include <set>
#include <sstream>
#include <iomanip>
#include <vikit/math_utils.h>
#include <vikit/robust_cost.h>
#include <vikit/vision.h>
#include <vikit/pinhole_camera.h>
#include <unordered_set>
#include <deque>

struct SubSparseMap
{
  vector<float> propa_errors;
  vector<float> errors;
  vector<vector<float>> warp_patch;
  vector<int> search_levels;
  vector<VisualPoint *> voxel_points;
  vector<double> inv_expo_list;
  vector<pointWithVar> add_from_voxel_map;
  vector<int> camera_ids; //!< camera each entry above was retrieved for, parallel to the other arrays

  SubSparseMap()
  {
    propa_errors.reserve(SIZE_LARGE);
    errors.reserve(SIZE_LARGE);
    warp_patch.reserve(SIZE_LARGE);
    search_levels.reserve(SIZE_LARGE);
    voxel_points.reserve(SIZE_LARGE);
    inv_expo_list.reserve(SIZE_LARGE);
    add_from_voxel_map.reserve(SIZE_SMALL);
    camera_ids.reserve(SIZE_LARGE);
  };

  void reset()
  {
    propa_errors.clear();
    errors.clear();
    warp_patch.clear();
    search_levels.clear();
    voxel_points.clear();
    inv_expo_list.clear();
    add_from_voxel_map.clear();
    camera_ids.clear();
  }
};

class Warp
{
public:
  Matrix2d A_cur_ref;
  int search_level;
  Warp(int level, Matrix2d warp_matrix) : search_level(level), A_cur_ref(warp_matrix) {}
  ~Warp() {}
};

class VOXEL_POINTS
{
public:
  std::vector<VisualPoint *> voxel_points;
  int count;
  double creation_timestamp_; //!< used by capVisualMap() to evict the oldest voxels first
  VOXEL_POINTS(int num) : count(num), creation_timestamp_(-1.0) {}
  ~VOXEL_POINTS()
  {
    for (VisualPoint* vp : voxel_points)
    {
      if (vp != nullptr) { delete vp; vp = nullptr; }
    }
  }
};

class VIOManager
{
public:
  int grid_size;
  int max_total_points = 300;     //!< cap on total retrieved+new points across all cameras per frame
  int points_per_camera_min = 50;
  int points_per_camera_max = 150;
  vector<vk::AbstractCamera *> cams; //!< one camera model per camera, indexed by cam_id
  StatesGroup *state;
  StatesGroup *state_propagat;
  bool raycast_en = false;
  std::vector<std::vector<std::vector<V3D>>> rays_with_sample_points; // [cam_id][grid_idx][sample_points]
  std::vector<std::vector<int>> border_flag;                          // [cam_id][grid_idx]

  // Per-camera extrinsics: Rci/Pci = IMU->camera, Rcl/Pcl = LiDAR->camera,
  // Rcw/Pcw = world->camera (recomputed every frame from the current state).
  std::vector<Eigen::Matrix3d> Rci_vec;
  std::vector<Eigen::Matrix3d> Rcl_vec;
  std::vector<Eigen::Matrix3d> Rcw_vec;
  std::vector<Eigen::Vector3d> Pci_vec;
  std::vector<Eigen::Vector3d> Pcl_vec;
  std::vector<Eigen::Vector3d> Pcw_vec;

  // Precomputed per-camera Jacobian factors (only depend on the fixed IMU->camera extrinsic).
  std::vector<Eigen::Matrix3d> Jdphi_dR_vec;
  std::vector<Eigen::Matrix3d> Jdp_dt_vec;
  std::vector<Eigen::Matrix3d> Jdp_dR_vec;

  M3D Rli; // IMU->LiDAR rotation (camera-independent)
  V3D Pli; // IMU->LiDAR translation (camera-independent)

  vector<int> grid_num;
  vector<int> map_index;
  vector<int> update_flag;
  vector<float> map_dist;
  vector<float> scan_value;
  vector<float> patch_buffer;
  bool normal_en, inverse_composition_en, exposure_estimate_en, has_ref_patch_cache;
  bool ncc_en = false, colmap_output_en = false;

  int width, height, grid_n_width, grid_n_height, length;
  double image_resize_factor;
  int patch_pyrimid_level, patch_size, patch_size_total, patch_size_half, border, warp_len;
  int max_iterations, total_points;

  double img_point_cov, outlier_threshold, ncc_thre;

  // Per-camera working state for patch selection, mirroring the single-camera
  // grid_num/map_dist/etc above but one instance per camera.
  std::vector<std::vector<int>> grid_num_per_cam_;
  std::vector<std::vector<float>> map_dist_per_cam_;
  std::vector<std::vector<VisualPoint*>> retrieve_voxel_points_per_cam_;
  std::vector<std::vector<int>> scan_grid_num_per_cam_;
  std::vector<std::vector<float>> scan_value_per_cam_;
  std::vector<std::vector<pointWithVar>> scan_append_points_per_cam_;

  std::vector<std::unordered_set<VisualPoint*>> retrieve_voxel_points_list_buffer_;
  std::vector<cv::Mat> depth_imgs_buffer_;
  std::vector<int> grid_num_buffer_;
  std::vector<float> map_dist_buffer_;
  std::vector<VisualPoint*> retrieve_voxel_points_buffer_;
  std::vector<float> scan_value_buffer_;

  SubSparseMap *visual_submap;

  double compute_jacobian_time, update_ekf_time;
  double ave_total = 0;

  int frame_count = 0;
  bool plot_flag;

  Eigen::Matrix<double, DIM_STATE, DIM_STATE> G, H_T_H;
  Eigen::MatrixXd K, H_sub_inv;

  ofstream fout_camera, fout_colmap;
  unordered_map<VOXEL_LOCATION, VOXEL_POINTS *> feat_map;
  unordered_map<VOXEL_LOCATION, int> sub_feat_map;
  unordered_map<int, Warp *> warp_map;
  vector<VisualPoint *> retrieve_voxel_points;
  vector<pointWithVar> append_voxel_points;
  FramePtr new_frame_;
  std::vector<cv::Mat> imgs_cp, imgs_rgb; //!< per-camera debug/RGB-output copies of the current frame
  cv::Mat panorama_image;                //!< mosaic of imgs_rgb across all cameras, for /rgb_img publishing

  // Per-camera photometric correction (exposure normalization + simple
  // vignetting model) applied before cross-camera photometric comparisons.
  struct CameraPhotoParams
  {
    double exposure_factor;
    std::vector<double> vignetting;
    bool parameters_initialized;
  };
  std::vector<CameraPhotoParams> camera_photo_params;

  void initializeCameraPhotoParams();
  float applyCameraPhotoCorrection(float intensity, int cam_id, const V2D &pixel_pos);
  void addCrossCameraConsistencyConstraint(VisualPoint *pt, int source_cam_id, int target_cam_id, Eigen::MatrixXd &H_sub, Eigen::VectorXd &z, int level, int &row_offset);

  int total_cross_camera_observations = 0;
  int successful_cross_camera_tracks = 0;
  bool enable_cross_camera_tracking = false;

  enum CellType
  {
    TYPE_MAP = 1,
    TYPE_POINTCLOUD,
    TYPE_UNKNOWN
  };

  VIOManager();
  ~VIOManager();
  void updateStateInverse(const std::vector<cv::Mat> &imgs, int level);
  void updateState(const std::vector<cv::Mat> &imgs, int level);
  void processFrame(const std::vector<cv::Mat> &imgs, vector<pointWithVar> &pg, const unordered_map<VOXEL_LOCATION, VoxelOctoTree *> &feat_map, double frame_timestamp);
  void retrieveFromVisualSparseMap(const std::vector<cv::Mat> imgs, vector<pointWithVar> &pg, const unordered_map<VOXEL_LOCATION, VoxelOctoTree *> &plane_map);
  void generateVisualMapPoints(const std::vector<cv::Mat> &imgs, vector<pointWithVar> &pg);
  void setImuToLidarExtrinsic(const V3D &transl, const M3D &rot);
  void setLidarToCameraExtrinsic(std::vector<std::vector<double>> &R, std::vector<std::vector<double>> &P);
  void initializeVIO();
  void initializeRaycast();
  void getImagePatch(cv::Mat img, V2D pc, float *patch_tmp, int level);
  void computeProjectionJacobian(int cam_idx, V3D p, MD(2, 3) & J);
  void computeJacobianAndUpdateEKF(const std::vector<cv::Mat> imgs);
  void resetGrid();
  void updateVisualMapPoints(std::vector<cv::Mat> &imgs);
  void getWarpMatrixAffine(const vk::AbstractCamera &cam, const Vector2d &px_ref, const Vector3d &f_ref, const double depth_ref, const SE3<double> &T_cur_ref,
                           const int level_ref,
                           const int pyramid_level, const int halfpatch_size, Matrix2d &A_cur_ref);
  void getWarpMatrixAffineHomography(const vk::AbstractCamera &cam, const V2D &px_ref,
                                     const V3D &xyz_ref, const V3D &normal_ref, const SE3<double> &T_cur_ref, const int level_ref, Matrix2d &A_cur_ref);
  void warpAffine(const Matrix2d &A_cur_ref, const cv::Mat &img_ref, const Vector2d &px_ref, const int level_ref, const int search_level,
                  const int pyramid_level, const int halfpatch_size, float *patch);
  void insertPointIntoVoxelMap(VisualPoint *pt_new);
  void setCurrentTimestamp(double timestamp) { current_timestamp_ = timestamp; }
  double current_timestamp_ = -1.0;
  void plotTrackedPoints();
  void updateFrameState(StatesGroup state);
  void projectPatchFromRefToCur(const unordered_map<VOXEL_LOCATION, VoxelOctoTree *> &plane_map);
  void updateReferencePatch(const unordered_map<VOXEL_LOCATION, VoxelOctoTree *> &plane_map);
  void precomputeReferencePatches(int level);
  void dumpDataForColmap();
  double calculateNCC(float *ref_patch, float *cur_patch, int patch_size);
  int getBestSearchLevel(const Matrix2d &A_cur_ref, const int max_level);
  V3F getInterpolatedPixel(cv::Mat img, V2D pc);

  // Bookkeeping to bound memory growth over a long run (voxel/point/frame caps).
  void cleanupCrossCameraData();
  void cleanupOldVisualPoints();
  void cleanupVisualMapByTimestamp(double oldest_kept_timestamp);
  int max_point_age_frames = 50;
  int cleanup_interval_frames = 10;
  int last_cleanup_frame_id = 0;

  void cleanupOldFrames();
  std::deque<FramePtr> frame_history_;
  int max_frame_history = 3;

  // Bound visual-map memory by capping feat_map voxel count. When enabled and
  // feat_map.size() > max_visual_voxels_, drop oldest voxels (by creation_timestamp_)
  // until the cap is satisfied.
  bool map_sliding_en_ = false;
  int max_visual_voxels_ = 10000;
  size_t capVisualMap();

  // Adaptive per-camera measurement covariance scaling based on recent
  // photometric error / point count (Omni-LIVO's "Adaptive Multi-View ESIKF").
  bool enable_dynamic_covariance_ = false;
  int dynamic_cov_warmup_frames = 200;
  double warmup_cov_scale = 500.0;
  double min_cov_scale = 10.0;
  double max_cov_scale = 2000.0;
  double dynamic_cov_error_max = 50.0;

private:
  std::vector<double> prev_cov_scale_per_cam_;
  std::vector<double> prev_avg_error_per_cam_;
  std::vector<int> prev_n_meas_per_cam_;

  double prev_cov_scale_ = 10.0;
  double prev_avg_error_ = 5.0;
  int prev_n_meas_ = 0;

  double calculateCoVarianceScale(double realtime_avg_error, int realtime_n_meas);
  double calculateCoVarianceScalePerCam(int cam_idx, double realtime_avg_error, int realtime_n_meas);

  bool isRealCrossCameraPoint(VisualPoint *pt, int current_cam_id);
  void updateCrossCameraHistory(VisualPoint *pt, int cam_id);
};
typedef std::shared_ptr<VIOManager> VIOManagerPtr;

#endif // VIO_H_