using namespace std;

#include "regression_search.h"

#include "globals.h"
#include "heuristic.h"
#include "option_parser.h"
#include "successor_generator.h"
#include "g_evaluator.h"
#include "sum_evaluator.h"
#include "plugin.h"

RegressionSearch::RegressionSearch(const Options &opts)
    : EagerSearch(opts),
      use_subsumption_check(opts.get<bool>("simple_subsumption")),
      use_subsumption_trie(opts.get<bool>("subsumption_trie")){}

void RegressionSearch::initialize() {
    //TODO children classes should output which kind of search
    cout << "Conducting best first search"
         << (reopen_closed_nodes ? " with" : " without")
         << " reopening closed nodes, (real) bound = " << bound
         << endl;
    assert(use_subsumption_check && use_subsumption_true == false);
    if (do_pathmax)
        cout << "Using pathmax correction" << endl;
    if (use_multi_path_dependence)
        cout << "Using multi-path dependence (LM-A*)" << endl;
    if (use_subsumption_check)
        cout << "Using subsumption check" << endl;
    if (use_subsumption_trie)
        cout << "Using subsumption trie" << endl;

    if (use_subsumption_trie) {
        subsumption_trie.attach_search_space(&search_space);
    }
    assert(open_list != NULL);

    set<Heuristic *> hset;
    open_list->get_involved_heuristics(hset);

    for (set<Heuristic *>::iterator it = hset.begin(); it != hset.end(); ++it) {
        estimate_heuristics.push_back(*it);
        search_progress.add_heuristic(*it);
    }

    // add heuristics that are used for preferred operators (in case they are
    // not also used in the open list)
    hset.insert(preferred_operator_heuristics.begin(),
                preferred_operator_heuristics.end());

    // add heuristics that are used in the f_evaluator. They are usually also
    // used in the open list and hence already be included, but we want to be
    // sure.
    if (f_evaluator) {
        f_evaluator->get_involved_heuristics(hset);
    }

    for (set<Heuristic *>::iterator it = hset.begin(); it != hset.end(); ++it) {
        heuristics.push_back(*it);
    }

    assert(!heuristics.empty());
    const GlobalState &initial_state = g_goal_state();
    for (size_t i = 0; i < heuristics.size(); ++i)
        heuristics[i]->evaluate(initial_state);
    open_list->evaluate(0, false);
    search_progress.inc_evaluated_states();
    search_progress.inc_evaluations(heuristics.size());

    if (open_list->is_dead_end()) {
        cout << "Initial state is a dead end." << endl;
    } else {
        search_progress.get_initial_h_values();
        if (f_evaluator) {
            f_evaluator->evaluate(0, false);
            search_progress.report_f_value(f_evaluator->get_value());
        }
        search_progress.check_h_progress(0);
        SearchNode node = search_space.get_node(initial_state);
        node.open_initial(heuristics[0]->get_value());

        open_list->insert(initial_state.get_id());
    }
}

void RegressionSearch::generate_applicable_ops_regression(
    const GlobalState &s, vector<const GlobalOperator *> &applicable_ops) {
    for (auto& op : g_operators) {
        bool applicable = false;
	
        // We need to save those for checking the conditions later.
        vector<int> vals_with_effects;
        for (auto effect : op.get_effects()) {
            vals_with_effects.push_back(effect.var);
            if (s[effect.var] == g_variable_domain[effect.var] - 1) {
                continue;
            }
            if (s[effect.var] == effect.val) {
                applicable = true;
            } else {
                applicable = false;
                break;
            }
        }
        if (applicable) {
            for (auto condition : op.get_preconditions()) {
                bool has_effect = false;
                for (auto var : vals_with_effects) {
                    if (var == condition.var) {
                        has_effect = true;
                        break;
                    }
                }
                if (!has_effect) {
                    if (s[condition.var] != g_variable_domain[condition.var] - 1
                        && s[condition.var] != condition.val) {
                        applicable = false;
                        break;
                    }
                }       
            }
        }
        if (applicable) {
            applicable_ops.push_back(&op);
        }
    }
}

bool RegressionSearch::check_goal_and_set_plan_regression(const GlobalState &state) {
    bool is_init_state = false;
    const GlobalState &initial_state = g_state_registry->get_initial_state();
    for (size_t i = 0; i < g_variable_domain.size(); ++i) {
        if (state[i] == g_variable_domain[i] - 1) {
            continue;
        }
        if (state[i] == initial_state[i]) {
            is_init_state = true;
        } else {
            is_init_state = false;
            break;
        }
    }
    if (is_init_state) {
        cout << "Solution found!" << endl;
        Plan plan;
        search_space.trace_path(state, plan);
        reverse(plan.begin(), plan.end());
        set_plan(plan);
        return true;
    }
    return false;
}

SearchStatus RegressionSearch::step() {
    pair<SearchNode, bool> n = fetch_next_node();
    if (!n.second) {
        return FAILED;
    }
    SearchNode node = n.first;

    GlobalState s = node.get_state();
    
    if (check_goal_and_set_plan_regression(s))
        return SOLVED;
    
    vector<const GlobalOperator *> applicable_ops;
    set<const GlobalOperator *> preferred_ops;

    generate_applicable_ops_regression(s, applicable_ops);
    // This evaluates the expanded state (again) to get preferred ops
    for (size_t i = 0; i < preferred_operator_heuristics.size(); ++i) {
        Heuristic *h = preferred_operator_heuristics[i];
        h->evaluate(s);
        if (!h->is_dead_end()) {
            // In an alternation search with unreliable heuristics, it is
            // possible that this heuristic considers the state a dead end.
            vector<const GlobalOperator *> preferred;
            h->get_preferred_operators(preferred);
            preferred_ops.insert(preferred.begin(), preferred.end());
        }
    }
    search_progress.inc_evaluations(preferred_operator_heuristics.size());
    if (use_subsumption_trie) {
        subsumption_trie.insert(s);
    }
    for (size_t i = 0; i < applicable_ops.size(); ++i) {
        const GlobalOperator *op = applicable_ops[i];
        if ((node.get_real_g() + op->get_cost()) >= bound)
            continue;
        
        GlobalState succ_state = g_state_registry->get_predecessor_state(s, *op);

        search_progress.inc_generated();
        bool is_preferred = (preferred_ops.find(op) != preferred_ops.end());
        
        SearchNode succ_node = search_space.get_node(succ_state);
        
        // Previously encountered dead end. Don't re-evaluate.
        if (succ_node.is_dead_end())
            continue;

        // update new path
        if (succ_node.is_new()) {
            bool h_is_dirty = false;
            for (size_t j = 0; j < heuristics.size(); ++j) {
                /*
                  Note that we can't break out of the loop when
                  h_is_dirty is set to true or use short-circuit
                  evaluation here. We must call reach_state for each
                  heuristic for its side effects.
                */
                if (heuristics[j]->reach_state(s, *op, succ_state))
                    h_is_dirty = true;
            }
            if (h_is_dirty)
                succ_node.set_h_dirty();
        }
        if (succ_node.is_new()) {
            // We have not seen this state before.
            // Evaluate and create a new node.
            for (size_t j = 0; j < heuristics.size(); ++j)
                heuristics[j]->evaluate(succ_state);
            succ_node.clear_h_dirty();
            search_progress.inc_evaluated_states();
            search_progress.inc_evaluations(heuristics.size());

            // Note that we cannot use succ_node.get_g() here as the
            // node is not yet open. Furthermore, we cannot open it
            // before having checked that we're not in a dead end. The
            // division of responsibilities is a bit tricky here -- we
            // may want to refactor this later.
            open_list->evaluate(node.get_g() + get_adjusted_cost(*op), is_preferred);
            bool dead_end = open_list->is_dead_end();
            if (dead_end) {
                succ_node.mark_as_dead_end();
                search_progress.inc_dead_ends();
                continue;
            }

            //TODO:CR - add an ID to each state, and then we can use a vector to save per-state information
            int succ_h = heuristics[0]->get_heuristic();
            if (do_pathmax) {
                if ((node.get_h() - get_adjusted_cost(*op)) > succ_h) {
                    //cout << "Pathmax correction: " << succ_h << " -> " << node.get_h() - get_adjusted_cost(*op) << endl;
                    succ_h = node.get_h() - get_adjusted_cost(*op);
                    heuristics[0]->set_evaluator_value(succ_h);
                    open_list->evaluate(node.get_g() + get_adjusted_cost(*op), is_preferred);
                    search_progress.inc_pathmax_corrections();
                }
            }
            succ_node.open(succ_h, node, op);

            if (use_subsumption_check || use_subsumption_trie) {
                bool subsuming = false;
                if (use_subsumption_check) {
                    subsuming = search_space.subsumption_check(succ_state);
                } else if (use_subsumption_trie) {
                    subsuming = subsumption_trie.lookup(succ_state);
                }
                if (subsuming) {
                    continue;
                }
            }
            open_list->insert(succ_state.get_id());
            if (search_progress.check_h_progress(succ_node.get_g())) {
                reward_progress();
            }
        } else if (succ_node.get_g() > node.get_g() + get_adjusted_cost(*op)) {
            // We found a new cheapest path to an open or closed state.
            if (reopen_closed_nodes) {
                //TODO:CR - test if we should add a reevaluate flag and if it helps
                // if we reopen closed nodes, do that
                if (succ_node.is_closed()) {
                    /* TODO: Verify that the heuristic is inconsistent.
                     * Otherwise, this is a bug. This is a serious
                     * assertion because it can show that a heuristic that
                     * was thought to be consistent isn't. Therefore, it
                     * should be present also in release builds, so don't
                     * use a plain assert. */
                    //TODO:CR - add a consistent flag to heuristics, and add an assert here based on it
                    search_progress.inc_reopened();
                }
                succ_node.reopen(node, op);
                heuristics[0]->set_evaluator_value(succ_node.get_h());
                // TODO: this appears fishy to me. Why is here only heuristic[0]
                // involved? Is this still feasible in the current version?
                open_list->evaluate(succ_node.get_g(), is_preferred);

                open_list->insert(succ_state.get_id());
            } else {
                // if we do not reopen closed nodes, we just update the parent pointers
                // Note that this could cause an incompatibility between
                // the g-value and the actual path that is traced back
                succ_node.update_parent(node, op);
            }
        }
    }

    return IN_PROGRESS;
}
	
static SearchEngine *_parse_regression(OptionParser &parser) {
	// TODO: Copied from astar, change stuff for regression
	parser.document_synopsis(
		"A* search (eager)",
		"A* is a special case of eager best first search that uses g+h "
		"as f-function. "
		"We break ties using the evaluator. Closed nodes are re-opened.");
	parser.document_note(
		"mpd option",
		"This option is currently only present for the A* algorithm and not "
		"for the more general eager search, "
		"because the current implementation of multi-path depedence "
		"does not support general open lists.");
	parser.document_note(
		"Equivalent statements using general eager search",
		"\n```\n--search astar(evaluator)\n```\n"
		"is equivalent to\n"
		"```\n--heuristic h=evaluator\n"
		"--search eager(tiebreaking([sum([g(), h]), h], unsafe_pruning=false),\n"
		"               reopen_closed=true, pathmax=false, progress_evaluator=sum([g(), h]))\n"
		"```\n", true);
	parser.add_option<ScalarEvaluator *>("eval", "evaluator for h-value");
	parser.add_option<bool>("pathmax",
							"use pathmax correction", "false");
	parser.add_option<bool>("mpd",
							"use multi-path dependence (LM-A*)", "false");
    parser.add_option<bool>("simple_subsumption", "use subsumption check", "false");
    parser.add_option<bool>("subsumption_trie", "use subsumption trie", "false");
	SearchEngine::add_options_to_parser(parser);
	Options opts = parser.parse();

	EagerSearch *engine = 0;
	if (!parser.dry_run()) {
		GEvaluator *g = new GEvaluator();
		vector<ScalarEvaluator *> sum_evals;
		sum_evals.push_back(g);
		ScalarEvaluator *eval = opts.get<ScalarEvaluator *>("eval");
		sum_evals.push_back(eval);
		ScalarEvaluator *f_eval = new SumEvaluator(sum_evals);

		// use eval for tiebreaking
		std::vector<ScalarEvaluator *> evals;
		evals.push_back(f_eval);
		evals.push_back(eval);
		OpenList<StateID> *open = \
			new TieBreakingOpenList<StateID>(evals, false, false);

		opts.set("open", open);
		opts.set("f_eval", f_eval);
		opts.set("reopen_closed", true);
		engine = new RegressionSearch(opts);
	}

	return engine;
}


static Plugin<SearchEngine> _plugin_regression("regression_search", _parse_regression);
