#include <costmap_2d/directional_layer.h>
#include <data_convert/data_convert.h>
#include <tf3/convert.h>
#include <tf3/utils.h>
#include <tf3_geometry_msgs/tf3_geometry_msgs.h>

#include <boost/dll/alias.hpp>

#define MAP_IDX(sx, i, j) ((sx) * (j) + (i))

using costmap_2d::CRITICAL_SPACE;
using costmap_2d::LETHAL_OBSTACLE;
using costmap_2d::NO_INFORMATION;

namespace costmap_2d
{
  DirectionalLayer::DirectionalLayer()
  {}
  
  DirectionalLayer::~DirectionalLayer() {}

  bool DirectionalLayer::laneFilter(const std::vector<geometry_msgs::PoseStamped> plan)
  {
    if (directional_areas_.empty())
      throw "Has no map data";
    nav_msgs::Path path;
    path.header.stamp = robot::Time::now();
    path.header.frame_id = map_frame_;
    path.poses = plan;
    return this->laneFilter(directional_areas_, path);
  }

  void DirectionalLayer::incomingMap(const nav_msgs::OccupancyGrid& new_map)
  {
    if(!map_shutdown_)
    {
      unsigned int size_x = new_map.info.width, size_y = new_map.info.height;

      printf("Received a %d X %d map at %f m/pix\n", size_x / 2, size_y / 2, new_map.info.resolution);

      // resize costmap if size, resolution or origin do not match
      Costmap2D *master = layered_costmap_->getCostmap();
      if (!layered_costmap_->isRolling() &&
          (master->getSizeInCellsX() != size_x / 2 ||
          master->getSizeInCellsY() != size_y / 2 ||
          master->getResolution() != new_map.info.resolution ||
          master->getOriginX() != new_map.info.origin.position.x ||
          master->getOriginY() != new_map.info.origin.position.y))
      {
        // Update the size of the layered costmap (and all layers, including this one)
        printf("Resizing costmap to %d X %d at %f m/pix\n", size_x / 2, size_y / 2, new_map.info.resolution);
        layered_costmap_->resizeMap(size_x / 2, size_y / 2, new_map.info.resolution, new_map.info.origin.position.x,
                                    new_map.info.origin.position.y,
                                    true /* set size_locked to true, prevents reconfigureCb from overriding map size*/);
      }
      else if (size_x_ != size_x / 2 || size_y_ != size_y / 2 ||
              resolution_ != new_map.info.resolution ||
              origin_x_ != new_map.info.origin.position.x ||
              origin_y_ != new_map.info.origin.position.y)
      {
        // only update the size of the costmap stored locally in this layer
        printf("Resizing static layer to %d X %d at %f m/pix\n", size_x / 2, size_y / 2, new_map.info.resolution);
        resizeMap(size_x / 2, size_y / 2, new_map.info.resolution,
                  new_map.info.origin.position.x, new_map.info.origin.position.y);
      }
      unsigned char values[4];
      directional_areas_.clear();
      directional_areas_.resize(size_x / 2 * size_y / 2);
      // initialize the costmap with static data
      for (unsigned int iy = 0; iy < size_y; iy += 2)
      {
        for (unsigned int ix = 0; ix < size_x; ix += 2)
        {
          values[0] = new_map.data[MAP_IDX(size_x, ix, iy)];
          values[1] = new_map.data[MAP_IDX(size_x, ix + 1, iy)];
          values[2] = new_map.data[MAP_IDX(size_x, ix, iy + 1)];
          values[3] = new_map.data[MAP_IDX(size_x, ix + 1, iy + 1)];
          uint32_t color_avg = (uint32_t)(values[0] | values[1] << 8u | values[2] << 16u | values[3] << 24u);
          int index = getIndex(ix / 2, iy / 2);
          directional_areas_[index][0] = (uint16_t)(color_avg & 0xfff);
          directional_areas_[index][1] = (uint16_t)(color_avg >> 12u);
          costmap_[index] = costmap_2d::NO_INFORMATION;
        }
      }
      map_frame_ = new_map.header.frame_id;

      // we have a new map, update full size of map
      x_ = y_ = 0;
      width_ = size_x_;
      height_ = size_y_;
      map_received_ = true;
      has_updated_data_ = true;

      new_map_.header = new_map.header;
      new_map_.header.stamp = robot::Time::now();
      new_map_.info = new_map.info;
      new_map_.info.width = width_;
      new_map_.info.height = height_;
    }
    else
    {
      printf("Stop receive new map!\n");
    }
    // shutdown the map subscrber if firt_map_only_ flag is on
    if (first_map_only_)
    {
      printf("Shutting down the map subscriber. first_map_only flag is on\n");
      map_shutdown_ = true;
      // map_sub_.shutdown();
    }
  }

  bool DirectionalLayer::laneFilter(std::vector<std::array<uint16_t, 2>> new_map, const nav_msgs::Path path)
  {
    boost::unique_lock<mutex_t> lock(*layered_costmap_->getCostmap()->getMutex());
    if (new_map.empty())
      return false;

    std::vector<std::pair<unsigned int, double>> X_List, Y_List;
    std::vector<double> Yaw_list;

    const geometry_msgs::PoseStamped &e = path.poses.back();
    
    bool get_success = getRobotPose(pose_x_, pose_y_, pose_yaw_);
    if(!get_success) return false;
    for (auto it = path.poses.begin(); it != path.poses.end(); ++it)
    {
      const geometry_msgs::PoseStamped &p = *it; // Get the pose from the iterator
      unsigned int mx, my;
      if (!worldToMap(p.pose.position.x, p.pose.position.y, mx, my))
      {
        printf("ERROR: trying to set a start cell %f %f that is outside of map\n", p.pose.position.x, p.pose.position.y);
        return false;
      }
      // Convert to yaw
      tf3::Quaternion quaternion = data_convert::convertQuaternion(p.pose.orientation);
      double theta = tf3::getYaw(quaternion);
      unsigned char value = laneFilter(mx, my, theta);
      if (get_success)
      {
        if (
            (!inSkipErea(pose_x_, pose_y_, p.pose.position.x, p.pose.position.y, distance_skip_) & 
            !inSkipErea(p.pose.position.x, p.pose.position.y, e.pose.position.x, e.pose.position.y, distance_skip_))
            || p == path.poses.back()
           )
        {
          if (value == costmap_2d::INSCRIBED_INFLATED_OBSTACLE || value == costmap_2d::LETHAL_OBSTACLE)
          {
            std::pair<unsigned int, double> x_val(mx, p.pose.position.x);
            std::pair<unsigned int, double> y_val(my, p.pose.position.y);
            X_List.push_back(x_val);
            Y_List.push_back(y_val);
            Yaw_list.push_back(theta);
            // costmap_[getIndex(mx, my)] = value;
          }
        }
      }
    }
    if (!Yaw_list.empty())
    {
      laneFilter(X_List, Y_List, Yaw_list);
      nav_msgs::OccupancyGrid lanes;
      convertToMap(costmap_, lanes, 0.65, 0.196);

      //////////////////////////////////
      //////////////////////////////////
      /////////THAY THẾ PUBLISH////////
      // lane_mask_pub_.publish(lanes);
      //////////////////////////////////
      //////////////////////////////////
      //////////////////////////////////

      
      return false;
    }
    return true;
    
  }

  void DirectionalLayer::convertToMap(unsigned char *costmap, nav_msgs::OccupancyGrid &dir_map, double occ_th, double free_th)
  {
    dir_map.header = new_map_.header;
    dir_map.header.stamp = robot::Time::now();
    dir_map.info = new_map_.info;
    int size_costmap = new_map_.info.width * new_map_.info.height;
    dir_map.data.resize(size_costmap);
    for (int i = 0; i < size_costmap; i++)
    {
      int8_t value;
      if (costmap[i] == costmap_2d::NO_INFORMATION)
      {
        value = -1;
      }
      else
      {
        double occ = (costmap[i]) / 255.0;
        if (occ > occ_th)
          value = +100;
        else if (occ < free_th)
          value = 0;
        else
        {
          double ratio = (occ - free_th) / (occ_th - free_th);
          value = 1 + 98 * ratio;
        }
      }
      dir_map.data[i] = value;
    }
  }

  unsigned char DirectionalLayer::laneFilter(unsigned int x, unsigned int y, double yaw_robot)
  {
    double yaw_lane;
    unsigned char cost = costmap_2d::NO_INFORMATION;

    int index = getIndex(x, y);
    for (auto &lane : directional_areas_[index])
    {
      if (lane > 359)
      {
        cost = std::min(cost, costmap_2d::NO_INFORMATION);
      }
      else
      {
        double yaw_lane = (double)lane / 180 * M_PI;
        if (0.4 <= cos(yaw_robot - yaw_lane) && cos(yaw_robot - yaw_lane) <= 1.0)
          cost = std::min(cost, costmap_2d::FREE_SPACE);
        else if (-1.0 <= cos(yaw_robot - yaw_lane) && cos(yaw_robot - yaw_lane) <= -0.4)
          cost = std::min(cost, costmap_2d::LETHAL_OBSTACLE);
        else
          cost = std::min(cost, costmap_2d::NO_INFORMATION);
      }
    }
    return cost;
  }

  void DirectionalLayer::laneFilter(std::vector<std::pair<unsigned int, double>> x, 
                                    std::vector<std::pair<unsigned int, double>> y, 
                                    std::vector<double> yaw_robot)
  {
    if(x.empty() || y.empty() || yaw_robot.empty()) 
      return;

    unsigned int x_min, y_min, x_max, y_max;
    double x_min_w, y_min_w, x_max_w, y_max_w;

    x_min = x.front().first; y_min = y.front().first;
    x_max = x.front().first; y_max = y.front().first;
    x_min_w = x.front().second; y_min_w = y.front().second;
    x_max_w = x.front().second; y_max_w = y.front().second;
    for (int i = 1; i < yaw_robot.size(); i++)
    {
      if(inSkipErea(pose_x_, pose_y_, x[i].second, y[i].second, distance_skip_))
        continue;
      x_min = std::min(x_min, x[i].first);
      y_min = std::min(y_min, y[i].first);
      x_max = std::max(x_max, x[i].first);
      y_max = std::max(y_max, y[i].first);

      x_min_w = std::min(x_min_w, x[i].second);
      y_min_w = std::min(y_min_w, y[i].second);
      x_max_w = std::max(x_max_w, x[i].second);
      y_max_w = std::max(y_max_w, y[i].second);

    }
    // printf("%d %d %d %d", x_min, y_min, x_max, y_max);
    // printf("%f %f %f %f", x_min_w, y_min_w, x_max_w, y_max_w);
    for (int i = 0; i < yaw_robot.size(); i++)
    {
      if(inSkipErea(pose_x_, pose_y_, x[i].second, y[i].second, distance_skip_))
        continue;

      double yaw_rad = (double)yaw_robot[i] / 180 * M_PI;
      if (fabs(cos(yaw_rad)) > fabs(sin(yaw_rad)))
      {
        int x_ = x[i].first;
        // Dưới lên trên
        for (int j = y[i].first; j <= y_max; j++)
        {
          int y_ = j;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }
        // Trên xuống dưới
        for (int k = y[i].first; k >= y_min; k--)
        {
          int y_ = k;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }

        int y_ = y[i].first;
        // Phải qua trái
        for (int j = x[i].first; j < layered_costmap_->getCostmap()->getSizeInCellsX(); j++)
        {
          int x_ = j;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }
        // Trái qua phải
        for (int k = x[i].first; k >= 0; k--)
        {
          int x_ = k;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }       
      }
      else if(fabs(cos(yaw_rad)) < fabs(sin(yaw_rad)))
      {
        int y_ = y[i].first;
        // Phải qua trái
        for (int j = x[i].first; j <= x_max; j++)
        {
          int x_ = j;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }
        // Trái qua phải
        for (int k = x[i].first; k >= x_min; k--)
        {
          int x_ = k;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }

        int x_ = x[i].first;
        // Dưới lên trên
        for (int j = y[i].first; j < layered_costmap_->getCostmap()->getSizeInCellsY(); j++)
        {
          int y_ = j;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }
        // Trên xuống dưới
        for (int k = y[i].first; k >= 0; k--)
        {
          int y_ = k;
          if (costmap_[getIndex(x_, y_)] == costmap_2d::LETHAL_OBSTACLE)
            continue;
          unsigned char value = laneFilter(x_, y_, yaw_robot[i]);
          if (value >= costmap_2d::INSCRIBED_INFLATED_OBSTACLE && value != costmap_2d::NO_INFORMATION)
            costmap_[getIndex(x_, y_)] = value;
          else
            break;
        }        
      }
    }
  }

  void DirectionalLayer::resetMap()
  {
    unsigned int size_x = layered_costmap_->getCostmap()->getSizeInCellsX();
    unsigned int size_y = layered_costmap_->getCostmap()->getSizeInCellsY();
    for (unsigned int iy = 0; iy < size_y; iy++)
    {
      for (unsigned int ix = 0; ix < size_x; ix++)
      {
        int index = getIndex(ix, iy);
        costmap_[index] = costmap_2d::NO_INFORMATION;
      }
    }
  }

  bool DirectionalLayer::getRobotPose(double &x, double &y, double &yaw)
  {
    tf3::TransformStampedMsg transformMsg;
    try
    {
      transformMsg = tf_->lookupTransform(map_frame_, base_frame_id_, tf3::Time());
        
      // listener_.lookupTransform(map_frame_, base_frame_id_, tf3::Time(), transform);
    }
    catch (tf3::TransformException &ex)
    {
      printf("%s\n", ex.what());
      return false;
    }
    // Copy map data given proper transformations
    tf3::Transform tf3_transform;
    tf3_transform = data_convert::convertTotf3Transform(transformMsg.transform);
    x = tf3_transform.getOrigin().x();
    y = tf3_transform.getOrigin().y();
    // Extract the rotation as a quaternion from the transform
    tf3::Quaternion rotation = tf3_transform.getRotation();
    // Convert the quaternion to a yaw angle (in radians)
    yaw = tf3::getYaw(rotation);
    return true;
  }

  bool DirectionalLayer::inSkipErea(double start_x, double start_y, double end_x, double end_y, double skip_distance)
  {
    return fabs(hypot(start_x - end_x, start_y - end_y)) <= skip_distance;
  }

  // Export factory function
  static boost::shared_ptr<Layer> create_directional_plugin() {
      return boost::make_shared<DirectionalLayer>();
  }

  // Alias cho Boost.DLL (nếu muốn dùng boost::dll::import_alias)
  BOOST_DLL_ALIAS(create_directional_plugin, create_directional_layer)
}