/* Copyright (c) AIRBUS and its affiliates.
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 * This is the scikit-decide implementation of MCTS and UCT from
 * "A Survey of Monte Carlo Tree Search Methods" by Browne et al
 * (IEEE Transactions on Computational Intelligence  and AI in games,
 * 2012). We additionnally implement a heuristic value estimate as in
 * "Monte-Carlo tree search and rapid action value estimation in
 * computer Go" by Gelly and Silver (Artificial Intelligence, 2011)
 * except that the heuristic estimate is called on states but not
 * on state-action pairs to be more in line with heuristic search
 * algorithms in the literature and other implementations of
 * heuristic search algorithms in scikit-decide.
 */
#ifndef SKDECIDE_MCTS_HH
#define SKDECIDE_MCTS_HH

#include <functional>
#include <memory>
#include <unordered_set>
#include <unordered_map>
#include <vector>
#include <queue>
#include <list>
#include <chrono>
#include <random>

#include <boost/range/irange.hpp>

#include "spdlog/spdlog.h"
#include "spdlog/sinks/stdout_color_sinks.h"

#include "utils/associative_container_deducer.hh"
#include "utils/string_converter.hh"
#include "utils/execution.hh"

namespace skdecide {

/** Use Environment domain knowledge for transitions */
template <typename Tsolver>
struct StepTransitionMode {
    void init_rollout(Tsolver& solver, const std::size_t* thread_id) const {
        solver.domain().reset(thread_id);
        std::for_each(solver.action_prefix().begin(), solver.action_prefix().end(),
                      [&solver, &thread_id](const typename Tsolver::Domain::Event& a){solver.domain().step(a, thread_id);});
    }

    typename Tsolver::Domain::TransitionOutcome random_next_outcome(
            Tsolver& solver,
            const std::size_t* thread_id,
            const typename Tsolver::Domain::State& state,
            const typename Tsolver::Domain::Event& action) const {
        return solver.domain().step(action, thread_id);
    }

    typename Tsolver::StateNode* random_next_node(
            Tsolver& solver,
            const std::size_t* thread_id,
            typename Tsolver::ActionNode& action) const {
        auto outcome = solver.domain().step(action.action, thread_id);
        typename Tsolver::StateNode* n = nullptr;

        solver.execution_policy().protect([&n, &solver, &outcome](){
            auto si = solver.graph().find(typename Tsolver::StateNode(outcome.state()));
            if (si != solver.graph().end()) {
                // we won't change the real key (ActionNode::action) so we are safe
                n = &const_cast<typename Tsolver::StateNode&>(*si);
            }
        });

        return n;
    }
};

/** Use Simulation domain knowledge for transitions */
template <typename Tsolver>
struct SampleTransitionMode {
    void init_rollout(Tsolver& solver, const std::size_t* thread_id) const {}

    typename Tsolver::Domain::TransitionOutcome random_next_outcome(
            Tsolver& solver,
            const std::size_t* thread_id,
            const typename Tsolver::Domain::State& state,
            const typename Tsolver::Domain::Event& action) const {
        return solver.domain().sample(state, action, thread_id);
    }

    typename Tsolver::StateNode* random_next_node(
            Tsolver& solver,
            const std::size_t* thread_id,
            typename Tsolver::ActionNode& action) const {
        typename Tsolver::StateNode* n = nullptr;

        solver.execution_policy().protect([&n, &action, &solver](){
            solver.execution_policy().protect([&n, &action, &solver](){
                n = action.dist_to_outcome[action.dist(solver.gen())]->first;
            }, solver.gen_mutex());
        }, action.parent->mutex);

        return n;
    }
};

/** Use uncertain transitions domain knowledge for transitions */
template <typename Tsolver>
struct DistributionTransitionMode {
    void init_rollout(Tsolver& solver, const std::size_t* thread_id) const {}

    typename Tsolver::Domain::TransitionOutcome random_next_outcome(
            Tsolver& solver,
            const std::size_t* thread_id,
            const typename Tsolver::Domain::State& state,
            const typename Tsolver::Domain::Event& action) const {
        return solver.domain().sample(state, action, thread_id);
    }

    typename Tsolver::StateNode* random_next_node(
            Tsolver& solver,
            const std::size_t* thread_id,
            typename Tsolver::ActionNode& action) const {
        typename Tsolver::StateNode* n = nullptr;

        solver.execution_policy().protect([&n, &action, &solver](){
            solver.execution_policy().protect([&n, &action, &solver](){
                n = action.dist_to_outcome[action.dist(solver.gen())]->first;
            }, solver.gen_mutex());
        }, action.parent->mutex);

        return n;
    }
};

/** Default tree policy as used in UCT */
template <typename Tsolver>
class DefaultTreePolicy {
public :
    typename Tsolver::StateNode* operator()(Tsolver& solver,
                                            const std::size_t* thread_id, // for parallelisation
                                            const typename Tsolver::Expander& expander,
                                            const typename Tsolver::ActionSelectorOptimization& action_selector,
                                            typename Tsolver::StateNode& n,
                                            std::size_t& d) const {
        try {
            if (solver.debug_logs()) {
                solver.execution_policy().protect([&n](){
                    spdlog::debug("Launching default tree policy from state " + n.state.print() +
                                  Tsolver::ExecutionPolicy::print_thread());
                }, n.mutex);
            }

            solver.transition_mode().init_rollout(solver, thread_id);
            typename Tsolver::StateNode* current_node = &n;

            while(!(current_node->terminal) && d < solver.max_depth()) {
                typename Tsolver::StateNode* next_node = expander(solver, thread_id, *current_node);
                d++;

                if (next_node == nullptr) { // node fully expanded
                    typename Tsolver::ActionNode* action = action_selector(solver, thread_id, *current_node);

                    if (action == nullptr) {
                        // It might happen in parallel execution mode when the current node's actions are all being
                        // expanded by concurrent threads that claim the node is expanded but not yet backpropagated
                        // and a new thread meantime comes and sees the node as expanded, thus all action visits counts
                        // are still equal to zero (implying action_selector to return nullptr).
                        // This shall NOT happen in sequential execution mode.
                        break;
                    } else {
                        next_node = solver.transition_mode().random_next_node(solver, thread_id, *action);

                        if (next_node == nullptr) { // might happen with step transition mode and stochastic environments
                            break;
                        } else {
                            current_node = next_node;
                        }
                    }
                } else {
                    current_node = next_node;
                    break;
                }
            }

            return current_node;
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&n, &e](){
                spdlog::error("SKDECIDE exception in MCTS when simulating the tree policy from state " + n.state.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, n.mutex);
            throw;
        }
    }
};

/** Test if a given node needs to be expanded by assuming that applicable actions and next states
 *  can be enumerated. Returns nullptr if all actions and outcomes have already been tried, otherwise a
 *  sampled unvisited outcome according to its probability (among only unvisited outcomes).
 *  REQUIREMENTS: returns nullptr if all actions have been already tried, and set the terminal
 *  flag of the returned next state
 */
template <typename Tsolver>
class FullExpand {
public :
    typedef std::function<std::pair<double, std::size_t>
                          (typename Tsolver::Domain&, const typename Tsolver::Domain::State&, const std::size_t*)> HeuristicFunctor;

    FullExpand(const HeuristicFunctor& heuristic = [](typename Tsolver::Domain& domain,
                                                      const typename Tsolver::Domain::State& state,
                                                      const std::size_t* thread_id) {
        return std::make_pair(0.0, 0);
    })
    : _heuristic(heuristic), _checked_transition_mode(false) {}

    FullExpand(const FullExpand& other)
    : _heuristic(other._heuristic), _checked_transition_mode(false) {}

    typename Tsolver::StateNode* operator()(Tsolver& solver,
                                            const std::size_t* thread_id,
                                            typename Tsolver::StateNode& n) const {
        try {
            if (solver.debug_logs()) {
                solver.execution_policy().protect([&n](){
                    spdlog::debug("Testing expansion of state " + n.state.print() +
                                  Tsolver::ExecutionPolicy::print_thread());
                }, n.mutex);
            }

            if (n.expanded) {
                if (solver.debug_logs()) { spdlog::debug("State already fully expanded" +
                                                         Tsolver::ExecutionPolicy::print_thread()); }
                return nullptr;
            }

            // Generate applicable actions if not already done
            solver.execution_policy().protect([&n, &solver, &thread_id](){
                if (n.actions.empty()) {
                    if (solver.debug_logs()) { spdlog::debug("State never expanded, generating all next actions" +
                                                             Tsolver::ExecutionPolicy::print_thread()); }
                    auto applicable_actions = solver.domain().get_applicable_actions(n.state, thread_id).get_elements();

                    for (const auto& a : applicable_actions) {
                        auto i = n.actions.emplace(typename Tsolver::ActionNode(a));

                        if (i.second) {
                            // we won't change the real key (ActionNode::action) so we are safe
                            const_cast<typename Tsolver::ActionNode&>(*i.first).parent = &n;
                        }
                    }
                }
            }, n.mutex);

            // Check for untried outcomes
            if (solver.debug_logs()) { spdlog::debug("Checking for untried outcomes" +
                                                     Tsolver::ExecutionPolicy::print_thread()); }
            std::vector<std::pair<typename Tsolver::ActionNode*, typename Tsolver::StateNode*>> untried_outcomes;
            std::vector<double> weights;

            solver.execution_policy().protect([&n, &untried_outcomes, &weights](){
                for (auto& a : n.actions) {
                    // we won't change the real key (ActionNode::action) so we are safe
                    typename Tsolver::ActionNode& ca = const_cast<typename Tsolver::ActionNode&>(a);

                    if (a.outcomes.empty()) {
                        // we won't change the real key (ActionNode::action) so we are safe
                        untried_outcomes.push_back(std::make_pair(&ca, nullptr));
                        weights.push_back(1.0);
                    } else {
                        // Check if there are next states that have been never visited
                        std::vector<double> probs = a.dist.probabilities();

                        for (std::size_t p = 0 ; p < probs.size() ; p++) {
                            typename Tsolver::StateNode* on = ca.dist_to_outcome[p]->first;

                            if (on->visits_count == 0) {
                                untried_outcomes.push_back(std::make_pair(&ca, on));
                                weights.push_back(probs[p]);
                            }
                        }
                    }
                }
            }, n.mutex);

            if (untried_outcomes.empty()) { // nothing to expand
                if (solver.debug_logs()) { spdlog::debug("All outcomes already tried" +
                                                         Tsolver::ExecutionPolicy::print_thread()); }
                n.expanded = true;
                return nullptr;
            } else {
                std::discrete_distribution<> odist(weights.begin(), weights.end());
                std::size_t outcome_id = 0;
                solver.execution_policy().protect([&outcome_id, &odist, &solver](){
                    outcome_id = odist(solver.gen());
                }, solver.gen_mutex());
                auto& uo = untried_outcomes[outcome_id];

                if (uo.second == nullptr) { // unexpanded action
                    if (solver.debug_logs()) { spdlog::debug("Found one unexpanded action: " + uo.first->action.print() +
                                                             Tsolver::ExecutionPolicy::print_thread()); }
                    return expand_action(solver, thread_id, solver.transition_mode(), n, *(uo.first));
                } else { // expanded action, just return the selected next state
                    if (solver.debug_logs()) {
                        solver.execution_policy().protect([&uo](){
                            spdlog::debug("Found one untried outcome: action " + uo.first->action.print() +
                                          " and next state " + uo.second->state.print() +
                                          Tsolver::ExecutionPolicy::print_thread());
                        }, uo.second->mutex);
                    }
                    return uo.second;
                }
            }
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&n, &e](){
                spdlog::error("SKDECIDE exception in MCTS when expanding state " + n.state.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, n.mutex);
            throw;
        }
    }

    template <typename Ttransition_mode,
              std::enable_if_t<std::is_same<Ttransition_mode, DistributionTransitionMode<Tsolver>>::value, int> = 0>
    typename Tsolver::StateNode* expand_action(Tsolver& solver,
                                               const std::size_t* thread_id,
                                               const Ttransition_mode& transition_mode,
                                               typename Tsolver::StateNode& state,
                                               typename Tsolver::ActionNode& action) const {
        try {
            // Generate the next states of this action
            auto next_states = solver.domain().get_next_state_distribution(state.state, action.action, thread_id).get_values();
            std::vector<typename Tsolver::StateNode*> untried_outcomes;
            std::vector<double> weights;
            std::vector<double> outcome_weights;

            for (auto ns : next_states) {
                std::pair<typename Tsolver::Graph::iterator, bool> i;

                solver.execution_policy().protect([&i, &solver, &ns](){
                    i = solver.graph().emplace(ns.state());
                });

                typename Tsolver::StateNode& next_node = const_cast<typename Tsolver::StateNode&>(*(i.first)); // we won't change the real key (StateNode::state) so we are safe
                double reward = 0.0;

                solver.execution_policy().protect([&reward, &solver, &state, &action, &next_node, &thread_id](){
                    reward = solver.domain().get_transition_reward(state.state, action.action, next_node.state, thread_id);
                }, next_node.mutex);

                solver.execution_policy().protect([&action, &next_node, &outcome_weights, &reward, &ns](){
                    auto ii = action.outcomes.insert(std::make_pair(&next_node, std::make_pair(reward, 1)));

                    if (ii.second) { // new outcome
                        action.dist_to_outcome.push_back(ii.first);
                        outcome_weights.push_back(ns.probability());
                    } else { // existing outcome (following code not efficient but hopefully very rare case if domain is well defined)
                        for (unsigned int oid = 0 ; oid < outcome_weights.size() ; oid++) {
                            if (action.dist_to_outcome[oid]->first == ii.first->first) { // found my outcome!
                                std::pair<double, std::size_t>& mp = ii.first->second;
                                mp.first = ((double) (outcome_weights[oid] * mp.first) + (reward * ns.probability())) / ((double) (outcome_weights[oid] + ns.probability()));
                                outcome_weights[oid] += ns.probability();
                                mp.second += 1; // useless in this mode a priori, but just keep track for coherency
                                break;
                            }
                        }
                    }
                }, action.parent->mutex);

                solver.execution_policy().protect([this, &next_node, &action, &i, &solver, &thread_id, &untried_outcomes, &weights, &ns](){
                    next_node.parents.insert(&action);

                    if (i.second) { // new node
                        next_node.terminal = solver.domain().is_terminal(next_node.state, thread_id);
                        std::pair<double, std::size_t> h = _heuristic(solver.domain(), next_node.state, thread_id);
                        next_node.value = h.first;
                        next_node.visits_count = h.second;
                    }

                    if (next_node.actions.empty()) {
                        if (solver.debug_logs()) spdlog::debug("Candidate next state: " + next_node.state.print() +
                                                               Tsolver::ExecutionPolicy::print_thread());
                        untried_outcomes.push_back(&next_node);
                        weights.push_back(ns.probability());
                    }
                }, next_node.mutex);
            }

            // Record the action's outcomes distribution
            solver.execution_policy().protect([&action, &outcome_weights](){
                action.dist = std::discrete_distribution<>(outcome_weights.begin(), outcome_weights.end());
            }, action.parent->mutex);

            // Pick a random next state
            if (untried_outcomes.empty()) {
                // All next states already visited => pick a random next state using action.dist
                std::size_t outcome_id = 0;

                solver.execution_policy().protect([&outcome_id, &action, &solver](){
                    outcome_id = action.dist(solver.gen());
                }, solver.gen_mutex());

                typename Tsolver::StateNode* outcome = nullptr;

                solver.execution_policy().protect([&action, &outcome, &outcome_id](){
                    outcome = action.dist_to_outcome[outcome_id]->first;
                }, action.parent->mutex);

                return outcome;
            } else {
                // Pick a random next state among untried ones
                std::discrete_distribution<> odist(weights.begin(), weights.end());
                std::size_t outcome_id = 0;

                solver.execution_policy().protect([&outcome_id, &odist, &solver](){
                    outcome_id = odist(solver.gen());
                }, solver.gen_mutex());

                return untried_outcomes[outcome_id];
            }
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&action, &e](){
                spdlog::error("SKDECIDE exception in MCTS when expanding action " + action.action.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, action.parent->mutex);
            throw;
        }
    }

    template <typename Ttransition_mode,
              std::enable_if_t<std::is_same<Ttransition_mode, StepTransitionMode<Tsolver>>::value ||
                               std::is_same<Ttransition_mode, SampleTransitionMode<Tsolver>>::value, int> = 0>
    typename Tsolver::StateNode* expand_action(Tsolver& solver,
                                               const std::size_t* thread_id,
                                               const Ttransition_mode& transition_mode,
                                               typename Tsolver::StateNode& state,
                                               typename Tsolver::ActionNode& action) const {
        try {
            if (!_checked_transition_mode) {
                spdlog::warn("Using MCTS full expansion mode with step() or sample() domain's transition mode assumes the domain is deterministic (unpredictable result otherwise).");
                _checked_transition_mode = true;
            }
            // Generate the next state of this action
            typename Tsolver::Domain::TransitionOutcome to = transition_mode.random_next_outcome(solver, thread_id, state.state, action.action);
            std::pair<typename Tsolver::Graph::iterator, bool> i;

            solver.execution_policy().protect([&i, &solver, &to](){
                i = solver.graph().emplace(to.state());
            });

            typename Tsolver::StateNode& next_node = const_cast<typename Tsolver::StateNode&>(*(i.first)); // we won't change the real key (StateNode::state) so we are safe

            solver.execution_policy().protect([&action, &next_node, &to](){
                auto ii = action.outcomes.insert(std::make_pair(&next_node, std::make_pair(to.reward(), 1)));
                action.dist_to_outcome.push_back(ii.first);
            }, action.parent->mutex);

            solver.execution_policy().protect([this, &next_node, &action, &i, &to, &solver, &thread_id](){
                next_node.parents.insert(&action);

                if (i.second) { // new node
                    next_node.terminal = to.terminal();
                    std::pair<double, std::size_t> h = _heuristic(solver.domain(), next_node.state, thread_id);
                    next_node.value = h.first;
                    next_node.visits_count = h.second;
                }
            }, next_node.mutex);

            // Record the action's outcomes distribution
            solver.execution_policy().protect([&action](){
                action.dist = std::discrete_distribution<>({1.0});
            }, action.parent->mutex);

            if (solver.debug_logs()) {
                solver.execution_policy().protect([&next_node](){
                    spdlog::debug("Candidate next state: " + next_node.state.print() +
                                  Tsolver::ExecutionPolicy::print_thread());
                }, next_node.mutex);
            }

            return &next_node;
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&action, &e](){
                spdlog::error("SKDECIDE exception in MCTS when expanding action " + action.action.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, action.parent->mutex);
            throw;
        }
    }

private :
    HeuristicFunctor _heuristic;
    mutable typename Tsolver::ExecutionPolicy::template atomic<bool> _checked_transition_mode;
};

/** Test if a given node needs to be expanded by sampling applicable actions and next states.
 *  Tries to sample new outcomes with a probability proportional to the number of actual expansions.
 *  Returns nullptr if we cannot sample new outcomes, otherwise a sampled unvisited outcome
 *  according to its probability (among only unvisited outcomes).
 *  REQUIREMENTS: returns nullptr if all actions have been already tried, and set the terminal
 *  flag of the returned next state
 */
template <typename Tsolver>
class PartialExpand {
public :
    typedef std::function<std::pair<double, std::size_t>
                          (typename Tsolver::Domain&, const typename Tsolver::Domain::State&, const std::size_t*)> HeuristicFunctor;

    PartialExpand(const HeuristicFunctor& heuristic = [](typename Tsolver::Domain& domain,
                                                         const typename Tsolver::Domain::State& state,
                                                         const std::size_t* thread_id) {
        return std::make_pair(0.0, 0);
    })
    : _heuristic(heuristic) {}

    PartialExpand(const PartialExpand& other)
    : _heuristic(other._heuristic) {}

    typename Tsolver::StateNode* operator()(Tsolver& solver,
                                            const std::size_t* thread_id,
                                            typename Tsolver::StateNode& n) const {
        try {
            if (solver.debug_logs()) {
                solver.execution_policy().protect([&n](){
                    spdlog::debug("Test expansion of state " + n.state.print() +
                                  Tsolver::ExecutionPolicy::print_thread());
                }, n.mutex);
            }

            // Sample an action
            std::bernoulli_distribution dist_state_expansion((n.visits_count > 0)?
                                                             (((double) n.expansions_count) / ((double) n.visits_count)):
                                                             1.0);
            typename Tsolver::ActionNode* action_node = nullptr;
            bool dist_res = false;

            solver.execution_policy().protect([&dist_res, &solver, &dist_state_expansion](){
                dist_res = dist_state_expansion(solver.gen());
            }, solver.gen_mutex());

            if (dist_res) {
                typename Tsolver::Domain::Action action = solver.domain().get_applicable_actions(n.state, thread_id).sample();
                solver.execution_policy().protect([&n, &action, &action_node, &solver](){
                    auto a = n.actions.emplace(typename Tsolver::ActionNode(action));

                    if (a.second) { // new action
                        n.expansions_count += 1;
                    }

                    action_node = &const_cast<typename Tsolver::ActionNode&>(*(a.first)); // we won't change the real key (ActionNode::action) so we are safe
                    if (solver.debug_logs()) { spdlog::debug("Tried to sample a new action: " + action_node->action.print() +
                                                             Tsolver::ExecutionPolicy::print_thread()); }
                }, n.mutex);
            } else {
                std::vector<typename Tsolver::ActionNode*> actions;

                solver.execution_policy().protect([&n, &actions](){
                    for (auto& a : n.actions) {
                        actions.push_back(&a);
                    }
                }, n.mutex);

                std::uniform_int_distribution<> dist_known_actions(0, actions.size()-1);
                std::size_t action_id = 0;

                solver.execution_policy().protect([&action_id, &solver, &dist_known_actions](){
                    action_id = dist_known_actions(solver.gen());
                }, solver.gen_mutex());

                action_node = actions[action_id];
                if (solver.debug_logs()) {
                    solver.execution_policy().protect([&action_node](){
                        spdlog::debug("Sampled among known actions: " + action_node->action.print() +
                                      Tsolver::ExecutionPolicy::print_thread());
                    }, action_node->parent.mutex);
                }
            }

            // Sample an outcome
            std::bernoulli_distribution dist_action_expansion((action_node->visits_count > 0)?
                                                              (((double) action_node->expansions_count) / ((double) action_node->visits_count)):
                                                              1.0);
            typename Tsolver::StateNode* ns = nullptr;

            solver.execution_policy().protect([&dist_res, &solver, &dist_action_expansion](){
                dist_res = dist_action_expansion(solver.gen());
            }, solver.gen_mutex());

            if (dist_res) {
                std::unique_ptr<typename Tsolver::Domain::TransitionOutcome> to = solver.transition_mode().random_next_outcome(solver, thread_id, n.state, action_node->action);
                std::pair<typename Tsolver::Graph::iterator, bool> s;

                solver.execution_policy().protect([&s, &solver, &to](){
                    s = solver.graph().emplace(to->state());
                });

                ns = &const_cast<typename Tsolver::StateNode&>(*(s.first)); // we won't change the real key (StateNode::state) so we are safe

                if (s.second) { // new state
                    solver.execution_policy().protect([this, &ns, &to, &solver, &thread_id](){
                        ns->terminal = to->termination();
                        std::pair<double, std::size_t> h = _heuristic(solver.domain(), ns->state, thread_id);
                        ns->value = h.first;
                        ns->visits_count = h.second;
                    }, ns->mutex);
                }

                std::pair<typename Tsolver::ActionNode::OutcomeMap::iterator, bool> ins;
                solver.execution_policy().protect([&action_node, &ns, &to, ins](){
                    ins = action_node->outcomes.emplace(std::make_pair(ns, std::make_pair(to->reward(), 1)));
                }, action_node->parent->mutex);

                // Update the outcome's reward and visits count
                if (ins.second) { // new outcome
                    solver.execution_policy().protect([&action_node, &ins](){
                        action_node->dist_to_outcome.push_back(ins.first);
                        action_node->expansions_count += 1;
                    }, action_node->parent->mutex);

                    solver.execution_policy().protect([&ns, &action_node](){
                        ns->parents.insert(action_node);
                    }, ns->mutex);
                } else { // known outcome
                    solver.execution_policy().protect([&ins, &to, &ns](){
                        std::pair<double, std::size_t>& mp = ins.first->second;
                        mp.first = ((double) (mp.second * mp.first) + to->reward()) / ((double) (mp.second + 1));
                        mp.second += 1;
                        ns = nullptr; // we have not discovered anything new
                    }, action_node->parent->mutex);
                }

                // Reconstruct the probability distribution
                solver.execution_policy().protect([&action_node](){
                    std::vector<double> weights(action_node->dist_to_outcome.size());

                    for (unsigned int oid = 0 ; oid < weights.size() ; oid++) {
                        weights[oid] = action_node->dist_to_outcome[oid]->second.second;
                    }

                    action_node->dist = std::discrete_distribution<>(weights.begin(), weights.end());
                }, action_node->parent->mutex);

                if (solver.debug_logs()) {
                    solver.execution_policy().protect([&ns](){
                        spdlog::debug("Tried to sample a new outcome: " + ns->state.print() +
                                      Tsolver::ExecutionPolicy::print_thread());
                    });
                }
            } else {
                ns = nullptr; // we have not discovered anything new

                if (solver.debug_logs()) {
                    solver.execution_policy().protect([&ns](){
                        spdlog::debug("Sampled among known outcomes: " + ns->state.print() +
                                      Tsolver::ExecutionPolicy::print_thread());
                    });
                }
            }

            return ns;
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&n, &e](){
                spdlog::error("SKDECIDE exception in MCTS when expanding state " + n.state.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, n.mutex);
            throw;
        }
    }

private :
    HeuristicFunctor _heuristic;
};

/** UCB1 Best Child */
template <typename Tsolver>
class UCB1ActionSelector {
public :
    // 1/sqrt(2) is a good compromise for rewards in [0;1]
    UCB1ActionSelector(double ucb_constant = 1.0 / std::sqrt(2.0))
    : _ucb_constant(ucb_constant) {}

    UCB1ActionSelector(const UCB1ActionSelector& other)
    : _ucb_constant((double) other._ucb_constant) {}

    typename Tsolver::ActionNode* operator()(Tsolver& solver,
                                             const std::size_t* thread_id,
                                             const typename Tsolver::StateNode& n) const {
        double best_value = -std::numeric_limits<double>::max();
        typename Tsolver::ActionNode* best_action = nullptr;

        solver.execution_policy().protect([this, &n, &best_value, &best_action, &solver](){
            for (const auto& a : n.actions) {
                if (a.visits_count > 0) {
                    double tentative_value = a.value + (2.0 * _ucb_constant * std::sqrt((2.0 * std::log((double) n.visits_count)) / ((double) a.visits_count)));

                    if (tentative_value > best_value) {
                        best_value = tentative_value;
                        best_action = &const_cast<typename Tsolver::ActionNode&>(a); // we won't change the real key (ActionNode::action) so we are safe
                    }
                }
            }

            if (solver.debug_logs()) { spdlog::debug("UCB1 selection from state " + n.state.print() +
                                                     ": value=" + StringConverter::from(best_value) +
                                                     ", action=" + ((best_action != nullptr)?(best_action->action.print()):("nullptr")) +
                                                     Tsolver::ExecutionPolicy::print_thread()); }
        }, n.mutex);

        return best_action;
    }

private :
    typename Tsolver::ExecutionPolicy::template atomic<double> _ucb_constant;
};

/** Select action with maximum Q-value */
template <typename Tsolver>
class BestQValueActionSelector {
public :
    typename Tsolver::ActionNode* operator()(Tsolver& solver,
                                             const std::size_t* thread_id,
                                             const typename Tsolver::StateNode& n) const {
        double best_value = -std::numeric_limits<double>::max();
        typename Tsolver::ActionNode* best_action = nullptr;

        solver.execution_policy().protect([&n, &best_value, &best_action, &solver](){
            for (const auto& a : n.actions) {
                if (a.visits_count > 0) {
                    if (a.value > best_value) {
                        best_value = a.value;
                        best_action = &const_cast<typename Tsolver::ActionNode&>(a); // we won't change the real key (ActionNode::action) so we are safe
                    }
                }
            }

            if (solver.debug_logs()) { spdlog::debug("Best Q-value selection from state " + n.state.print() +
                                                     ": value=" + StringConverter::from(best_value) +
                                                     ", action=" + ((best_action != nullptr)?(best_action->action.print()):("nullptr")) +
                                                     Tsolver::ExecutionPolicy::print_thread()); }
        }, n.mutex);

        return best_action;
    }
};

/** Default rollout policy */
template <typename Tsolver>
class DefaultRolloutPolicy {
public :
    typedef std::function<typename Tsolver::Domain::Action
                          (typename Tsolver::Domain&, const typename Tsolver::Domain::State&, const std::size_t*)> PolicyFunctor;

    DefaultRolloutPolicy(const PolicyFunctor& policy = [](typename Tsolver::Domain& domain,
                                                          const typename Tsolver::Domain::State& state,
                                                          const std::size_t* thread_id) {
        return domain.get_applicable_actions(state, thread_id).sample();
    })
    : _policy(policy) {}

    void operator()(Tsolver& solver,
                    const std::size_t* thread_id,
                    typename Tsolver::StateNode& n,
                    std::size_t d) const {
        try {
            typename Tsolver::Domain::State current_state;

            solver.execution_policy().protect([&solver, &n, &current_state](){
                if (solver.debug_logs()) { spdlog::debug("Launching default rollout policy from state " + n.state.print() +
                                                         Tsolver::ExecutionPolicy::print_thread()); }
                current_state = n.state;
            }, n.mutex);

            bool termination = false;
            std::size_t current_depth = d;
            double reward = 0.0;
            double gamma_n = 1.0;

            while(!termination && current_depth < solver.max_depth()) {
                typename Tsolver::Domain::Action action = _policy(solver.domain(), current_state, thread_id);
                typename Tsolver::Domain::TransitionOutcome o = solver.transition_mode().random_next_outcome(solver, thread_id, current_state, action);
                reward += gamma_n * (o.reward());
                gamma_n *= solver.discount();
                current_state = o.state();
                termination = o.terminal();
                current_depth++;
                if (solver.debug_logs()) { spdlog::debug("Sampled transition: action=" + action.print() +
                                                         ", next state=" + current_state.print() +
                                                         ", reward=" + StringConverter::from(o.reward()) +
                                                         Tsolver::ExecutionPolicy::print_thread()); }
            }

            // since we can come to state n after exhausting the depth, n might be already visited
            // so don't erase its value but rather update it
            solver.execution_policy().protect([&n, &reward](){
                n.value = ((n.visits_count * n.value)  + reward) / ((double) (n.visits_count + 1));
                n.visits_count += 1;
            }, n.mutex);
        } catch (const std::exception& e) {
            solver.execution_policy().protect([&n, &e](){
                spdlog::error("SKDECIDE exception in MCTS when simulating the random default policy from state " + n.state.print() + ": " + e.what() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, n.mutex);
            throw;
        }
    }

private :
    PolicyFunctor _policy;
};

/** Graph backup: update Q values using the graph ancestors (rather than only the trajectory leading to n) */
template <typename Tsolver>
struct GraphBackup {
    void operator()(Tsolver& solver,
                    const std::size_t* thread_id,
                    typename Tsolver::StateNode& n) const {
        if (solver.debug_logs()) {
            solver.execution_policy().protect([&n](){
                spdlog::debug("Back-propagating values from state " + n.state.print() +
                              Tsolver::ExecutionPolicy::print_thread());
            }, n.mutex);
        }

        std::size_t depth = 0; // used to prevent infinite loop in case of cycles
        std::unordered_set<typename Tsolver::StateNode*> frontier;
        frontier.insert(&n);

        while (!frontier.empty() && depth <= solver.max_depth()) {
            depth++;
            std::unordered_set<typename Tsolver::StateNode*> new_frontier;

            for (auto& f : frontier) {
                update_frontier(solver, new_frontier, f, &solver.execution_policy());
            }

            frontier = new_frontier;
        }
    }

    template <typename Texecution_policy,
              std::enable_if_t<std::is_same<Texecution_policy, SequentialExecution>::value, int> = 0>
    static void update_frontier(Tsolver& solver,
                                std::unordered_set<typename Tsolver::StateNode*>& new_frontier, typename Tsolver::StateNode* f,
                                [[maybe_unused]] Texecution_policy* execution_policy) {
        for (auto& a : f->parents) {
            double q_value = a->outcomes[f].first + (solver.discount() * (f->value));
            a->value = (((a->visits_count) * (a->value))  + q_value) / ((double) (a->visits_count + 1));
            a->visits_count += 1;
            typename Tsolver::StateNode* parent_node = a->parent;
            parent_node->value = (((parent_node->visits_count) * (parent_node->value))  + (a->value)) / ((double) (parent_node->visits_count + 1));
            parent_node->visits_count += 1;
            new_frontier.insert(parent_node);
            if (solver.debug_logs()) { spdlog::debug("Updating state " + parent_node->state.print() +
                                                    ": value=" + StringConverter::from(parent_node->value) +
                                                    ", visits=" + StringConverter::from(parent_node->visits_count) +
                                                    Tsolver::ExecutionPolicy::print_thread()); }
        }
    }

    template <typename Texecution_policy,
              std::enable_if_t<std::is_same<Texecution_policy, ParallelExecution>::value, int> = 0>
    static void update_frontier(Tsolver& solver,
                                std::unordered_set<typename Tsolver::StateNode*>& new_frontier, typename Tsolver::StateNode* f,
                                [[maybe_unused]] Texecution_policy* execution_policy) {
        std::list<typename Tsolver::ActionNode*> parents;
        solver.execution_policy().protect([&f, &parents](){
            std::copy(f->parents.begin(), f->parents.end(), std::inserter(parents, parents.end()));
        }, f->mutex);
        for (auto& a : parents) {
            solver.execution_policy().protect([&a, &solver, &f, &new_frontier](){
                double q_value = a->outcomes[f].first + (solver.discount() * (f->value));
                a->value = (((a->visits_count) * (a->value))  + q_value) / ((double) (a->visits_count + 1));
                a->visits_count += 1;
                typename Tsolver::StateNode* parent_node = a->parent;
                parent_node->value = (((parent_node->visits_count) * (parent_node->value))  + (a->value)) / ((double) (parent_node->visits_count + 1));
                parent_node->visits_count += 1;
                new_frontier.insert(parent_node);
                if (solver.debug_logs()) { spdlog::debug("Updating state " + parent_node->state.print() +
                                                        ": value=" + StringConverter::from(parent_node->value) +
                                                        ", visits=" + StringConverter::from(parent_node->visits_count) +
                                                        Tsolver::ExecutionPolicy::print_thread()); }
            }, a->parent->mutex);
        }
    }
};

template <typename Tdomain,
          typename TexecutionPolicy = SequentialExecution,
          template <typename Tsolver> class TtransitionMode = DistributionTransitionMode,
          template <typename Tsolver> class TtreePolicy = DefaultTreePolicy,
          template <typename Tsolver> class Texpander = FullExpand,
          template <typename Tsolver> class TactionSelectorOptimization = UCB1ActionSelector,
          template <typename Tsolver> class TactionSelectorExecution = BestQValueActionSelector,
          template <typename Tsolver> class TrolloutPolicy = DefaultRolloutPolicy,
          template <typename Tsolver> class TbackPropagator = GraphBackup>
class MCTSSolver {
public :
    typedef MCTSSolver<Tdomain, TexecutionPolicy,
                       TtransitionMode, TtreePolicy, Texpander,
                       TactionSelectorOptimization, TactionSelectorExecution,
                       TrolloutPolicy, TbackPropagator> Solver;

    typedef Tdomain Domain;
    typedef typename Domain::State State;
    typedef typename Domain::Event Action;
    typedef TexecutionPolicy ExecutionPolicy;
    typedef TtransitionMode<Solver> TransitionMode;
    typedef TtreePolicy<Solver> TreePolicy;
    typedef Texpander<Solver> Expander;
    typedef TactionSelectorOptimization<Solver> ActionSelectorOptimization;
    typedef TactionSelectorExecution<Solver> ActionSelectorExecution;
    typedef TrolloutPolicy<Solver> RolloutPolicy;
    typedef TbackPropagator<Solver> BackPropagator;

    typedef typename ExecutionPolicy::template atomic<std::size_t> atomic_size_t;
    typedef typename ExecutionPolicy::template atomic<double> atomic_double;
    typedef typename ExecutionPolicy::template atomic<bool> atomic_bool;

    struct StateNode;

    struct ActionNode {
        Action action;
        typedef std::unordered_map<StateNode*, std::pair<double, std::size_t>> OutcomeMap; // next state nodes owned by _graph
        OutcomeMap outcomes;
        std::vector<typename OutcomeMap::iterator> dist_to_outcome;
        std::discrete_distribution<> dist;
        atomic_size_t expansions_count; // used only for partial expansion mode
        atomic_double value;
        atomic_size_t visits_count;
        StateNode* parent;

        ActionNode(const Action& a)
            : action(a), expansions_count(0), value(0.0),
              visits_count(0), parent(nullptr) {}

        ActionNode(const ActionNode& a)
            : action(a.action), outcomes(a.outcomes), dist_to_outcome(a.dist_to_outcome),
              dist(a.dist), expansions_count((std::size_t) a.expansions_count),
              value((double) a.value), visits_count((std::size_t) a.visits_count),
              parent(a.parent) {}

        struct Key {
            const Action& operator()(const ActionNode& an) const { return an.action; }
        };
    };

    struct StateNode {
        typedef typename SetTypeDeducer<ActionNode, Action>::Set ActionSet;
        State state;
        atomic_bool terminal;
        atomic_bool expanded; // used only for full expansion mode
        atomic_size_t expansions_count; // used only for partial expansion mode
        ActionSet actions;
        atomic_double value;
        atomic_size_t visits_count;
        std::unordered_set<ActionNode*> parents;
        mutable typename ExecutionPolicy::Mutex mutex;

        StateNode(const State& s)
            : state(s), terminal(false), expanded(false),
              expansions_count(0), value(0.0), visits_count(0) {}

        StateNode(const StateNode& s)
            : state(s.state), terminal((bool) s.terminal), expanded((bool) s.expanded),
              expansions_count((std::size_t) s.expansions_count), actions(s.actions),
              value((double) s.value), visits_count((std::size_t) s.visits_count),
              parents(s.parents) {}

        struct Key {
            const State& operator()(const StateNode& sn) const { return sn.state; }
        };
    };

    typedef typename SetTypeDeducer<StateNode, State>::Set Graph;

    MCTSSolver(Domain& domain,
               std::size_t time_budget = 3600000,
               std::size_t rollout_budget = 100000,
               std::size_t max_depth = 1000,
               double discount = 1.0,
               bool online_node_garbage = false,
               bool debug_logs = false,
               std::unique_ptr<TreePolicy> tree_policy = std::make_unique<TreePolicy>(),
               std::unique_ptr<Expander> expander = std::make_unique<Expander>(),
               std::unique_ptr<ActionSelectorOptimization> action_selector_optimization = std::make_unique<ActionSelectorOptimization>(),
               std::unique_ptr<ActionSelectorExecution> action_selector_execution = std::make_unique<ActionSelectorExecution>(),
               std::unique_ptr<RolloutPolicy> rollout_policy = std::make_unique<RolloutPolicy>(),
               std::unique_ptr<BackPropagator> back_propagator = std::make_unique<BackPropagator>())
    : _domain(domain),
      _time_budget(time_budget), _rollout_budget(rollout_budget),
      _max_depth(max_depth), _discount(discount), _nb_rollouts(0),
      _online_node_garbage(online_node_garbage),
      _debug_logs(debug_logs),
      _tree_policy(std::move(tree_policy)),
      _expander(std::move(expander)),
      _action_selector_optimization(std::move(action_selector_optimization)),
      _action_selector_execution(std::move(action_selector_execution)),
      _rollout_policy(std::move(rollout_policy)),
      _back_propagator(std::move(back_propagator)),
      _current_state(nullptr) {
        if (debug_logs) {
            spdlog::set_level(spdlog::level::debug);
        } else {
            spdlog::set_level(spdlog::level::info);
        }

        std::random_device rd;
        _gen = std::make_unique<std::mt19937>(rd());
    }

    // clears the solver (clears the search graph, thus preventing from reusing
    // previous search results)
    void clear() {
        _graph.clear();
    }

    // solves from state s
    void solve(const State& s) {
        try {
            spdlog::info("Running " + ExecutionPolicy::print_type() + " MCTS solver from state " + s.print());
            auto start_time = std::chrono::high_resolution_clock::now();
            _nb_rollouts = 0;

            // Get the root node
            auto si = _graph.emplace(s);
            StateNode& root_node = const_cast<StateNode&>(*(si.first)); // we won't change the real key (StateNode::state) so we are safe

            boost::integer_range<std::size_t> parallel_rollouts(0, _domain.get_parallel_capacity());

            std::for_each(ExecutionPolicy::policy, parallel_rollouts.begin(), parallel_rollouts.end(), [this, &start_time, &root_node] (const std::size_t& thread_id) {

                while (elapsed_time(start_time) < _time_budget && _nb_rollouts < _rollout_budget) {

                    std::size_t depth = 0;
                    StateNode* sn = (*_tree_policy)(*this, &thread_id, *_expander, *_action_selector_optimization, root_node, depth);
                    (*_rollout_policy)(*this, &thread_id, *sn, depth);
                    (*_back_propagator)(*this, &thread_id, *sn);
                    _nb_rollouts++;

                }

            });

            auto end_time = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(end_time - start_time).count();
            spdlog::info("MCTS finished to solve from state " + s.print() +
                         " in " + StringConverter::from((double) duration / (double) 1e9) + " seconds with " +
                         StringConverter::from(_nb_rollouts) + " rollouts.");
        } catch (const std::exception& e) {
            spdlog::error("MCTS failed solving from state " + s.print() + ". Reason: " + e.what());
            throw;
        }
    }

    bool is_solution_defined_for(const State& s) {
        auto si = _graph.find(s);
        if (si == _graph.end()) {
            return false;
        } else {
            return (*_action_selector_execution)(*this, nullptr, *si) != nullptr;
        }
    }

    Action get_best_action(const State& s) {
        auto si = _graph.find(s);
        ActionNode* action = nullptr;
        if (si != _graph.end()) {
            action = (*_action_selector_execution)(*this, nullptr, *si);
        }
        if (action == nullptr) {
            spdlog::error("SKDECIDE exception: no best action found in state " + s.print());
            throw std::runtime_error("SKDECIDE exception: no best action found in state " + s.print());
        } else {
            if (_debug_logs) {
                std::string str = "(";
                for (const auto& o : action->outcomes) {
                    str += "\n    " + o.first->state.print();
                }
                str += "\n)";
                spdlog::debug("Best action's known outcomes:\n" + str);
            }
            if (_online_node_garbage && _current_state) {
                std::unordered_set<StateNode*> root_subgraph, child_subgraph;
                compute_reachable_subgraph(_current_state, root_subgraph);
                compute_reachable_subgraph(const_cast<StateNode*>(&(*si)), child_subgraph); // we won't change the real key (StateNode::state) so we are safe
                remove_subgraph(root_subgraph, child_subgraph);
            }
            _current_state = const_cast<StateNode*>(&(*si)); // we won't change the real key (StateNode::state) so we are safe
            _action_prefix.push_back(action->action);
            return action->action;
        }
    }

    double get_best_value(const State& s) {
        auto si = _graph.find(StateNode(s));
        ActionNode* action = nullptr;
        if (si != _graph.end()) {
            action = (*_action_selector_execution)(*this, nullptr, *si);
        }
        if (action == nullptr) {
            spdlog::error("SKDECIDE exception: no best action found in state " + s.print());
            throw std::runtime_error("SKDECIDE exception: no best action found in state " + s.print());
        } else {
            return action->value;
        }
    }

    std::size_t nb_of_explored_states() const {
        return _graph.size();
    }

    std::size_t nb_rollouts() const {
        return _nb_rollouts;
    }

    typename MapTypeDeducer<State, std::pair<Action, double>>::Map policy() {
        typename MapTypeDeducer<State, std::pair<Action, double>>::Map p;
        for (auto& n : _graph) {
            ActionNode* action = (*_action_selector_execution)(*this, nullptr, n);
            if (action != nullptr) {
                p.insert(std::make_pair(n.state, std::make_pair(action->action, (double) action->value)));
            }
        }
        return p;
    }

    Domain& domain() { return _domain; }

    std::size_t time_budget() const { return _time_budget; }

    std::size_t rollout_budget() const { return _rollout_budget; }

    std::size_t max_depth() const { return _max_depth; }

    double discount() const { return _discount; }

    ExecutionPolicy& execution_policy() { return _execution_policy; }

    TransitionMode& transition_mode() { return _transition_mode; }

    const TreePolicy& tree_policy() { return *_tree_policy; }

    const Expander& expander() { return *_expander; }

    const ActionSelectorOptimization& action_selector_optimization() { return *_action_selector_optimization; }

    const ActionSelectorExecution& action_selector_execution() { return *_action_selector_execution; }

    const RolloutPolicy& rollout_policy() { return *_rollout_policy; }

    const BackPropagator& back_propagator() { return *_back_propagator; }

    Graph& graph() { return _graph; }

    const std::list<Action>& action_prefix() const { return _action_prefix; }

    std::mt19937& gen() { return *_gen; }

    typename ExecutionPolicy::Mutex& gen_mutex() { return _gen_mutex; }

    bool debug_logs() const { return _debug_logs; }

private :

    Domain& _domain;
    atomic_size_t _time_budget;
    atomic_size_t _rollout_budget;
    atomic_size_t _max_depth;
    atomic_double _discount;
    atomic_size_t _nb_rollouts;
    bool _online_node_garbage;
    atomic_bool _debug_logs;

    ExecutionPolicy _execution_policy;
    TransitionMode _transition_mode;

    std::unique_ptr<TreePolicy> _tree_policy;
    std::unique_ptr<Expander> _expander;
    std::unique_ptr<ActionSelectorOptimization> _action_selector_optimization;
    std::unique_ptr<ActionSelectorExecution> _action_selector_execution;
    std::unique_ptr<RolloutPolicy> _rollout_policy;
    std::unique_ptr<BackPropagator> _back_propagator;

    Graph _graph;
    StateNode* _current_state;
    std::list<Action> _action_prefix;

    std::unique_ptr<std::mt19937> _gen;
    typename ExecutionPolicy::Mutex _gen_mutex;
    typename ExecutionPolicy::Mutex _time_mutex;

    void compute_reachable_subgraph(StateNode* node, std::unordered_set<StateNode*>& subgraph) {
        std::unordered_set<StateNode*> frontier;
        frontier.insert(node);
        subgraph.insert(node);
        while(!frontier.empty()) {
            std::unordered_set<StateNode*> new_frontier;
            for (auto& n : frontier) {
                for (auto& action : n->actions) {
                    for (auto& outcome : action.outcomes) {
                        if (subgraph.find(outcome.first) == subgraph.end()) {
                            new_frontier.insert(outcome.first);
                            subgraph.insert(outcome.first);
                        }
                    }
                }
            }
            frontier = new_frontier;
        }
    }

    void remove_subgraph(std::unordered_set<StateNode*>& root_subgraph, std::unordered_set<StateNode*>& child_subgraph) {
        std::unordered_set<StateNode*> removed_subgraph;
        // First pass: look for nodes in root_subgraph but not child_subgraph and remove
        // those nodes from their children's parents
        // Don't actually remove those nodes in the first pass otherwise some children to remove
        // won't exist anymore when looking for their parents
        for (auto& n : root_subgraph) {
            if (child_subgraph.find(n) == child_subgraph.end()) {
                for (auto& action : n->actions) {
                    for (auto& outcome : action.outcomes) {
                        // we won't change the real key (ActionNode::action) so we are safe
                        outcome.first->parents.erase(&const_cast<ActionNode&>(action));
                    }
                }
                removed_subgraph.insert(n);
            }
        }
        // Second pass: actually remove nodes in root_subgraph but not in child_subgraph
        for (auto& n : removed_subgraph) {
            _graph.erase(StateNode(n->state));
        }
    }

    std::size_t elapsed_time(const std::chrono::time_point<std::chrono::high_resolution_clock>& start_time) {
        std::size_t milliseconds_duration;
        _execution_policy.protect([&milliseconds_duration, &start_time](){
            milliseconds_duration = static_cast<std::size_t>(
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::high_resolution_clock::now() - start_time
                ).count()
            );
        }, _time_mutex);
        return milliseconds_duration;
    }
}; // MCTSSolver class

/** UCT is MCTS with the default template options */
template <typename Tdomain,
          typename Texecution_policy,
          template <typename Tsolver> typename TtransitionMode,
          template <typename Tsolver> typename ...T>
using UCTSolver = MCTSSolver<Tdomain, Texecution_policy, TtransitionMode, T...>;

} // namespace skdecide

#endif // SKDECIDE_MCTS_HH


// @authors:     Ahmad Humayun
// @contact:     ahumayun@cc.gatech.edu
// @affiliation: Georgia Institute of Technology
// @date:        Fall 2013 - Summer 2014

/* precomptrees.cpp */

#ifndef __PRECOMPTREESIMPL_HPP__
#define __PRECOMPTREESIMPL_HPP__

#include <iostream>
#include <fstream>
#include <boost/format.hpp>

#include "graph.h"

using boost::format;

/*
  special constants for node->parent. Duplicated in graph.cpp, both should match!
*/
#define TERMINAL ( (arc *) 1 )    /* to terminal */
#define ORPHAN   ( (arc *) 2 )    /* orphan */

template<typename captype, typename tcaptype, typename flowtype>
 void Graph<captype, tcaptype, flowtype>::transform_seed_trees(
    const int keep_src_id, const SrcSeedList* remove_seed_list)
{
  node *i;

  typename SrcSeedList::const_iterator it_orig_cap, it_new_cap;

  // useful for maxflow_reuse_trees_init()
  curr_src_origin_idx = keep_src_id;

  // iterate over all the fg seed nodes which need to be converted from Seed S
  // trees to T sink trees. This loop looks at the fg seeds which need to be
  // discarded: it changes its residual capacities; adjusts the flow and adds
  // to the queue for being marked as the sink tree
  for (it_orig_cap = remove_seed_list[0].begin(),
       it_new_cap = remove_seed_list[1].begin();
       it_orig_cap != remove_seed_list[0].end(); ++it_orig_cap, ++it_new_cap) {
    // get the pointer to the current node
    i = nodes + it_orig_cap->first;

    tcaptype new_sink_residual_cap, new_src_residual_cap;

    // get all the capacities
    tcaptype residual_cap = get_trcap(i - nodes);
    tcaptype orig_src_cap = it_orig_cap->second.first;
    //tcaptype orig_sink_cap = it_orig_cap->second.second;
    tcaptype new_src_cap = it_new_cap->second.first;
    tcaptype new_sink_cap = it_new_cap->second.second;

    tcaptype flow_s = orig_src_cap - residual_cap;

    // we can directly set the residual capacity on the sink link because
    // we know there was no capacity on it originally i.e.
    // c_{it} = r_{it} = f_{it} = 0   /   orig_sink_cap := 0
    new_sink_residual_cap = new_sink_cap;

    // if the capacity is reduced only to the extent that it can still
    // accomodate the old flow (see Kohli and Torr, PAMI 2007)
    if (new_src_cap >= flow_s) {
      new_src_residual_cap = new_src_cap - flow_s;
    } else {
      /* if not then add an alpha value to the sink capacity */
      new_src_residual_cap = 0;
      new_sink_residual_cap += flow_s - new_src_cap;
      flow -= flow_s - new_src_cap;
    }

    set_trcap(it_orig_cap->first, new_src_residual_cap - new_sink_residual_cap);

    //std::cout << "Node " << (node_id)(i - nodes) << ": Old " << residual_cap << ", New " << get_trcap(i - nodes) << std::endl;

    // push to queue according to whether they have positive residual capacity
    // or negative. If they have negative residual capacity on its t link then
    // the node and its children can be converted to Sink T tree in the
    // upcoming while loop. Otherwise it needs to be added to the Source S
    // tree
    /*
    i->parent = TERMINAL;
    if (get_trcap(i - nodes) < 0)
      t_tree_children.push(i);
    else if (get_trcap(i - nodes) > 0)
      s_tree_children.push(i);
    */
    if (get_trcap(i - nodes) != 0)
      i->parent = TERMINAL;
    else {
      i->src_origin_idx = INVALID_SRC_ID;
      i->parent = NULL;
    }
  }

  node* n_ptr;
  NodeEditSet tree_nodes;

  // iterate over all nodes and add terminal and free nodes to the tree_nodes
  // list (and indicate as root in orig_root_nodes)
  for (n_ptr = nodes; n_ptr < node_last; ++n_ptr) {
    // skip node if its not a terminal or free node
    if (n_ptr->parent == TERMINAL || n_ptr->parent == NULL) {
      tree_nodes.insert(n_ptr);
    }
  }

  normalize_tree(tree_nodes);

  // iterate over all the nodes and set as free if they were neither part of
  // seed src or the converted sink trees
  /*for (i = nodes; i < node_last; i++) {
    if (i->src_origin_idx != keep_src_id && i->src_origin_idx != INVALID_SRC_ID &&
        i->is_sink != 1 && i->tr_cap == 0) {
      //if (PRINT_DEBUG && ((node_id)(i - nodes) == 41 || (node_id)(i - nodes) == 69 || (node_id)(i - nodes) == 99 || (node_id)(i - nodes) == 70))
      //  std::cout << "\tMarking " << i - nodes << " as free" << std::endl;

      i->parent = NULL;
      i->src_origin_idx = INVALID_SRC_ID;
      //mark_node(i - nodes);

      // mark all the neighboring nodes as active too
      for (arc* a = i->first; a; a = a->next) {
        mark_node(a->head - nodes);
        //std::cout << a->head - nodes << ", ";
      }
    }
  } */
}

template<typename captype, typename tcaptype, typename flowtype>
  void Graph<captype, tcaptype, flowtype>::set_active_normalization()
{
  node* n_ptr;
  node* neigh_ptr;

  // iterate over all nodes to see which ones should be set to active
  for (n_ptr = nodes; n_ptr < node_last; ++n_ptr) {
    //set_active(n_ptr);
    // if node doesn't belong to a tree after normalization, don't mark it
    //  active. A neighboring node, which belongs to a tree, and hence can
    //  grow into it would be marked when it sees this node
    if (n_ptr->parent == NULL || n_ptr->parent == ORPHAN) {
      continue;
    } else {
      // we now know that n_ptr belongs to a valid tree

      // iterate over each neighbor to see if there is a reason to set this
      // node active
      for (arc* a = n_ptr->first; a; a = a->next) {
        neigh_ptr = a->head;

        // neighbor is free or orphan (i.e. it doesnt belong to a tree) set
        // this node as active, so that the neighbor can be grown into
        if (neigh_ptr->parent == NULL || neigh_ptr->parent == ORPHAN) {
          set_active(n_ptr);
          break;
        } else {
          // if the neighbor belongs to an opposite tree
          if (n_ptr->is_sink != neigh_ptr->is_sink) {
            // if there is an unsaturated path between the nodes
            //  if u is in src then the arc u->v should be unsaturated
            //  if u is in sink then the arc v->u should be unsaturated
            if ((!n_ptr->is_sink && a->r_cap) ||
                (n_ptr->is_sink && a->sister->r_cap)) {
              // if the neighbor is active, then no need to set this node
              //  to active. It will get noticed anyways by neighbor anyways
              //if (!neigh_ptr->next) {
                set_active(n_ptr);
                break;
              //}
            }
          }
        }
      }
    }
  }
}

template<typename captype, typename tcaptype, typename flowtype>
  void Graph<captype, tcaptype, flowtype>::normalize_node(node* const n_ptr)
{
  captype up_rcap, dwn_rcap;

  bool is_root = (n_ptr->parent == TERMINAL || n_ptr->parent == NULL || n_ptr->parent == ORPHAN);

  // note down the original sink/src value
  int old_is_sink = -1;
  if (n_ptr->parent != NULL && n_ptr->parent != ORPHAN)
    old_is_sink = n_ptr->is_sink;

  // check if the parent was sink/src
  int parent_is_sink = -1;
  if (!is_root && n_ptr->parent->head->parent != NULL &&
                  n_ptr->parent->head->parent != ORPHAN)
    parent_is_sink = n_ptr->parent->head->is_sink;

  // note the residual capacities both on tr and n links
  if (!is_root) {
    up_rcap = n_ptr->parent->r_cap;
    dwn_rcap = n_ptr->parent->sister->r_cap;
  }
  tcaptype trcap = get_trcap(n_ptr-nodes);

  /*
  if (trcap != 0) {
    n_ptr->next = NULL;
    n_ptr->is_marked = 0;
    std::cout << " | " << (!n_ptr->next);
    set_active(n_ptr);
  }*/

  // if the node can be connected to the parent, then update values,
  // otherwise either set to free or connect to terminal if t residual
  // capacity preset TIME
  if ((parent_is_sink == 0 && dwn_rcap > 0 && trcap >= 0) ||
      (parent_is_sink == 1 && up_rcap > 0 && trcap <= 0)) {
    n_ptr->is_sink = parent_is_sink;
    n_ptr->src_origin_idx = n_ptr->parent->head->src_origin_idx;
    n_ptr->TS = TIME + 1;
    n_ptr->DIST = n_ptr->parent->head->DIST + 1;
  } else if (trcap != 0) {
    // directly connect to terminal
    n_ptr->is_sink = trcap < 0;
    n_ptr->src_origin_idx = n_ptr->is_sink ? INVALID_SRC_ID : curr_src_origin_idx;
    n_ptr->parent = TERMINAL;
    n_ptr->TS = TIME + 1;
    n_ptr->DIST = 1;
  } else {
    // set free
    n_ptr->parent = NULL;
    n_ptr->src_origin_idx = INVALID_SRC_ID;

    // also mark node so it comes in active list
    //mark_node(n_ptr - nodes);
  }

  // get the new src/sink value
  int new_is_sink = -1;
  if (n_ptr->parent != NULL && n_ptr->parent != ORPHAN)
    new_is_sink = n_ptr->is_sink;

  // if the src/sink value change, then mark itself and all the neighboring
  // nodes
  if (new_is_sink != old_is_sink) {
    //mark_node(n_ptr - nodes);
    add_to_changed_list(n_ptr);

    // mark all neighboring nodes with different src/sink id as active too
   /* for (arc* a = n_ptr->first; a; a = a->next) {
      int neigh_is_sink = -1;
      if (a->head->parent != NULL && a->head->parent != ORPHAN)
        neigh_is_sink = a->head->is_sink;

      if (neigh_is_sink != new_is_sink)
        mark_node(a->head - nodes);
    }*/
  }
}

template<typename captype, typename tcaptype, typename flowtype>
  void Graph<captype, tcaptype, flowtype>::normalize_tree(NodeEditSet& tree_nodes)
{
  node* n_ptr;
  node* tree_ptr;
  captype up_rcap, dwn_rcap;

  // BFS over all trees
  for (n_ptr = nodes; n_ptr < node_last; ++n_ptr) {
    // check if
    bool is_root = (n_ptr->parent == TERMINAL || n_ptr->parent == NULL);

    // skip node if its not a terminal node
    if (is_root) {
      // if src root node, then explore the whole tree
      std::queue<node*> tree_queue;

      tree_queue.push(n_ptr);

      // BFS over whole tree
      while (!tree_queue.empty()) {
        tree_ptr = tree_queue.front();
        tree_queue.pop();

        bool node_normalized = false;

        // if it was suggested that the node was changed
        if (tree_nodes.count(tree_ptr) != 0) {
          normalize_node(tree_ptr);
          node_normalized = true;

          tree_nodes.erase(tree_ptr);
        }

        // iterate over all the nodes connected to this node, and add children
        for (arc* a=tree_ptr->first; a; a=a->next) {
          if (a->head->parent != TERMINAL && a->head->parent != ORPHAN &&
              a->head->parent != NULL && a->head->parent->head == tree_ptr) {
            // if the current arc points to a node which is a child
            tree_queue.push(a->head);

            if (node_normalized) {
              tree_nodes.insert(a->head);
            }
          }
        }
      }
    }
  }

  // mark the relevant nodes active
  set_active_normalization();
}

template<typename captype, typename tcaptype, typename flowtype>
  void Graph<captype, tcaptype, flowtype>::update_capacities(
                                            const tcaptype* const prv_cap_s,
                                            const tcaptype* const prv_cap_t,
                                            const tcaptype* const new_cap_s,
                                            const tcaptype* const new_cap_t,
                                            const bool do_normalize_tree)
{
  NodeEditSet tree_nodes;

  /* update unary capacity edges (t-links) */
  for (node_id var_i = 0; var_i < node_num; ++var_i) {
    // the previous unary capacities
    tcaptype old_cap_diff = prv_cap_s[var_i] - prv_cap_t[var_i];

    // the new unary capacities
    tcaptype new_cap_diff = new_cap_s[var_i] - new_cap_t[var_i];

    // reparameterize the unary
    reparam_unary(var_i, old_cap_diff, new_cap_diff);

    if (do_normalize_tree) {
      if ((nodes[var_i].tr_cap > 0 && nodes[var_i].is_sink) ||
          (nodes[var_i].tr_cap < 0 && !nodes[var_i].is_sink)) {
        tree_nodes.insert(nodes + var_i);
        add_to_changed_list(nodes + var_i);
      }
    } else {
      mark_node(var_i);
    }
  }

  if (do_normalize_tree)
    normalize_tree(tree_nodes);
}

template<typename captype, typename tcaptype, typename flowtype>
  bool Graph<captype, tcaptype, flowtype>::check_tree_integrity(const std::string& filename) const
{
  //std::cout << "Checking integrity of trees [should be run after maxflow() "
  //          << "or transform_seed_trees()] - " << filename << std::endl;

  std::ofstream tree_out;
  if (!filename.empty()) tree_out.open(filename.c_str(), std::ios::trunc);

  /* vector to check if each node is accounted for */
  std::vector<int> nodes_checked(node_num, 0);
  int num_nodes_checked = 0;
  tcaptype trcap;
  node_id curr_i;

  /* iterate over all nodes */
  for (node_id i = 0; i < node_num; ++i) {
    node* n_ptr = nodes + i;
    const int curr_src_idx = n_ptr->src_origin_idx;

    // skip node if its not a terminal node
    if (n_ptr->parent == TERMINAL) {
      if (!n_ptr->is_sink) {
        // if src tree
        if (!filename.empty()) tree_out << "Source (" << n_ptr->src_origin_idx << ") rooted at ";

        // if src root node, then explore the whole tree
        std::queue<node*> src_nodes;

        curr_i = n_ptr - nodes;
        src_nodes.push(n_ptr);

        // this node shouldn't have been encountered before
        if (nodes_checked[curr_i] != 0)
          std::cerr << "I have encountered node " << curr_i << " before"
                    << std::endl;
        // check if the src_idx is a valid one
        if (curr_src_idx == INVALID_SRC_ID) {
          std::cerr << curr_i << " should be part of some source "
                    << " tree, not " << INVALID_SRC_ID << std::endl;
          nodes_checked[curr_i] = -1;
        }

        // BFS over whole tree
        while (!src_nodes.empty()) {
          n_ptr = src_nodes.front();
          curr_i = n_ptr - nodes;
          src_nodes.pop();

          if (!filename.empty()) tree_out << curr_i << ", ";

          // iterate over all the arcs of this node and search for children
          for (arc* a=n_ptr->first; a; a=a->next) {
            if (a->head->parent != TERMINAL && a->head->parent != ORPHAN &&
                a->head->parent != NULL && a->head->parent->head == n_ptr) {
              // if the current arc points to a node which is a child
              node_id conn_i = (node_id)(a->head - nodes);
              src_nodes.push(a->head);

              // this node shouldn't have been encountered before
              if (nodes_checked[conn_i] != 0)
                std::cerr << "I have encountered node " << conn_i << " before"
                          << std::endl;

              // there should be positive residual on the edge connecting
              // to the parent
              if (a->r_cap <= 0) {
                std::cerr << curr_i << "->" << conn_i << " needs to have "
                          << "positive residual capacity because the edge is"
                          << " part of the source tree (" << a->r_cap << ")"
                          << std::endl;
                nodes_checked[conn_i] = -1;
              }
            }
          }

          trcap = get_trcap(curr_i);
          // check if the node connecting to terminal has positive tr cap
          if (trcap < 0) {
            std::cerr << curr_i << " cannot have residual capacity to the sink "
                      << "(" << trcap << ") since it is part of the src tree"
                      << std::endl;
            nodes_checked[curr_i] = -1;
          }
          // check the id of the src origin idx
          if (n_ptr->src_origin_idx != curr_src_idx) {
            std::cerr << curr_i << " should be part of the " << curr_src_idx
                      << " source tree, not " << n_ptr->src_origin_idx << std::endl;
            nodes_checked[curr_i] = -1;
          }
          // check if the node is not sink
          if (n_ptr->is_sink) {
            std::cerr << curr_i << " should not be part of the sink side of the cut"
                      << " because it is part of the src tree" << std::endl;
            nodes_checked[curr_i] = -1;
          }

          // if node still good, mark it as having good integrity
          if (nodes_checked[curr_i] == 0)
            nodes_checked[curr_i] = 1;

          ++num_nodes_checked;
        }

        if (!filename.empty()) tree_out << std::endl;
      } else {
        // if sink tree
        if (!filename.empty()) tree_out << "Sink rooted at ";

        // if sink root node, then explore the whole tree
        std::queue<node*> sink_nodes;

        curr_i = n_ptr - nodes;
        sink_nodes.push(n_ptr);

        // this node shouldn't have been encountered before
        if (nodes_checked[curr_i] != 0)
          std::cerr << "I have encountered node " << curr_i << " before"
                    << std::endl;
        // check if the src_idx is invalid since we are sink
        if (curr_src_idx != INVALID_SRC_ID) {
          std::cerr << curr_i << " should be part of sink tree, not "
                    << INVALID_SRC_ID << std::endl;
          nodes_checked[curr_i] = -1;
        }

        // BFS over whole tree
        while (!sink_nodes.empty()) {
          n_ptr = sink_nodes.front();
          curr_i = n_ptr - nodes;
          sink_nodes.pop();

          if (!filename.empty()) tree_out << curr_i << ", ";

          // iterate over all the arcs of this node and search for children
          for (arc* a=n_ptr->first; a; a=a->next) {
            if (a->head->parent != TERMINAL && a->head->parent != ORPHAN &&
                a->head->parent != NULL && a->head->parent->head == n_ptr) {
              // if the current arc points to a node which is a child
              node_id conn_i = (node_id)(a->head-nodes);
              sink_nodes.push(a->head);

              // this node shouldn't have been encountered before
              if (nodes_checked[conn_i] != 0)
                std::cerr << "I have encountered node " << conn_i << " before"
                          << std::endl;

              // there should be positive residual on the edge connecting
              // to the parent
              if (a->sister->r_cap <= 0) {
                std::cerr << conn_i << "->" << curr_i << " needs to have "
                          << "positive residual capacity because the edge is"
                          << " part of the sink tree (" << a->r_cap << ")"
                          << std::endl;
                nodes_checked[conn_i] = -1;
              }
            }
          }

          trcap = get_trcap(curr_i);
          // check if the node connecting to terminal has negative tr cap
          if (trcap > 0) {
            std::cerr << curr_i << " cannot have residual capacity to the src "
                      << "(" << trcap << ") since it is part of the sink tree"
                      << std::endl;
            nodes_checked[curr_i] = -1;
          }
          // check the id of the src origin idx
          if (n_ptr->src_origin_idx != curr_src_idx) {
            std::cerr << curr_i << " should be part of the " << curr_src_idx
                      << " sink tree, not " << n_ptr->src_origin_idx << std::endl;
            nodes_checked[curr_i] = -1;
          }
          // check if the node is not sink
          if (!n_ptr->is_sink) {
            std::cerr << curr_i << " should not be part of the src side of the cut"
                      << " because it is part of the sink tree" << std::endl;
            nodes_checked[curr_i] = -1;
          }

          // if node still good, mark it as having good integrity
          if (nodes_checked[curr_i] == 0)
            nodes_checked[curr_i] = 1;

          ++num_nodes_checked;
        }

        if (!filename.empty()) tree_out << std::endl;
      }
    }
  }

  /* all remaining nodes which are not part of some tree should be free */
  for (node_id i=0; i < node_num; ++i) {
    node* n_ptr = nodes + i;

    if (nodes_checked[i] == 0) {
      nodes_checked[i] = 1;

      if (n_ptr->parent != NULL) {
        std::cerr << i << " should be a free node because it is not part of "
                  << "any tree" << std::endl;
        nodes_checked[i] = -1;
      }

      tcaptype trcap = get_trcap(i);
      if (trcap != 0) {
        std::cerr << i << " is a free node at the end of maxflow so should "
                  << "have zero t residual capacity (" << trcap << ")"
                  << std::endl;
        nodes_checked[i] = -1;
      }

      // check if the src_idx is invalid since we are free
      if (n_ptr->src_origin_idx != INVALID_SRC_ID) {
        std::cerr << curr_i << " should be free, not part of tree "
                  << INVALID_SRC_ID << std::endl;
        nodes_checked[curr_i] = -1;
      }

      if (!filename.empty()) tree_out << "Free node " << i << std::endl;

      ++num_nodes_checked;
    }
  }

  int nodes_wrong = 0;
  for (node_id i=0; i < node_num; ++i) {
    if (nodes_checked[i] == -1)
      ++nodes_wrong;
  }

  if (!filename.empty()) tree_out.close();

  test_consistency();

  gassert(nodes_wrong == 0, "errors while checking tree integrity");

  return (nodes_wrong == 0);
}

#endif


#pragma once

#include <string>
#include <memory>
#include <string>
#include <armadillo>

class Config;
class Physics;
class BatchTracking;
class Pipeline;
class GoverningEquationSolverBase;
class BoundaryConditions;
class TimeStep;

/*!
 * \brief The Solver class combines GoverningEquationSolver and BatchTracking
 * to advance the governing equations forward in time, and advect the
 * gas composition.
 *
 * It operates on Pipeline instances, and returns a copy of the input, with
 * updated flow, pressure, temperature, and composition (if batch tracking is
 * enabled).
 */
class Solver
{
public:
    //! Declared to avoid the inline compiler-generated default destructor.
    virtual ~Solver();

    /*!
     * \brief Construct from Config.
     * \param nGridPoints Number of grid points
     * \param config Config instance
     */
    Solver(
            const arma::uword nGridPoints,
            const Config& config);

    /*!
     * \brief Construct from number of grid points and string to select
     * discretizer.
     * \param nGridPoints Number of grid points.
     * \param energyEquation Type of energy equation ("InternalEnergy" or "Enthalpy")
     * \param relaxationFactors Relaxation factors, typically around 1.
     * \param toleranceType Tolerance type ("absolute" or "relative")
     * \param tolerances Convergence criteria
     * \param bruteForce Will always do maxIterations iterations and not check
     * for convergence
     * \param maxIterations The maximum number of iterations to perform
     */
    explicit Solver(
            const arma::uword nGridPoints,
            const std::string& energyEquation = "InternalEnergy",
            const arma::vec& relaxationFactors = {1, 1, 2/3.0},
            const std::string& toleranceType = "relative",
            const arma::vec& tolerances = {0.001, 0.001, 0.001},
            const bool bruteForce = false,
            const arma::uword maxIterations = 200);

    /*!
     * \brief Solve the governing equations.
     *
     * \see solveWithIterations()
     *
     * \param dt Time step [s]
     * \param current Current pipeline state
     * \param boundaryConditions Boundary conditions with time stamp
     * \param physics Physics instance to update derived properties and heat transfer
     * \return Copy of current with updated properties.
     */
    virtual Pipeline solve(
            const arma::uword dt,
            const Pipeline& current,
            const TimeStep& boundaryConditions,
            const Physics& physics) const;

    /*!
     * \brief Solve the governing equations.
     *
     * \see solveWithIterations()
     *
     * \param dt Time step [s]
     * \param current Current pipeline state
     * \param boundaryConditions Boundary conditions
     * \param physics Physics instance to update derived properties and heat transfer
     * \return Copy of current with updated properties.
     */
    virtual Pipeline solve(
            const arma::uword dt,
            const Pipeline& current,
            const BoundaryConditions& boundaryConditions,
            const Physics& physics) const;

    /*!
     * \brief Check if differences between an old and a new state are within
     * tolerances. This will either compare relative or absolute differences.
     *
     * This is used to see if the iteration over solution of the governing
     * equations have converged, by comparing differences between the two last
     * iterations.
     *
     * In the case of relative differences, we can run into issues if we have
     * zero flow, so this is handled explicitly (temperature and pressure is
     * never expected to reach zero in well-behaved simulations).
     *
     * \param guess
     * \param previous
     * \param tolerances
     * \param toleranceType
     * \param relaxationFactors
     * \return
     */
    static arma::uword differencesWithinTolerance(
            const Pipeline& guess,
            const Pipeline& previous,
            const arma::vec& tolerances,
            const std::string& toleranceType,
            const arma::vec& relaxationFactors);

    /*!
     * \brief Solve the governing equations.
     *
     * TODO: Some more documentation here
     *
     * \param dt Time step [s]
     * \param current Pipeline state
     * \param boundaryConditions Boundary conditions
     * \param physics Physics instance to update derived properties and heat transfer
     * \return Copy of current with updated properties.
     */
    Pipeline solveWithIterations(
            const arma::uword dt,
            const Pipeline& current,
            const BoundaryConditions& boundaryConditions,
            const Physics& physics) const;

    //! Enable brute force solver, which always does m_maxIterations iterations.
    void enableBruteForce();

    //! Set max iterations.
    void setMaxIterations(const arma::uword maxIterations);

    //! Return the number of iterations performed during previous solution attempt
    arma::uword nIterations() const { return m_nIterations; }

    //! Get (const ref) relaxation factors
    const arma::vec& relaxationFactors() const { return m_relaxationFactor;}
    //! Get (const ref) tolerance type
    const std::string& toleranceType() const { return m_toleranceType; }
    //! Get (const ref) tolerances
    const arma::vec& tolerances() const { return m_tolerances; }
    //! Get (const ref) GoverningEquationSolver
    const GoverningEquationSolverBase& governingEquationSolver() const { return *m_governingEquationSolver; }

private:
    //! Relaxation factors for each property (flow, pressure and temperature).
    arma::vec m_relaxationFactor;

    //! What kind of tolerance to use when checking convergence ("relative" or
    //! "absolute").
    std::string m_toleranceType;

    //! Convergence tolerances for each property (flow, pressure and temperature).
    arma::vec m_tolerances;

//    arma::uword m_maxNumberOfIterations = 20; // not used
//    bool m_hardLimitOnNumberOfIterations = true; // not used
//    bool m_bruteForce = false; // not used

    //! If we should use brute-force solving, which means not checking for
    //! convergence, but always performing m_maxIterations iterations.
    bool m_bruteForce;
    //! The maximum number of iterations to perform, or (if m_bruteForce is
    //! true) the exact number of iterations to perform.
    arma::uword m_maxIterations;
    //! The number of iterations performed during the previous solution attempt.
    //! This is mutable, and is updated in Solver::solve().
    mutable arma::uword m_nIterations = 0;

    //! GoverningEquationSolver instance used for solving the governing equations.
    std::unique_ptr<GoverningEquationSolverBase> m_governingEquationSolver;

    //! BatchTracking instance used to advect composition in the pipeline.
    std::unique_ptr<BatchTracking> m_compositionSolver;

    /*!
     * \brief Private method for making GoverningEquationSolverBase instance
     * from Config and nGridPoints.
     * \param config Config instance
     * \param nGridPoints Number of grid points
     * \return unique_ptr to GoverningEquationSolverBase
     */
    std::unique_ptr<GoverningEquationSolverBase> makeGoverningEquationSolver(
            const arma::uword nGridPoints,
            const Config& config);

    /*!
     * \brief Private method for making GoverningEquationSolverBase from string
     * and nGridPoints.
     * \param discretizer Discretizer type ("InternalEnergy" or "Enthalpy")
     * \param nGridPoints Number of grid points
     * \return unique_ptr to GoverningEquationSolverBase
     */
    std::unique_ptr<GoverningEquationSolverBase> makeGoverningEquationSolver(
            const arma::uword nGridPoints,
            const std::string& discretizer);
};