/*
 * Copyright (C) 2025, Itamar Mishani, Yorai Shaoul
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *     * Neither the name of the Carnegie Mellon University nor the names of its
 *       contributors may be used to endorse or promote products derived from
 *       this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */
/*!
 * \file   factories.cpp
 * \author Itamar Mishani (imishani@cmu.edu), Yoarai Shaoul (yorai@cmu.edu)
 * \date   2025-02-20
*/


#include <srmp_moveit_plugin/heuristics/manip_heuristics.hpp>
#include <srmp_moveit_plugin/common/moveit_distance_field_wrapper.hpp>
#include <srmp_moveit_plugin/common/factories.hpp>
#include <manipulation_planning/action_space/manipulation_action_space.hpp>
#include <manipulation_planning/action_space/manipulation_edge_action_space.hpp>
#include <manipulation_planning/action_space/mramp_action_space.hpp>
#include <manipulation_planning/action_space/manipulation_experience_accelerated_constrained_action_space.hpp>


namespace ims {
PlanningWorld::PlanningWorld() {
    // Set the single and multi agent caches to nullptr.
    single_agent_cache_ = nullptr;
    multi_agent_cache_ = nullptr;
}

void PlanningWorld::setPlanningSceneMoveIt(const planning_scene::PlanningSceneConstPtr& planning_scene) {
    planning_scene_moveit_ = planning_scene;
}

void PlanningWorld::makePlanner(const std::vector<std::string> & articulation_names,
                                         const std::map<std::string, ims::PlannerContextValueType>  & planner_context) {
    // Get planner name.
    std::string planner_name = "";
    getValueFromVariantMap(planner_context, "planner_id", planner_name);
    // Check if this should be a single or multi-robot planner.
    bool is_multi_agent = articulation_names.size() > 1;
    if (!is_multi_agent) {
        // Get the planner name default if needed.
        planner_name = planner_name.empty() ? "wAstar" : planner_name;

        if (registered_single_agent_planner_factories.find(planner_name) == registered_single_agent_planner_factories.end()) {
            std::cout << RED << "PlanningWorld::makePlanner: Planner name " << planner_name << " is not valid." << RESET << std::endl;
            return;
        }

        // Create scene interface.
        std::string articulation_name = articulation_names.front();
        planning_scene::PlanningScenePtr planning_scene_moveit = planning_scene_moveit_->diff();
        std::shared_ptr<MoveitInterface> scene_interface = std::make_shared<MoveitInterface>(articulation_name,
                                                                                             planning_scene_moveit);
        auto planner_factory_single_agent = registered_single_agent_planner_factories.find(planner_name)->second;
        planner_factory_single_agent(articulation_name,
                                     planner_context,
                                     scene_interface,
                                     single_agent_cache_);
    }
    else{
        // Get the planner name default if needed.
        planner_name = planner_name.empty() ? "xECBS" : planner_name;
        if (registered_multi_agent_planner_factories.find(planner_name) == registered_multi_agent_planner_factories.end()) {
            ROS_ERROR_STREAM("PlanningWorld::makePlanner: Planner name " << planner_name << " is not valid for multi-agent planning.");
            return;
        }

        // Create scene interface.
        std::vector<std::shared_ptr<MoveitInterface>> scene_interfaces;
        for (const auto &articulation_name : articulation_names) {
            planning_scene::PlanningScenePtr planning_scene_moveit = planning_scene_moveit_->diff();
            std::shared_ptr<MoveitInterface> scene_interface = std::make_shared<MoveitInterface>(articulation_name,
                                                                                                 planning_scene_moveit);
            scene_interfaces.push_back(scene_interface);
        }

        auto planner_factory_multi_agent = registered_multi_agent_planner_factories.find(planner_name)->second;
        planner_factory_multi_agent(articulation_names,
                                    planner_context,
                                    scene_interfaces,
                                    multi_agent_cache_);
    }
}

PathType PlanningWorld::plan(const StateType& start_state, const StateType& goal_state) {
    // Discretize and round the start and goal states.
    StateType start_state_dis = start_state;
    StateType goal_state_dis = goal_state;

    // Get the discretization from the action type.
    if (single_agent_cache_->planner_name == "wPASE") {
        std::shared_ptr <ims::ManipulationEdgeActionSpace> action_space =
                std::dynamic_pointer_cast<ims::ManipulationEdgeActionSpace>(single_agent_cache_->action_space);
        if (!action_space) {
            throw std::runtime_error("PlanningWorld::plan: Action space is not of type ManipulationEdgeActionSpace.");
        }
        std::static_pointer_cast<MoveitInterface>(action_space->getSceneInterface())->planning_scene_ = planning_scene_moveit_->diff();  // Important.

        std::vector <std::pair<double, double>> joint_limits;
        action_space->getSceneInterface()->getJointLimits(joint_limits);
        ims::normalizeAngles(start_state_dis, joint_limits);
        ims::normalizeAngles(goal_state_dis, joint_limits);
        StateType state_discretization = action_space->getStateDiscretization();
        roundStateToDiscretization(start_state_dis, state_discretization);
        roundStateToDiscretization(goal_state_dis, state_discretization);

        // Set this goal state in the snap heuristic of the action space.
        std::cout << BOLDGREEN << "PlanningWorld::plan: Setting new goal for BFS heuristic..." << RESET << std::endl;
        std::chrono::steady_clock::time_point start_time = std::chrono::steady_clock::now();
        // Option I: Do this by straight up creating a new BFS heuristic and setting the goal state.
//    std::shared_ptr<distance_field::PropagationDistanceField> df = action_space->getSceneInterface()->getDistanceFieldMoveItInterface();
//    ims::BFSHeuristic* snap_heuristic = new ims::BFSHeuristic(df, single_agent_cache_->group_name);
//    snap_heuristic->setGoalStateValue(goal_state_dis);
//    action_space->setBFSHeuristic(snap_heuristic);

// Option II: Do this by setting a new goal for the BFS heuristic. This may not capture changes in the distance field.
        // PUT BACK IN.
//    std::shared_ptr<distance_field::PropagationDistanceField> df = action_space->getSceneInterface()->getDistanceFieldMoveItInterface(); // ~50 ms.
//    action_space->getBFSHeuristic()->setDistanceField(df);
        action_space->getBFSHeuristic()->setGoalStateValue(goal_state_dis); // ~50 ms.

        std::chrono::steady_clock::time_point end_time = std::chrono::steady_clock::now();

        std::cout << BOLDGREEN << "====== OK: New SDF set. Took "
                  << std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count() << " ms."
                  << RESET << std::endl;
        // Initialize the planner.
        action_space->resetPlanningData();
        single_agent_cache_->planner->resetPlanningData();


        // Some planners need special
        single_agent_cache_->planner->initializePlanner(
                std::static_pointer_cast<ims::EdgeActionSpace>(action_space), start_state_dis, goal_state_dis);

    }
    else {
        std::shared_ptr <ims::ManipulationActionSpace> action_space =
                std::dynamic_pointer_cast<ims::ManipulationActionSpace>(single_agent_cache_->action_space);
        std::shared_ptr <ims::ManipulationEdgeActionSpace> edge_action_space =
                std::dynamic_pointer_cast<ims::ManipulationEdgeActionSpace>(single_agent_cache_->action_space);
        if (!action_space) {
            throw std::runtime_error("PlanningWorld::plan: Action space is not of type ManipulationActionSpace.");
        }
        std::static_pointer_cast<MoveitInterface>(action_space->getSceneInterface())->planning_scene_ = planning_scene_moveit_->diff();  // Important.

        std::vector <std::pair<double, double>> joint_limits;
        action_space->getSceneInterface()->getJointLimits(joint_limits);
        ims::normalizeAngles(start_state_dis, joint_limits);
        ims::normalizeAngles(goal_state_dis, joint_limits);
        StateType state_discretization = action_space->getStateDiscretization();
        roundStateToDiscretization(start_state_dis, state_discretization);
        roundStateToDiscretization(goal_state_dis, state_discretization);

        // Set this goal state in the snap heuristic of the action space.
        std::cout << BOLDGREEN << "PlanningWorld::plan: Setting new goal for BFS heuristic..." << RESET << std::endl;
        std::chrono::steady_clock::time_point start_time = std::chrono::steady_clock::now();
        // Option I: Do this by straight up creating a new BFS heuristic and setting the goal state.
//    std::shared_ptr<distance_field::PropagationDistanceField> df = action_space->getSceneInterface()->getDistanceFieldMoveItInterface();
//    ims::BFSHeuristic* snap_heuristic = new ims::BFSHeuristic(df, single_agent_cache_->group_name);
//    snap_heuristic->setGoalStateValue(goal_state_dis);
//    action_space->setBFSHeuristic(snap_heuristic);

// Option II: Do this by setting a new goal for the BFS heuristic. This may not capture changes in the distance field.
        // PUT BACK IN.
//    std::shared_ptr<distance_field::PropagationDistanceField> df = action_space->getSceneInterface()->getDistanceFieldMoveItInterface(); // ~50 ms.
//    action_space->getBFSHeuristic()->setDistanceField(df);
        action_space->getBFSHeuristic()->setGoalStateValue(goal_state_dis); // ~50 ms.

        std::chrono::steady_clock::time_point end_time = std::chrono::steady_clock::now();

        std::cout << BOLDGREEN << "====== OK: New SDF set. Took "
                  << std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count() << " ms."
                  << RESET << std::endl;
        // Initialize the planner.
        action_space->resetPlanningData();
        single_agent_cache_->planner->resetPlanningData();
        single_agent_cache_->planner->initializePlanner(single_agent_cache_->action_space, start_state_dis,
                                                        goal_state_dis);
    }

    // Plan.
    PathType path;
    if (single_agent_cache_->planner->plan(path)) {

        // Process the path.
        // If the path is empty, return true and do nothing.
        std::cout << "E"  << std::endl;
        if (path.empty()) {
            std::cout << "F"  << std::endl;
            return path;
        }
        std::cout << "G path length " << path.size() << std::endl;
        // Replace the first state of the path with the actual start state.
        path.front() = start_state;
        // Replace the last state of the path with the actual goal state.
        path.back() = goal_state;
        // TODO(yoraish): can add our own profiling and shortcutting here.

        std::cout << "G2 path length " << path.size() << std::endl;

        return path;
    }
    else {
        std::cout << RED << "PlanningWorld::plan: Planning failed." << RESET << std::endl;
    }

    return path;
}

std::vector<PathType> PlanningWorld::plan(const std::vector<StateType>& start_states, const std::vector<StateType>& goal_states) {
    // Discretize and round the start and goal states.
    std::vector<StateType> start_states_dis = start_states;
    std::vector<StateType> goal_states_dis = goal_states;

    // Casted action spaces.
    std::vector<std::shared_ptr<ims::MrampManipulationActionSpace>> mramp_action_spaces;
    std::vector<std::shared_ptr<ims::SubcostConstrainedActionSpace>> sc_action_spaces;
    std::vector<std::shared_ptr<ims::ConstrainedActionSpace>> c_action_spaces;

    for (size_t i = 0; i < start_states.size(); i++) {
        // Get the discretization from the action type.
        std::shared_ptr<ims::MrampManipulationActionSpace> action_space =
                std::dynamic_pointer_cast<ims::MrampManipulationActionSpace>(multi_agent_cache_->action_spaces.at(i));
        if (!action_space) {
            throw std::runtime_error("PlanningWorld::plan: Action space is not of type MrampManipulationActionSpace.");
        }
        mramp_action_spaces.push_back(action_space);
        sc_action_spaces.push_back(action_space);
        c_action_spaces.push_back(action_space);
        std::static_pointer_cast<MoveitInterface>(action_space->getSceneInterface())->planning_scene_ = planning_scene_moveit_->diff();  // Important.


        std::vector<std::pair<double, double>> joint_limits;
        action_space->getSceneInterface()->getJointLimits(joint_limits);
        std::vector<bool> agent_state_angle_valid_mask(joint_limits.size(), true);
        agent_state_angle_valid_mask.push_back(false);
        ims::normalizeAngles(start_states_dis.at(i), joint_limits, agent_state_angle_valid_mask);
        ims::normalizeAngles(goal_states_dis.at(i), joint_limits, agent_state_angle_valid_mask);
        StateType state_discretization = action_space->getStateDiscretization();

        roundStateToDiscretization(start_states_dis.at(i), state_discretization);
        roundStateToDiscretization(goal_states_dis.at(i), state_discretization);
    }

    // Reset planning data in the action spaces and planner.
    multi_agent_cache_->planner->resetPlanningData();
    for (size_t i = 0; i < start_states.size(); i++) {
        mramp_action_spaces.at(i)->resetPlanningData();
        mramp_action_spaces.at(i)->resetPlanningDataAndExperiences();
    }

    // Set this goal state in the snap heuristic of the action spaces.
    std::chrono::steady_clock::time_point start_time_sdf = std::chrono::steady_clock::now();
    for (size_t i = 0; i < start_states.size(); i++) {
        // PUT BACK IN.
        auto df_wrapper = mramp_action_spaces.at(i)->getSceneInterface()->getDistanceFieldMoveItInterface();
        auto df = std::static_pointer_cast<MoveItDistanceFieldWrapper>(df_wrapper)->getMoveItDistanceFieldShared();
        std::static_pointer_cast<BFSHeuristic>(std::shared_ptr<BFSHeuristicBase>(mramp_action_spaces.at(i)->getBFSHeuristic(), [](BFSHeuristicBase*){}))->setDistanceField(df);
        mramp_action_spaces.at(i)->getBFSHeuristic()->setGoalStateValue(goal_states_dis.at(i)); // ~50 ms.
    }
    std::chrono::steady_clock::time_point end_time_sdf = std::chrono::steady_clock::now();
    std::cout << BOLDGREEN << "====== OK: New SDFs set. Took " << std::chrono::duration_cast<std::chrono::milliseconds>(end_time_sdf - start_time_sdf).count() << " ms." << RESET << std::endl;


    // Cast the planner to the right type.
    // HERE HERE HERE. Add base class initialize planner so that casting would not be needed.
    // Hack for now.
    std::string planner_name = multi_agent_cache_->planner_name;
    ims::MultiAgentPaths paths;
    bool is_plan_successful = false;
    if (multi_agent_cache_->planner_name == "ECBS") {
        std::shared_ptr<ims::ECBS> multi_agent_planner = std::dynamic_pointer_cast<ims::ECBS>(multi_agent_cache_->planner);
        if (!multi_agent_planner) {
            throw std::runtime_error("PlanningWorld::plan: Planner is not of type ECBS despite claiming so.");
        }
        // Initialize the planner.
        multi_agent_planner->initializePlanner(sc_action_spaces, multi_agent_cache_->agent_group_names, start_states_dis, goal_states_dis);
        is_plan_successful = multi_agent_planner->plan(paths);
    }
    else if (multi_agent_cache_->planner_name == "xECBS") {
        std::shared_ptr<ims::EAECBS> multi_agent_planner = std::dynamic_pointer_cast<ims::EAECBS>(multi_agent_cache_->planner);
        if (!multi_agent_planner) {
            throw std::runtime_error("PlanningWorld::plan: Planner is not of type xECBS despite claiming so.");
        }
        // Initialize the planner.
        multi_agent_planner->initializePlanner(sc_action_spaces, multi_agent_cache_->agent_group_names, start_states_dis, goal_states_dis);
        is_plan_successful = multi_agent_planner->plan(paths);
    }
    else if (multi_agent_cache_->planner_name == "PP") {
        std::shared_ptr<ims::PrioritizedPlanning> multi_agent_planner = std::dynamic_pointer_cast<ims::PrioritizedPlanning>(multi_agent_cache_->planner);
        if (!multi_agent_planner) {
            throw std::runtime_error("PlanningWorld::plan: Planner is not of type PP despite claiming so.");
        }
        // Initialize the planner.
        multi_agent_planner->initializePlanner(c_action_spaces, multi_agent_cache_->agent_group_names, start_states_dis, goal_states_dis);
        is_plan_successful = multi_agent_planner->plan(paths);
    }
    else {
            throw std::runtime_error("PlanningWorld::plan: Planner is not of a known type.");
    }

    // Plan.
    if (is_plan_successful) {
        // Process the paths.
        // Set the starts and goals to the actual starts and goals (not discretized).
        for (size_t i = 0; i < start_states.size(); i++) {
            paths.at(i).front() = start_states.at(i);
            paths.at(i).back() = goal_states.at(i);
            // Fix the start and goal time.
            paths.at(i).front().back() = 0;
            paths.at(i).back().back() = paths.at(i).size() - 1;
        }

        // TODO(yoraish): can add our own profiling and shortcutting here.

        // Convert to the right format.
        std::vector<PathType> paths_out(paths.size());
        for (size_t i = 0; i < paths.size(); i++) {
            paths_out.at(i) = paths.at(i);
        }

        return paths_out;
    }
    else {
        std::cout << RED << "PlanningWorld::plan: Planning failed." << RESET << std::endl;
    }

    return {};

}


// ====================
// Factories.
// ====================


} // namespace ims