#include <cstdio>
#include <iostream>
#include <fstream>

#include <memory>

#include <rclcpp/rclcpp.hpp>

#include <boost/thread.hpp>

#include <ompl-1.7/ompl/geometric/planners/rrt/RRTConnect.h>

#include <ompl-1.7/ompl/geometric/PathSimplifier.h>
#include <ompl-1.7/ompl/base/SpaceInformation.h>

#include <ompl-1.7/ompl/base/objectives/StateCostIntegralObjective.h>
#include <ompl-1.7/ompl/base/objectives/PathLengthOptimizationObjective.h>

#include <ompl-1.7/ompl/base/samplers/ObstacleBasedValidStateSampler.h>

#include <grid_map_geo/grid_map_geo.hpp>

#include <ompl_planner_owen_path_cost/PathValidityChecker.hpp>

#include "ompl_planner_owen_path_cost/ompl_planner.hpp"

#include <ompl-1.7/ompl/base/spaces/Dubins3DMotionValidator.h>

#include <grid_map_core/iterators/CircleIterator.hpp>
#include "grid_map_core/GridMapMath.hpp"

namespace og = ompl::geometric;
namespace ob = ompl::base;

#define DIM 3                         // Dimension of the state space

/**
 * @brief Planner class
 */
class OMPL_Planner
{
  std::shared_ptr<ob::OwenStateSpace> space;
  std::shared_ptr<GridMapGeo> terrain_map;
  std::shared_ptr<ompl::PathValidityChecker> validity_checker;
  ob::ValidStateSamplerPtr ge_sampler;

  double Range = 100;
  double Max_wave = 2.0; // Size of the biggest wave allowed to go over
  double Cord = 1.0;     // Cord of the wing [m]

  std::shared_ptr<ob::ScopedState<ob::OwenStateSpace>> start, goal;

  bool boundsSet = false, startSet = false, goalSet = false, mapSet = false, obstaclesSet = false, pathSet = false;

  std::vector<Pose> Path;

  std::vector<Obstacle> Obstacles;

  // Functions

  /**
   * @brief Path extraction from ProblemDefinition
   */
  nav_msgs::msg::Path extractPath(ob::ProblemDefinition *probDef);

  /**
   * @brief Use terrain map to set bounds
   */
  void setBoundsFromMap(const grid_map::GridMap &map);

  /**
   * @brief Calculate distance to objective path
   */
  ob::Cost distanceToPath();

public:
  /**
   * @brief Basic constructor
   */
  OMPL_Planner()
  {
    setBounds(0, 0, 0, 0, 0, 0);
  }

  /**
   * @brief Constructor
   * @param map GridMap of the area, may include layers to show maximum and minimum allowable height (see param setLayers)
   * @param start Start position
   * @param goal Goal position
   * @param setLayers false if GridMap already comes with bounding layers, true if it doesn't
   */
  OMPL_Planner(const grid_map::GridMap &map, geometry_msgs::msg::PoseStamped start, geometry_msgs::msg::PoseStamped goal, bool setLayers = false);

  /**
   * @brief Constructor
   * @param range Desired range of the planner
   * @param width width of the map (X coordinates)
   * @param height height of the map (Y coordinates), not to mistake with elevation
   * @param start Start position
   * @param goal Goal position
   */
  OMPL_Planner(double range, int width, int height, geometry_msgs::msg::PoseStamped start, geometry_msgs::msg::PoseStamped goal);

  ~OMPL_Planner()
  {
  }

  /**
   * @brief Path planning function using gridmap
   */
  nav_msgs::msg::Path pathPlanner(double time = 1.0);

  /**
   * @brief Set Start
   * @param poseMsg start position as a ROS2 message
   */
  void setStart(geometry_msgs::msg::PoseStamped poseMsg);

  /**
   * @brief Set Goal
   * @param poseMsg start position as a ROS2 message
   */
  void setGoal(geometry_msgs::msg::PoseStamped poseMsg);

  /**
   * @brief Set Bounds
   */
  void setBounds(double X_min, double X_max, double Y_min, double Y_max, double Z_min, double Z_max);

  /**
   * @brief Set Obstacles
   * @param obstacles_ Pointer to the list of obstacles as a BoundingBox3DArray
   */
  void setObstacles(const vision_msgs::msg::BoundingBox3DArray &obstacles_);

  /**
   * @brief Set Map
   */
  void setMap(const grid_map::GridMap &map, bool setLayers = false);

  /**
   * @brief Get Map
   * @return Map as a GridMap
   */
  grid_map::GridMap getMap() { return terrain_map->getGridMap(); }

  /**
   * \brief Set Path
   */
  void setPath(std::vector<Pose> Path_msg)
  {
    Path = Path_msg;
    pathSet = true;
  }

  /**
   * @brief Get Path
   */
  std::vector<Pose> getPath()
  {
    return Path;
  }
};



// Constructor
OMPL_Planner::OMPL_Planner(const grid_map::GridMap &map, geometry_msgs::msg::PoseStamped start, geometry_msgs::msg::PoseStamped goal, bool setLayers)
{
  grid_map::GridMap temp_map_ = map;
  if (Range > map.getResolution())
    Range = map.getResolution();
  setMap(map, setLayers);
  setStart(start);
  setGoal(goal);
}

// Constructor
OMPL_Planner::OMPL_Planner(double range, int width, int height, geometry_msgs::msg::PoseStamped start, geometry_msgs::msg::PoseStamped goal)
{
  Range = range;
  setBounds(-range * width / 2, range * width / 2, -range * height / 2, range * height / 2, 0.0, Max_wave + Cord);
  setStart(start);
  setGoal(goal);
}

// Path planning function using gridmap
nav_msgs::msg::Path OMPL_Planner::pathPlanner(double time)
{
  std::shared_ptr<ob::OwenStateSpace> space; // In here it is stated to use owen paths
  std::shared_ptr<GridMapGeo> terrain_map; // Uses GridMapGeo for the map
  auto si(std::make_shared<ob::SpaceInformation>(space));

  // It is possible for obstacles to exist, composed of bounding boxes
  if (obstaclesSet)
  {
    validity_checker = std::make_shared<ompl::PathValidityChecker>(si, terrain_map->getGridMap(), Path);
    si->setStateValidityChecker(ob::StateValidityCheckerPtr(validity_checker));
  }

  // If there are no obstacles
  else
  {
    validity_checker = std::make_shared<ompl::PathValidityChecker>(si, terrain_map->getGridMap(), Path, Obstacles);
    si->setStateValidityChecker(ob::StateValidityCheckerPtr(validity_checker));
  }
  si->setMotionValidator(std::make_shared<ob::Dubins3DMotionValidator<ob::OwenStateSpace>>(si));

  auto probDef(std::make_shared<ob::ProblemDefinition>(si));
  probDef->setStartAndGoalStates(*start, *goal);
  probDef->setOptimizationObjective(getOptObj(si));

  auto planner(std::make_shared<og::RRTConnect>(si));// In here it is stated to use RRTConnect
  planner->setRange(Range);
  // planner->setTreePruning(false);
  // planner->setDelayCC(true);
  planner->setProblemDefinition(probDef);
  planner->setup();
  printf("..................PLANNER HAS THE FOLLOWING SETTINGS..................\n");
  si->printSettings(std::cout);
  printf("......................................................................\n");
  probDef->print(std::cout);
  printf("......................................................................\n");

  ob::PlannerStatus solved = planner->ob::Planner::solve(time);

  nav_msgs::msg::Path return_path;

  if (solved == ob::PlannerStatus::EXACT_SOLUTION || solved == ob::PlannerStatus::APPROXIMATE_SOLUTION)
  {
    return_path = extractPath(probDef.get());
    std::cout << "INFO: ompl_planner_owen-- Found " << solved.asString() << std::endl;
  }
  else
  {
    std::cout << "\n----------------------PlannerStatus::" << solved.asString() << "----------------------\n";
  }

  std::cout << "Exracted path" << std::endl;

  return return_path;
}

nav_msgs::msg::Path OMPL_Planner::extractPath(ob::ProblemDefinition *probDef)
{
  nav_msgs::msg::Path extractedPath;

  extractedPath.header.frame_id = "/map";

  ob::PathPtr path_ = probDef->getSolutionPath();

  const og::PathGeometric *pathPtr = path_.get()->as<og::PathGeometric>();

  og::PathGeometric path = *pathPtr;

  double distance = 0;
  for (unsigned int i = 0; i < path.getStateCount(); i++)
  {
    const ob::State *state = path.getState(i);
    const auto &st = *state->as<ob::OwenStateSpace::StateType>();

    const double coordX = st[0];
    const double coordY = st[1];
    const double coordZ = st[2];

    const Quaternion quat = Euler2Quaternion(0, 0, st.yaw());

    geometry_msgs::msg::PoseStamped poseMsg;
    poseMsg.pose.position.x = coordX;
    poseMsg.pose.position.y = coordY;
    poseMsg.pose.position.z = coordZ;
    poseMsg.pose.orientation.w = quat.w;
    poseMsg.pose.orientation.x = quat.x;
    poseMsg.pose.orientation.y = quat.y;
    poseMsg.pose.orientation.z = quat.z;
    poseMsg.header.frame_id = "/map";
    poseMsg.header.stamp = rclcpp::Clock().now();
    extractedPath.poses.push_back(poseMsg);
    if (i < path.getStateCount() - 1)
    {
      const ob::State *stateNext = path.getState(i + 1);
      const auto &stNext = *stateNext->as<ob::OwenStateSpace::StateType>();
      const auto coordXN = stNext[0];
      const auto coordYN = stNext[1];
      const auto coordZN = stNext[2];
      distance += sqrt((coordX - coordXN) * (coordX - coordXN) +
                       (coordY - coordYN) * (coordY - coordYN) +
                       (coordZ - coordZN) * (coordZ - coordZN));
    }
  }
  std::cout << "\nINFO: ompl_planner_owen::extractPath--  Distance: " << distance << std::endl;
  return extractedPath;
}

// Set Start
void OMPL_Planner::setStart(geometry_msgs::msg::PoseStamped poseMsg)
{
  start.reset(new ob::ScopedState<ob::OwenStateSpace>(space));

  (*start.get())[0] = (poseMsg.pose.position.x);
  (*start.get())[1] = (poseMsg.pose.position.y);
  (*start.get())[2] = (poseMsg.pose.position.z);

  Euler eul = Quaternion2Euler(poseMsg.pose.orientation.w, poseMsg.pose.orientation.x, poseMsg.pose.orientation.y, poseMsg.pose.orientation.z);
  (*start.get())->yaw() = eul.yaw;

  grid_map::Position pos{(*start.get())[0], (*start.get())[1]};

  //(*start.get()).random();
  startSet = true;
}

// Set Goal
void OMPL_Planner::setGoal(geometry_msgs::msg::PoseStamped poseMsg)
{
  goal.reset(new ob::ScopedState<ob::OwenStateSpace>(space));
  (*goal.get())[0] = (poseMsg.pose.position.x);
  (*goal.get())[1] = (poseMsg.pose.position.y);
  (*goal.get())[2] = (poseMsg.pose.position.z);
  Euler eul = Quaternion2Euler(poseMsg.pose.orientation.w, poseMsg.pose.orientation.x, poseMsg.pose.orientation.y, poseMsg.pose.orientation.z);
  (*goal.get())->yaw() = eul.yaw;

  grid_map::Position pos{(*goal.get())[0], (*goal.get())[1]};

  //(*goal.get()).random();
  goalSet = true;
}

// Set Bounds
void OMPL_Planner::setBounds(double X_min, double X_max, double Y_min, double Y_max, double Z_min, double Z_max)
{
  ob::RealVectorBounds bounds(DIM);
  bounds.setLow(0, X_min);
  bounds.setHigh(0, X_max);
  bounds.setLow(1, Y_min);
  bounds.setHigh(1, Y_max);
  bounds.setLow(2, Z_min);
  bounds.setHigh(2, Z_max);
  auto BoundsState(std::make_shared<ob::OwenStateSpace>(500.0, 0.123));
  space = BoundsState;
  BoundsState->setBounds(bounds);
  boundsSet = true;
}

// Set Obstacles using BoundingBox3DArray
void OMPL_Planner::setObstacles(const vision_msgs::msg::BoundingBox3DArray &obstacles_)
{
  Obstacles.erase(Obstacles.begin(), Obstacles.end());
  for (int i = 0; i < (int)obstacles_.boxes.size(); i++)
  {
    // Obtain transformation matrix for the new reference system of each box
    Eigen::Quaternion q(obstacles_.boxes[i].center.orientation.w,
                        obstacles_.boxes[i].center.orientation.x,
                        obstacles_.boxes[i].center.orientation.y,
                        obstacles_.boxes[i].center.orientation.z);
    Eigen::Matrix3d Rmat = q.normalized().toRotationMatrix();
    Eigen::Matrix4d Tmat{
        {Rmat(0, 0), Rmat(0, 1), Rmat(0, 2), obstacles_.boxes[i].center.position.x},
        {Rmat(1, 0), Rmat(1, 1), Rmat(1, 2), obstacles_.boxes[i].center.position.y},
        {Rmat(2, 0), Rmat(2, 1), Rmat(2, 2), obstacles_.boxes[i].center.position.z},
        {0, 0, 0, 1}};
    Obstacle BOXi;
    BOXi.transformation_mat = Tmat;

    // Get the size of the obstacle
    BOXi.boundaries[0] = obstacles_.boxes[i].size.x;
    BOXi.boundaries[1] = obstacles_.boxes[i].size.y;
    BOXi.boundaries[2] = obstacles_.boxes[i].size.z;

    Obstacles.push_back(BOXi);
    if (Range > obstacles_.boxes[i].size.x)
      Range = obstacles_.boxes[i].size.x;
    if (Range > obstacles_.boxes[i].size.y)
      Range = obstacles_.boxes[i].size.y;
    if (Range > obstacles_.boxes[i].size.z)
      Range = obstacles_.boxes[i].size.z;
  }
}

// Set Map
void OMPL_Planner::setMap(const grid_map::GridMap &map, bool setLayers)
{
  std::cout << "INFO: ompl_planner_owen::setMap-- SetMap\n";
  grid_map::GridMap gridmap = map;
  terrain_map = std::make_shared<GridMapGeo>();
  terrain_map->setGridMap(gridmap);

  if (setLayers)
  {
    terrain_map->AddLayerDistanceTransform(15.0 + Cord, "max_elevation");

    terrain_map->AddLayerDistanceTransform(0.25 * Cord, "distance_surface");

    double radius = 10.0;

    terrain_map->AddLayerHorizontalDistanceTransform(radius, "ics_+", "distance_surface");

    terrain_map->AddLayerHorizontalDistanceTransform(-radius, "ics_-", "max_elevation");
  }

  if (Range > map.getResolution())
    Range = map.getResolution();

  setBoundsFromMap(terrain_map->getGridMap());

  mapSet = true;
}

// Use terrain map to set bounds
void OMPL_Planner::setBoundsFromMap(const grid_map::GridMap &map)
{
  double max_point = std::numeric_limits<double>::min(), min_point = std::numeric_limits<double>::max();

  for (grid_map::GridMapIterator it(map); !it.isPastEnd(); ++it)
  {
    const grid_map::Index map_index = *it;
    const double min_point_it = map.at("ics_+", map_index), max_point_it = map.at("ics_-", map_index);
    if (min_point_it < min_point)
      min_point = min_point_it;
    if (max_point_it > max_point)
      max_point = max_point_it;
  }
  setBounds(-map.getLength().x() / 2, map.getLength().x() / 2, -map.getLength().y() / 2, map.getLength().y() / 2, min_point, max_point);
}

#include <nav_msgs/msg/path.hpp>
#include <nav_msgs/msg/occupancy_grid.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <vision_msgs/msg/bounding_box3_d_array.hpp>

#include <ompl-1.7/ompl/base/spaces/OwenStateSpace.h>

#include <grid_map_core/GridMap.hpp>

#include <eigen3/Eigen/Dense>

#include <ompl-1.7/ompl/base/StateValidityChecker.h>

#include <grid_map_core/iterators/CircleIterator.hpp>
#include "grid_map_core/GridMapMath.hpp"
#include <limits.h>


// Structs

// Struct for a signle obstacle
struct Obstacle
{
    Eigen::Matrix4d transformation_mat;
    Eigen::Vector3d boundaries;
};

// Struct for position and orientation
struct Pose
{
    Eigen::Vector3d position;
    Eigen::Vector4d orientation;
};

struct Quaternion
{
    double w, x, y, z;
};

// Function definitions

/**
 * @brief Convert Euler angles to quaternion
 * @param roll Bank (rad)
 * @param pitch Attitude (rad)
 * @param yaw Heading (rad)
 * @return Quaternion struct
 */
Quaternion Euler2Quaternion(double roll, double pitch, double yaw)
{
    double cr = cos(roll * 0.5);
    double sr = sin(roll * 0.5);
    double cp = cos(pitch * 0.5);
    double sp = sin(pitch * 0.5);
    double cy = cos(yaw * 0.5);
    double sy = sin(yaw * 0.5);

    Quaternion q;
    q.w = cr * cp * cy + sr * sp * sy;
    q.x = sr * cp * cy - cr * sp * sy;
    q.y = cr * sp * cy + sr * cp * sy;
    q.z = cr * cp * sy - sr * sp * cy;

    return q;
}

// Variables

grid_map::Length mapLength_;
grid_map::Position mapPosition_, itPos;
double resolution_;
grid_map::Size bufferSize_;
grid_map::Index bufferStartIndex_;

namespace ompl
{
  class PathValidityChecker : public base::StateValidityChecker
    {
    public:
        /**
         * @brief Constructor of the owen validity checker
         * @param Map Map of the area as a GridMap
         */
        PathValidityChecker(const base::SpaceInformationPtr &space_info, const grid_map::GridMap &map_, const std::vector<Pose> Path)
            : base::StateValidityChecker(space_info), gridmap(map_), path(Path) {}

        /**
         * @brief Constructor of the owen validity checker including obstacles
         * @param Map Map of the area as a GridMap
         * @param Obstacles List of obstacles as an array
         */
        PathValidityChecker(const base::SpaceInformationPtr &space_info, const grid_map::GridMap &map_, const std::vector<Pose> Path, const std::vector<Obstacle> Obstacles)
            : base::StateValidityChecker(space_info), gridmap(map_), path(Path), obstacle_list(Obstacles) { existObstacles = true; }

        /**
         * @brief Check if state is valid
         * @param state
         * @return true if state is valid
         * @return false if state is not valid
         */
        virtual bool isValid(const base::State *state) const override;

        /**
         * @brief Check if that position is in collision
         * @param state Position as an eigen 3D vector
         * @return true if there is collision
         * @return false if there is no collision
         */
        bool checkCollision(const Eigen::Vector3d state) const;

        /**
         * @brief Check if that position is in collision
         * @param state Position as a Pose
         * @return true if there is collision
         * @return false if there is no collision
         */
        bool checkCollision(const Pose state) const;

        /**
         * @brief Check minimum distance from current state to the path
         * @return Minimum distance to the path
         */
        double clearance(const base::State *s) const;

    protected:
        const grid_map::GridMap &gridmap;
        std::vector<Obstacle> obstacle_list;
        std::vector<Pose> path;
        bool existObstacles = false;
    };

  bool PathValidityChecker::isValid(const base::State *state) const
    {
        const auto &st = (*state->as<base::OwenStateSpace::StateType>());
        Pose state_pos;
        state_pos.position(0) = st[0];
        state_pos.position(1) = st[1];
        state_pos.position(2) = st[2];

        Quaternion quat = Euler2Quaternion(0, 0, st.yaw());
        state_pos.orientation(0) = quat.x;
        state_pos.orientation(1) = quat.y;
        state_pos.orientation(2) = quat.z;
        state_pos.orientation(3) = quat.w;

        // std::cout << "Checking State [" << st[0] << ", " << st[1] << ", " << st[2] << "] yaw: " << st.yaw() << std::endl;

        return (!checkCollision(state_pos));
    }

    bool PathValidityChecker::checkCollision(const Eigen::Vector3d state) const
    {
        // Check collision with map
        grid_map::Position pos(state(0), state(1));
        // std::cout << "INFO: ValidityChecker-- Check position [" << pos(0) << ", " << pos(1) << "]\n";
        if (gridmap.isInside(pos))
        {
            if (gridmap.atPosition("ics_+", pos) >= state(2) || gridmap.atPosition("ics_-", pos) <= state(2))
                return true;
        }
        // Check collision with obstacles
        if (existObstacles)
        {
            Eigen::Matrix4d PosTmat{
                {1, 0, 0, state(0)},
                {0, 1, 0, state(1)},
                {0, 0, 1, state(2)},
                {0, 0, 0, 1}};
            for (int i = 0; i < (int)obstacle_list.size(); i++)
            {
                Eigen::Vector3d obs_position(obstacle_list[i].transformation_mat(0, 3), obstacle_list[i].transformation_mat(1, 3), obstacle_list[i].transformation_mat(2, 3));
                float distance = (obs_position - state).norm();
                if (distance <= obstacle_list[i].boundaries.norm() / 2)
                {
                    Eigen::Matrix4d Pos2Box = obstacle_list[i].transformation_mat.inverse() * PosTmat;
                    if ((abs(Pos2Box(0, 3)) <= (obstacle_list[i].boundaries(0) / 2)) &&
                        (abs(Pos2Box(1, 3)) <= (obstacle_list[i].boundaries(1) / 2)) &&
                        (abs(Pos2Box(2, 3)) <= (obstacle_list[i].boundaries(2) / 2)))
                        return true;
                }
            }
        }
        return false;
    }
}