#include "visualization.h"

#include <iostream>
#include <algorithm>
#include <numeric>
#include <time.h>

#include <tf2/LinearMath/Quaternion.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

#include "utils.h"

using namespace std;

Visualizer::Visualizer(rclcpp::Node::SharedPtr node, VisualizationConfig& config)
    : node_(node), frame_id_(config.frame_id) {

  kcc_pose_pub_ = node_->create_publisher<nav_msgs::msg::Path>(config.kcc_pose_topic, 10);
  frame_pose_pub_ = node_->create_publisher<nav_msgs::msg::Path>(config.frame_pose_topic, 10);
  map_pub_ = node_->create_publisher<nav_msgs::msg::OccupancyGrid>(config.map_topic, 1);
  image_pub_ = node_->create_publisher<sensor_msgs::msg::Image>(config.image_topic, 1);

  auto current_time = node_->get_clock()->now();
  odom_pose_msgs_.header.stamp = current_time;
  odom_pose_msgs_.header.frame_id = frame_id_;
  kcc_pose_msgs_.header.stamp = current_time;
  kcc_pose_msgs_.header.frame_id = frame_id_;
  frame_pose_msgs_.header.stamp = current_time;
  frame_pose_msgs_.header.frame_id = frame_id_;

  occupancy_map_msgs_.header.stamp = current_time;
  occupancy_map_msgs_.header.frame_id = frame_id_;
}

void Visualizer::AddNewPoseToPath(
    Eigen::Vector3d& pose, double time_double,
    nav_msgs::msg::Path& path, std::string& id) {

  auto current_time = node_->get_clock()->now();

  geometry_msgs::msg::PoseStamped pose_stamped;

  if (time_double < 0) {
    pose_stamped.header.stamp = current_time;
  } else {
    auto sec = static_cast<int64_t>(time_double);
    auto nsec = static_cast<uint32_t>((time_double - sec) * 1e9);
    pose_stamped.header.stamp = rclcpp::Time(sec, nsec);
  }
  pose_stamped.header.frame_id = id;

  pose_stamped.pose.position.x = pose(0);
  pose_stamped.pose.position.y = pose(1);
  pose_stamped.pose.position.z = 0;

  // tf2 equivalent of tf::createQuaternionMsgFromYaw
  tf2::Quaternion q;
  q.setRPY(0, 0, pose(2));
  pose_stamped.pose.orientation = tf2::toMsg(q);

  path.poses.push_back(pose_stamped);
}

void Visualizer::UpdateOdomPose(Eigen::Vector3d& pose, double time_double) {
  AddNewPoseToPath(pose, time_double, odom_pose_msgs_, frame_id_);
  // odom_pose_pub not used currently
}

void Visualizer::UpdateKccPose(Eigen::Vector3d& pose, double time_double) {
  AddNewPoseToPath(pose, time_double, kcc_pose_msgs_, frame_id_);
  kcc_pose_pub_->publish(kcc_pose_msgs_);
}

void Visualizer::UpdateFramePose(
    Aligned<std::vector, Eigen::Vector3d>& frame_poses,
    std::vector<double>& timestamps) {
  frame_pose_msgs_.poses.clear();
  for (size_t i = 0; i < frame_poses.size(); i++) {
    AddNewPoseToPath(frame_poses[i], timestamps[i], frame_pose_msgs_, frame_id_);
  }
  frame_pose_pub_->publish(frame_pose_msgs_);
}

void Visualizer::ConvertMapToOccupancyMsgs(
    OccupancyData& map, nav_msgs::msg::OccupancyGrid& msgs) {
  if (map.size() < 1) return;

  int max_x = std::numeric_limits<int>::min();
  int max_y = std::numeric_limits<int>::min();
  int min_x = std::numeric_limits<int>::max();
  int min_y = std::numeric_limits<int>::max();

  for (auto& kv : map) {
    int x = kv.first.x;
    int y = kv.first.y;
    max_x = std::max(max_x, x);
    max_y = std::max(max_y, y);
    min_x = std::min(min_x, x);
    min_y = std::min(min_y, y);
  }

  if ((min_x > max_x) || (min_y > max_y)) return;

  int size = map.begin()->second.size;
  msgs.info.width = (max_x - min_x + 1) * size;
  msgs.info.height = (max_y - min_y + 1) * size;

  msgs.info.origin.position.x = (min_x * size);
  msgs.info.origin.position.y = (min_y * size);
  msgs.info.origin.position.z = 0;

  int num_grids = msgs.info.width * msgs.info.height;
  std::vector<int8_t> data(num_grids, -1);

  for (auto& kv : map) {
    for (int i = 0; i < kv.second.size; i++) {
      for (int j = 0; j < kv.second.size; j++) {
        if (kv.second.weight(i, j) < 1) continue;

        int x = (kv.first.x - min_x) * kv.second.size + j;
        int y = (kv.first.y - min_y) * kv.second.size + i;
        int idx = y * msgs.info.width + x;
        int pixel = static_cast<int>(kv.second.data(i, j));

        data[idx] = 100 - static_cast<int8_t>(pixel);
      }
    }
  }
  msgs.data = data;
}

void Visualizer::UpdateMap(MapBuilder& map_builder) {
  occupancy_map_msgs_.info.resolution = map_builder.GetMapResolution();
  OccupancyData map_data = map_builder.GetMapData();
  if (map_data.size() < 1) return;
  ConvertMapToOccupancyMsgs(map_data, occupancy_map_msgs_);

  Eigen::Matrix<double, 7, 1> real_origin;
  Eigen::Vector3d pixel_origin;
  pixel_origin << occupancy_map_msgs_.info.origin.position.x,
                  occupancy_map_msgs_.info.origin.position.y,
                  occupancy_map_msgs_.info.origin.position.z;
  map_builder.GetOccupancyMapOrigin(pixel_origin, real_origin);
  occupancy_map_msgs_.info.origin.orientation.w = real_origin(0, 0);
  occupancy_map_msgs_.info.origin.orientation.x = real_origin(1, 0);
  occupancy_map_msgs_.info.origin.orientation.y = real_origin(2, 0);
  occupancy_map_msgs_.info.origin.orientation.z = real_origin(3, 0);
  occupancy_map_msgs_.info.origin.position.x = real_origin(4, 0);
  occupancy_map_msgs_.info.origin.position.y = real_origin(5, 0);
  occupancy_map_msgs_.info.origin.position.z = real_origin(6, 0);

  map_pub_->publish(occupancy_map_msgs_);
}

void Visualizer::PublishImage(cv::Mat& image, double time_double) {
  auto image_msg = cv_bridge::CvImage(std_msgs::msg::Header(), "mono8", image).toImageMsg();
  if (time_double > 0) {
    auto sec = static_cast<int64_t>(time_double);
    auto nsec = static_cast<uint32_t>((time_double - sec) * 1e9);
    image_msg->header.stamp = rclcpp::Time(sec, nsec);
  }
  image_pub_->publish(*image_msg);
}

void Visualizer::GetTrajectoryTxt(
    std::vector<std::vector<std::string>>& lines,
    TrajectoryType trajectory_type) {
  nav_msgs::msg::Path* path_ptr;
  switch (trajectory_type) {
    case TrajectoryType::Frame:
      path_ptr = &frame_pose_msgs_;
      break;
    case TrajectoryType::KCC:
      path_ptr = &kcc_pose_msgs_;
      break;
    case TrajectoryType::Odom:
      path_ptr = &odom_pose_msgs_;
      break;
    default:
      std::cout << "please select trajectory_type from Frame, KCC and Odom !" << std::endl;
      return;
  }

  for (geometry_msgs::msg::PoseStamped& pose_stamped : (*path_ptr).poses) {
    std::vector<std::string> line;

    double time_double = static_cast<double>(pose_stamped.header.stamp.sec) +
                         static_cast<double>(pose_stamped.header.stamp.nanosec) / 1e9;
    std::string s_time = std::to_string(time_double);

    line.emplace_back(s_time);
    line.emplace_back(std::to_string(pose_stamped.pose.position.x));
    line.emplace_back(std::to_string(pose_stamped.pose.position.y));
    line.emplace_back(std::to_string(pose_stamped.pose.position.z));
    line.emplace_back(std::to_string(pose_stamped.pose.orientation.x));
    line.emplace_back(std::to_string(pose_stamped.pose.orientation.y));
    line.emplace_back(std::to_string(pose_stamped.pose.orientation.z));
    line.emplace_back(std::to_string(pose_stamped.pose.orientation.w));
    lines.push_back(line);
  }
}
