//
// Created by itamar on 4/13/23.
//

#ifndef MANIPULATION_PLANNING_MANIPHEURISTICS_HPP
#define MANIPULATION_PLANNING_MANIPHEURISTICS_HPP

// standard includes
#include <memory>
#include <utility>
#include <vector>
#include <eigen3/Eigen/Dense>

// moveit and ros includes
#include <moveit/distance_field/propagation_distance_field.h>

// search includes
#include <search/common/intrusive_heap.h>
#include <search/planners/bfs3d.h>
#include <search/action_space/action_space.hpp>

#include <search/common/experience_graph.hpp>
#include <search/common/scene_interface.hpp>
#include <search/common/types.hpp>
#include <search/heuristics/base_heuristic.hpp>
#include <search/planners/dijkstra.hpp>

// project includes
#include <srmp_moveit_plugin/common/smpl_grid/grid.h>
#include <srmp_moveit_plugin/common/utils.hpp>
#include <manipulation_planning/manip_heuristics/manip_heuristics_base.hpp>

namespace ims {

class BFSHeuristic : public BFSHeuristicBase {
protected:
    // The full pose of the goal EE. Orientation and translation.
    Eigen::Isometry3d ee_goal_state_;

private:
    std::shared_ptr<distance_field::PropagationDistanceField> distance_field_;
    std::unique_ptr<::smpl::BFS_3D> bfs_;
    double inflation_radius_ = 0.02;
    int cost_per_cell_ = 100;
    int start_cells[3]{};

    // Flag for whether the goal is set.
    bool is_goal_set_ = false;

    struct CellCoord {
        int x, y, z;
        CellCoord() = default;
        CellCoord(int x, int y, int z) : x(x), y(y), z(z) {}
    };
    std::vector<CellCoord> m_goal_cells;

    void syncGridAndBfs() override {
        const int xc = distance_field_->getXNumCells();
        const int yc = distance_field_->getYNumCells();
        const int zc = distance_field_->getZNumCells();
        bfs_ = std::make_unique<::smpl::BFS_3D>(xc, yc, zc);
        const int cell_count = xc * yc * zc;
        int wall_count = 0;
        for (int x = 0; x < xc; ++x) {
            for (int y = 0; y < yc; ++y) {
                for (int z = 0; z < zc; ++z) {
                    const double radius = inflation_radius_;
                    if (distance_field_->getDistance(x, y, z) <= radius) {
                        bfs_->setWall(x, y, z);
                        ++wall_count;
                    }
                }
            }
        }

        std::cout << wall_count << " " << cell_count << " " << 100.0 * (double)wall_count / cell_count << " walls in the bfs heuristic" << std::endl;
    }

    int getBfsCostToGoal(const ::smpl::BFS_3D& bfs, int x, int y, int z) const override {
        if (!bfs.inBounds(x, y, z) || bfs.getDistance(x, y, z) == ::smpl::BFS_3D::WALL)
            return INF_INT;
        else {
            return cost_per_cell_ * bfs.getDistance(x, y, z);
        }
    }

    std::string tip_link_;
    std::shared_ptr<moveit::planning_interface::MoveGroupInterface> move_group;
    moveit::core::RobotModelConstPtr robot_model;
    const moveit::core::JointModelGroup* joint_model_group;
    moveit::core::RobotStatePtr kinematic_state;

public:
    BFSHeuristic() : BFSHeuristic(getDistanceFieldMoveIt()) {
    }

    explicit BFSHeuristic(std::shared_ptr<distance_field::PropagationDistanceField> distance_field,
                          const std::string& group_name = "manipulator_1") {
        distance_field_ = std::move(distance_field);
        // setup robot move group and planning scene
        move_group = std::make_shared<moveit::planning_interface::MoveGroupInterface>(group_name);
        // robot model
        robot_model = move_group->getRobotModel();
        // joint model group
        joint_model_group = robot_model->getJointModelGroup(group_name);

        kinematic_state = std::make_shared<moveit::core::RobotState>(robot_model);
        auto names = joint_model_group->getLinkModelNames();
        // get the planning group tip link
        tip_link_ = joint_model_group->getLinkModelNames().back();

        syncGridAndBfs();

        // Flag for whether the goal is set.
        is_goal_set_ = false;
    }

    void setInflationRadius(double radius) override {
        inflation_radius_ = radius;
    }

    void setCostPerCell(int cost_per_cell) override {
        cost_per_cell_ = cost_per_cell;
    }
    /**
     * @brief Set the goal configuration. This is NOT the mapped state but the actual configuration.
     * @param goal The goal is specified in configuration space.
     */
    void setGoalStateValue(const StateType& goal) override {
        // Set the goal condition to point to the passed configuration (joint values).
        goal_condition_ = std::make_shared<ims::StateEqualityCondition>(goal);

        // Compute the goal position in world space.
        kinematic_state->setJointGroupPositions(joint_model_group, goal);
        ee_goal_state_ = kinematic_state->getGlobalLinkTransform(tip_link_);
        goal_state_mapped_val = {ee_goal_state_.translation().x(),
                                 ee_goal_state_.translation().y(),
                                 ee_goal_state_.translation().z(),
                                 0, 0, 0};  // Orientation not used here.

        auto goal_position = ee_goal_state_.translation();
        int x, y, z;
        distance_field_->worldToGrid(goal_position.x(), goal_position.y(), goal_position.z(),
                                     x, y, z);
        if (!bfs_->inBounds(x, y, z))
            throw std::runtime_error("goal is out of bounds");

        m_goal_cells.emplace_back(x, y, z);

        bfs_->run(x, y, z);

        // Flag for whether the goal is set.
        is_goal_set_ = true;
    }

    /**
     * @brief Set the goal condition for the BFS3D heuristic.
     * @param goal_condition This condition should have a representative EE pose [x, y, z, roll, pitch, yaw] (m, rad).
     */
    void setGoalCondition(std::shared_ptr<Condition> goal_condition) override {
        goal_condition_ = goal_condition;
        // Variable to populate with the mapped state.
        StateType goal_state_mapped_val;
        // Check that this condition has a representative mapped state (we need it for this heuristic).
        auto goal_condition_repstate_mapped_mixin = dynamic_cast<ConditionRepresentativeStateMappedMixin*>(goal_condition.get());
        if (goal_condition_repstate_mapped_mixin != nullptr){
            goal_state_mapped_val = goal_condition_repstate_mapped_mixin->getRepresentativeStateMapped();
        }
        else {
            auto goal_condition_repstate_mixin = dynamic_cast<ConditionRepresentativeStateMixin *>(goal_condition.get());
            if (!goal_condition_repstate_mixin) {
                throw std::runtime_error("Goal condition does not have a representative state or mapped state mixin.");
            }
            // Get the representative state.
            StateType goal_state_val = goal_condition_repstate_mixin->getRepresentativeState();
            // Map the state to the EE pose (mapped state).
            kinematic_state->setJointGroupPositions(joint_model_group, goal_state_val);
            auto ee_goal_state = kinematic_state->getGlobalLinkTransform(tip_link_);
            goal_state_mapped_val = {ee_goal_state.translation().x(),
                                     ee_goal_state.translation().y(),
                                     ee_goal_state.translation().z(),
                                     0, 0, 0};  // Orientation not used here.
        }

        // Convert to an Eigen isometry3d.
        ee_goal_state_.translation().x() = goal_state_mapped_val.at(0);
        ee_goal_state_.translation().y() = goal_state_mapped_val.at(1);
        ee_goal_state_.translation().z() = goal_state_mapped_val.at(2); // Orientation not used here.

        // Run the BFS.
        auto goal_position = ee_goal_state_.translation();
        int x, y, z;
        distance_field_->worldToGrid(goal_position.x(), goal_position.y(), goal_position.z(),
                                     x, y, z);
        if (!bfs_->inBounds(x, y, z))
            throw std::runtime_error("goal is out of bounds");

        m_goal_cells.emplace_back(x, y, z);

        bfs_->run(x, y, z);

        // Flag for whether the goal is set.
        is_goal_set_ = true;
    }

    void setStart(const StateType& start) override {
        start_ = start;
        kinematic_state->setJointGroupPositions(joint_model_group, start);
        auto ee_start_state = kinematic_state->getGlobalLinkTransform(tip_link_);

        auto start_position = ee_start_state.translation();
        distance_field_->worldToGrid(start_position.x(), start_position.y(), start_position.z(),
                                     start_cells[0], start_cells[1], start_cells[2]);
        //            if (!bfs_->inBounds(x, y, z))
        //                throw std::runtime_error("start is out of bounds");
    }

    /// @brief Get the metric distance to the goal state
    /// @param x The x position of the state
    /// @param y The y position of the state
    /// @param z The z position of the state
    /// @return The distance to the goal state
    double getMetricGoalDistance(double x, double y, double z) const override {
        int gx, gy, gz;
        distance_field_->worldToGrid(x, y, z, gx, gy, gz);
        if (!bfs_->inBounds(gx, gy, gz))
            return (double)::smpl::BFS_3D::WALL * distance_field_->getResolution();
        else
            return (double)bfs_->getDistance(gx, gy, gz) * distance_field_->getResolution();
    }

    /**
     * @brief Get the heuristic distance between a mapped state and the goal.
     * @param state_mapped_val a mapped state of form [x, y, z, r, p, y] (m, rads).
     * @param dist The heuristic distance.
     * @return
     */
    bool getHeuristicMapped(const StateType& state_mapped_val, double& dist) override {
        if (m_goal_cells.empty() || !is_goal_set_) {
            std::cout << "Goal is not set in BFS heuristic! It must be set prior to calling getHeuristic" << std::endl;
            return false;
        }
        dist = getBfsCostToGoal(*bfs_, state_mapped_val.at(0), state_mapped_val.at(1), state_mapped_val.at(2));
        return true;
    }

    /// @brief Get the metric distance to the start state
    /// @param x The x position of the state
    /// @param y The y position of the state
    /// @param z The z position of the state
    /// @return The distance to the start state
    double getMetricStartDistance(double x, double y, double z) override {
        int sx, sy, sz;
        distance_field_->worldToGrid(x, y, z, sx, sy, sz);
        // manhattan distance
        return (std::abs(sx - start_cells[0]) + std::abs(sy - start_cells[1]) + std::abs(sz - start_cells[2])) * distance_field_->getResolution();
    }

    bool getHeuristic(const StateType& s1, const StateType& s2, double& dist) override {
        // check if s2 is a goal state
        kinematic_state->setJointGroupPositions(joint_model_group, s2);
        auto ee_check_state = kinematic_state->getGlobalLinkTransform(tip_link_);

        auto check_position = ee_check_state.translation();
        int x, y, z;
        distance_field_->worldToGrid(check_position.x(), check_position.y(), check_position.z(),
                                     x, y, z);
        // check if s2 is a goal state
        for (const auto& goal_cell : m_goal_cells) {
            if (goal_cell.x == x && goal_cell.y == y && goal_cell.z == z) {
                return getHeuristic(s1, dist);
            }
        }

        std::cout << "getHeuristic between two random states with BFS is not supported!" << std::endl;
        return false;
    }

    bool getHeuristic(const StateType& s, double& dist) override {
        if (m_goal_cells.empty() || !is_goal_set_) {
            std::cout << "Goal is not set in BFS heuristic! It must be set prior to calling getHeuristic" << std::endl;
            return false;
        }

        kinematic_state->setJointGroupPositions(joint_model_group, s);
        auto pose_ee = kinematic_state->getGlobalLinkTransform(tip_link_);

        auto s_position = pose_ee.translation();
        int x, y, z;
        distance_field_->worldToGrid(s_position.x(), s_position.y(), s_position.z(),
                                     x, y, z);

        dist = getBfsCostToGoal(*bfs_, x, y, z);

        return true;
    }

    Eigen::Isometry3d getGoalPoseEE() const override {
        return ee_goal_state_;
    }

    StateType getGoalPoseEExyzrpy() const override {
        // Get the goal pose in a vector containing the xyz (meters) and rpy (radians) of the goal pose.
        StateType ee_goal_state_xyzrpy(6);
        ee_goal_state_xyzrpy[0] = ee_goal_state_.translation().x();
        ee_goal_state_xyzrpy[1] = ee_goal_state_.translation().y();
        ee_goal_state_xyzrpy[2] = ee_goal_state_.translation().z();
        // The stored Isometry3d uses zyx euler angles.
        ee_goal_state_xyzrpy[3] = ee_goal_state_.rotation().eulerAngles(2, 1, 0).z();
        ee_goal_state_xyzrpy[4] = ee_goal_state_.rotation().eulerAngles(2, 1, 0).y();
        ee_goal_state_xyzrpy[5] = ee_goal_state_.rotation().eulerAngles(2, 1, 0).x();
        
        return ee_goal_state_xyzrpy;
    }

    void setDistanceField(std::shared_ptr<distance_field::PropagationDistanceField> distance_field) {
        distance_field_ = std::move(distance_field);
        syncGridAndBfs();
    }
};

/// @brief Same as the BFS Heuristic, but with states that include time in thier last element of the state vector. The methods here remove the time component and call the BFS heuristic methods as normal.
class BFSRemoveTimeHeuristic : public BFSHeuristic {
public:
    BFSRemoveTimeHeuristic() : BFSHeuristic() {
    }

    explicit BFSRemoveTimeHeuristic(std::shared_ptr<distance_field::PropagationDistanceField> distance_field,
                                    const std::string& group_name = "panda0_arm") : BFSHeuristic(distance_field, group_name) {
    }

    void setGoalStateValue(const StateType& goal) override {
        StateType goal_wo_time = goal;
        goal_wo_time.pop_back();
        BFSHeuristic::setGoalStateValue(goal_wo_time);
        goal_condition_ = std::make_shared<ims::StateEqualityCondition>(goal_wo_time);
    }

    void setStart(const StateType& start) override {
        StateType start_wo_time = start;
        start_wo_time.pop_back();
        BFSHeuristic::setStart(start_wo_time);
    }

    bool getHeuristic(const StateType& s1, const StateType& s2, double& dist) override {
        StateType s1_wo_time = s1;
        s1_wo_time.pop_back();
        StateType s2_wo_time = s2;
        s2_wo_time.pop_back();
        return BFSHeuristic::getHeuristic(s1_wo_time, s2_wo_time, dist);
    }

    bool getHeuristic(const StateType& s, double& dist) override {
        StateType s_wo_time = s;
        s_wo_time.pop_back();
        return BFSHeuristic::getHeuristic(s_wo_time, dist);
    }
};
}  // namespace ims

#endif  // MANIPULATION_PLANNING_MANIPHEURISTICS_HPP
