/*
 * Copyright (C) 2023 Dominik Drexler and Simon Stahlberg
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "tyr/planning/algorithms/gbfs_lazy.hpp"

#include "tyr/common/chrono.hpp"
#include "tyr/common/segmented_vector.hpp"
#include "tyr/formalism/planning/repository.hpp"
#include "tyr/formalism/planning/views.hpp"
#include "tyr/planning/algorithms/gbfs_lazy/event_handler.hpp"
#include "tyr/planning/algorithms/openlists/alternating.hpp"
#include "tyr/planning/algorithms/strategies/goal.hpp"
#include "tyr/planning/algorithms/strategies/pruning.hpp"
#include "tyr/planning/algorithms/utils.hpp"
#include "tyr/planning/applicability.hpp"
#include "tyr/planning/ground_task.hpp"
#include "tyr/planning/ground_task/node.hpp"
#include "tyr/planning/ground_task/state_repository.hpp"
#include "tyr/planning/ground_task/state_view.hpp"
#include "tyr/planning/ground_task/successor_generator.hpp"
#include "tyr/planning/ground_task/unpacked_state.hpp"
#include "tyr/planning/heuristic.hpp"
#include "tyr/planning/lifted_task.hpp"
#include "tyr/planning/lifted_task/node.hpp"
#include "tyr/planning/lifted_task/state_repository.hpp"
#include "tyr/planning/lifted_task/state_view.hpp"
#include "tyr/planning/lifted_task/successor_generator.hpp"
#include "tyr/planning/lifted_task/unpacked_state.hpp"
#include "tyr/planning/search_node.hpp"
#include "tyr/planning/search_space.hpp"
#include "tyr/planning/state_index.hpp"

#include <algorithm>

namespace tyr::planning::gbfs_lazy
{

/**
 * GBFS search node
 */

template<typename Task>
struct SearchNode
{
    float_t g_value;
    Index<State<Task>> parent_state;
    SearchNodeStatus status;
    bool preferred;
};

static_assert(sizeof(SearchNode<LiftedTask>) == 16);
static_assert(sizeof(SearchNode<GroundTask>) == 16);

template<typename Task>
using SearchNodeVector = SegmentedVector<SearchNode<Task>>;

template<typename Task>
static SearchNode<Task>& get_or_create_search_node(Index<State<Task>> state_index, SearchNodeVector<Task>& search_nodes)
{
    static auto default_node = SearchNode { std::numeric_limits<float_t>::infinity(), Index<State<Task>>::max(), SearchNodeStatus::NEW, false };

    while (uint_t(state_index) >= search_nodes.size())
    {
        search_nodes.push_back(default_node);
    }
    return search_nodes[uint_t(state_index)];
}

/**
 * GBFS queue
 */

template<typename Task>
struct QueueEntry
{
    using KeyType = std::tuple<float_t, float_t, uint_t, SearchNodeStatus>;
    using ItemType = Index<State<Task>>;

    float_t g_value;
    float_t h_value;
    Index<State<Task>> state;
    uint_t step;
    SearchNodeStatus status;

    KeyType get_key() const { return std::make_tuple(h_value, g_value, step, status); }
    ItemType get_item() const { return state; }
};

static_assert(sizeof(QueueEntry<LiftedTask>) == 32);
static_assert(sizeof(QueueEntry<GroundTask>) == 32);

template<typename Task>
using Queue = PriorityQueue<QueueEntry<Task>>;

template<typename Task>
SearchResult<Task> find_solution(Task& task, SuccessorGenerator<Task>& successor_generator, Heuristic<Task>& heuristic, const Options<Task>& options)
{
    const auto available_threads = static_cast<uint_t>(oneapi::tbb::info::default_concurrency());

    if (options.num_threads == 0)
        throw std::invalid_argument("num_threads must be at least 1.");
    if (options.num_threads > available_threads)
        throw std::invalid_argument("Requested " + std::to_string(options.num_threads) + " threads, but only " + std::to_string(available_threads)
                                    + " are available by default.");

    auto control = oneapi::tbb::global_control(oneapi::tbb::global_control::max_allowed_parallelism, options.num_threads);

    const auto start_node = (options.start_node) ? options.start_node.value() : successor_generator.get_initial_node();
    const auto& start_state = start_node.get_state();
    const auto start_state_index = start_state.get_index();
    const auto event_handler = (options.event_handler) ? options.event_handler : DefaultEventHandler<Task>::create(0);
    const auto pruning_strategy = (options.pruning_strategy) ? options.pruning_strategy : PruningStrategy<Task>::create();
    const auto goal_strategy = (options.goal_strategy) ? options.goal_strategy : TaskGoalStrategy<Task>::create(task);
    auto rng = std::mt19937_64(options.random_seed);
    auto& state_repository = *successor_generator.get_state_repository();

    auto step = uint_t(0);
    auto result = SearchResult<Task>();
    auto search_nodes = SearchNodeVector<Task>();
    auto preferred_openlist = Queue<Task>();
    auto standard_openlist = Queue<Task>();
    auto openlist = AlternatingOpenList<Queue<Task>, Queue<Task>>(preferred_openlist, standard_openlist, std::array<size_t, 2> { 1, 1 });
    const auto start_h_value = heuristic.evaluate(start_state);
    auto best_h_value = start_h_value;
    const auto start_preferred = false;
    auto& start_search_node = get_or_create_search_node(start_state_index, search_nodes);
    start_search_node.status = (start_h_value == std::numeric_limits<float_t>::infinity()) ? SearchNodeStatus::DEAD_END : SearchNodeStatus::OPEN;
    start_search_node.g_value = start_node.get_metric();
    start_search_node.preferred = start_preferred;

    event_handler->on_start_search(start_node, start_h_value);

    /* Test static goal. */

    if (!goal_strategy->is_static_goal_satisfied())
    {
        event_handler->on_end_search();
        event_handler->on_unsolvable();

        result.status = SearchStatus::UNSOLVABLE;
        return result;
    }

    /* Test whether initial state is goal. */

    if (goal_strategy->is_dynamic_goal_satisfied(start_state))
    {
        event_handler->on_end_search();

        result.plan = Plan(start_node, LabeledNodeList<Task> {});
        result.goal_node = start_node;
        result.status = SearchStatus::SOLVED;

        event_handler->on_solved(result.plan.value());

        return result;
    }

    if (std::isnan(start_node.get_metric()))
    {
        event_handler->on_end_search();

        throw std::runtime_error("find_solution(...): start node metric value is NaN.");
    }

    /* Test whether start state is deadend. */

    if (start_search_node.status == SearchNodeStatus::DEAD_END)
    {
        event_handler->on_end_search();
        event_handler->on_unsolvable();

        result.status = SearchStatus::UNSOLVABLE;
        return result;
    }

    /* Test whether initial state should be pruned. */

    if (pruning_strategy->should_prune_state(start_state))
    {
        event_handler->on_end_search();
        event_handler->on_unsolvable();

        result.status = SearchStatus::EXHAUSTED;
        return result;
    }

    auto labeled_succ_nodes = std::vector<LabeledNode<Task>> {};

    standard_openlist.insert(QueueEntry { start_node.get_metric(), start_h_value, start_state_index, step++, start_search_node.status });

    auto stopwatch = options.max_time ? std::optional<CountdownWatch>(options.max_time.value()) : std::nullopt;

    auto& openlist_weights = openlist.get_weights();

    while (!openlist.empty())
    {
        if (stopwatch && stopwatch->has_finished())
        {
            event_handler->on_end_search();

            result.status = SearchStatus::OUT_OF_TIME;
            return result;
        }

        const auto state_index = openlist.top();
        const auto state = state_repository.get_registered_state(state_index);

        openlist.pop();
        // Weight decay of prefered queue
        openlist_weights[0] = std::max(openlist_weights[0] - 1, size_t { 1 });

        auto& search_node = get_or_create_search_node(state_index, search_nodes);
        auto node = Node<Task>(state, search_node.g_value);

        /* Close state. */

        if (search_node.status == SearchNodeStatus::CLOSED || search_node.status == SearchNodeStatus::DEAD_END)
        {
            continue;
        }

        /* Expand the successors of the node. */

        event_handler->on_expand_node(node);

        const auto state_h_value = heuristic.evaluate(state);
        if (state_h_value == std::numeric_limits<float_t>::infinity())
        {
            search_node.status = SearchNodeStatus::DEAD_END;
            continue;
        }

        if (state_h_value < best_h_value)
        {
            best_h_value = state_h_value;
            event_handler->on_new_best_h_value(best_h_value);

            // Boost prefered queue
            openlist_weights[0] += options.boost_preferred_queue;
        }

        const auto& preferred_actions = heuristic.get_preferred_actions();

        /* Ensure that the state is closed */

        search_node.status = SearchNodeStatus::CLOSED;

        successor_generator.get_labeled_successor_nodes(node, labeled_succ_nodes);

        if (options.shuffle_labeled_succ_nodes)
            std::shuffle(labeled_succ_nodes.begin(), labeled_succ_nodes.end(), rng);

        for (const auto& labeled_succ_node : labeled_succ_nodes)
        {
            const auto& succ_node = labeled_succ_node.node;
            const auto& succ_state = succ_node.get_state();
            const auto succ_state_index = succ_state.get_index();

            auto& successor_search_node = get_or_create_search_node(succ_state_index, search_nodes);

            assert(!std::isnan(succ_node.get_metric()));

            const auto is_preferred = preferred_actions.contains(labeled_succ_node.label.get_index());
            const auto is_new_successor_state = (successor_search_node.status == SearchNodeStatus::NEW);

            if (is_new_successor_state && search_nodes.size() >= options.max_num_states)
            {
                event_handler->on_end_search();

                result.status = SearchStatus::OUT_OF_STATES;
                return result;
            }

            /* Skip previously generated state. */

            if (!is_new_successor_state)
                continue;

            /* Open new state. */

            successor_search_node.status = SearchNodeStatus::OPEN;
            successor_search_node.parent_state = state_index;
            successor_search_node.g_value = succ_node.get_metric();
            successor_search_node.preferred = is_preferred;

            /* Early goal test. */

            const auto successor_is_goal_state = goal_strategy->is_dynamic_goal_satisfied(succ_state);

            if (successor_is_goal_state)
            {
                successor_search_node.status = SearchNodeStatus::GOAL;

                event_handler->on_expand_goal_node(succ_node);

                event_handler->on_end_search();

                result.plan = extract_total_ordered_plan(successor_search_node, succ_node, search_nodes, successor_generator);
                result.goal_node = succ_node;
                result.status = SearchStatus::SOLVED;

                event_handler->on_solved(result.plan.value());

                return result;
            }

            /* Apply pruning strategy */

            if (pruning_strategy->should_prune_successor_state(state, succ_state, is_new_successor_state))
            {
                event_handler->on_prune_node(succ_node);
                continue;
            }

            event_handler->on_generate_node(labeled_succ_node);

            /* Exploration strategy */

            if (is_preferred)
                preferred_openlist.insert(QueueEntry { succ_node.get_metric(), state_h_value, succ_state_index, step++, successor_search_node.status });
            else
                standard_openlist.insert(QueueEntry { succ_node.get_metric(), state_h_value, succ_state_index, step++, successor_search_node.status });
        }
    }

    event_handler->on_end_search();
    event_handler->on_exhausted();

    result.status = SearchStatus::EXHAUSTED;
    return result;
}

template SearchResult<LiftedTask> find_solution<LiftedTask>(LiftedTask& task,
                                                            SuccessorGenerator<LiftedTask>& successor_generator,
                                                            Heuristic<LiftedTask>& heuristic,
                                                            const Options<LiftedTask>& options);

template SearchResult<GroundTask> find_solution<GroundTask>(GroundTask& task,
                                                            SuccessorGenerator<GroundTask>& successor_generator,
                                                            Heuristic<GroundTask>& heuristic,
                                                            const Options<GroundTask>& options);

}
