/*!
 * Copyright (c) <2023> <Andreas Niskanen, University of Helsinki>
 * 
 * 
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * 
 * 
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * 
 * 
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#ifdef DEBUG
#include <iostream>
#endif

#include <cstdlib>

#include "mu-toksia.h"
#include "Encodings.h"
#include "ipasir.h"

using namespace std;
using namespace Encodings;

AFSolver::AFSolver()
{
	sat_solver = nullptr;
	solver_encoding = UNKNOWN_SEM;
	current_state = INPUT;
	enumerate = false;
}

AFSolver::~AFSolver()
{
	if (sat_solver) ipasir_release(sat_solver);
}

void AFSolver::set_semantics(sigma s)
{
	dynamic_af.sem = s;
}

void AFSolver::add_argument(int32_t arg)
{
	bool success = dynamic_af.add_argument(arg);
	if (!success) current_state = ERROR;
}

void AFSolver::del_argument(int32_t arg)
{
	bool success = dynamic_af.del_argument(arg);
	if (!success) current_state = ERROR;
}

void AFSolver::add_attack(int32_t source, int32_t target)
{
	bool success = dynamic_af.add_attack(source, target);
	if (!success) current_state = ERROR;
}

void AFSolver::del_attack(int32_t source, int32_t target)
{
	bool success = dynamic_af.del_attack(source, target);
	if (!success) current_state = ERROR;
}

void AFSolver::assume_in(int32_t arg)
{
	if (dynamic_af.arg_to_int.count(abs(arg)) == 0) {
		current_state = ERROR;
		return;
	}
	int32_t arg_index = dynamic_af.arg_to_int.at(abs(arg));
	int32_t var = dynamic_af.accepted_var[arg_index];
	status_assumptions.push_back(arg > 0 ? var : -var);
}

void AFSolver::assume_out(int32_t arg)
{
	if (dynamic_af.arg_to_int.count(abs(arg)) == 0) {
		current_state = ERROR;
		return;
	}
	int32_t arg_index = dynamic_af.arg_to_int.at(abs(arg));
	int32_t var = dynamic_af.rejected_var[arg_index];
	status_assumptions.push_back(arg > 0 ? var : -var);
}

void AFSolver::assume_current_structure(void * solver)
{
	if (dynamic_af.static_mode) return;
	for (uint32_t i = 0; i < dynamic_af.args; i++) {
		if (dynamic_af.arg_exists[i]) {
			assume_lit(dynamic_af, solver, dynamic_af.arg_exists_var[i]);
			assume_lit(dynamic_af, solver, dynamic_af.attackers_var[i]);
		}
	}
}

void AFSolver::get_assignment(void * solver)
{
	assignment.clear();
	assignment.resize(dynamic_af.count+1);
#ifdef DEBUG
	cout << "ipasir_val ";
#endif
	for (uint32_t i = 1; i <= dynamic_af.count; i++) {
		assignment[i] = ipasir_val(solver, i);
#ifdef DEBUG
		cout << (assignment[i] > 0 ? "+" : "-") << dynamic_af.var_to_str[abs(assignment[i])] << " ";;
#endif
	}
#ifdef DEBUG
	cout << "\n";
#endif
}

void AFSolver::assume_grounded(void * solver)
{
	return;
	/*
	void * propagator = new ipasir_init();
	Encodings::add_complete(dynamic_af, propagator);
	assume_current_structure(propagator);
	vector<int32_t> grounded;
	ipasir_propagate();
	for (uint32_t i = 0; i < grounded.size(); i++)
		assume_lit(dynamic_af, solver, grounded[i]);
	if (propagator) delete propagator;
	*/
}

void AFSolver::assume_range(void * solver)
{
	for (uint32_t i = 0; i < dynamic_af.args; i++) {
		if (dynamic_af.arg_exists[i]) {
			assume_lit(dynamic_af, solver, dynamic_af.range_var[i]);
		}
	}
}

void AFSolver::set_objective_vars()
{
	objective_vars.clear();
	objective_vars.reserve(dynamic_af.args);
	if (dynamic_af.sem == PR) {
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if (dynamic_af.arg_exists[i])
				objective_vars.push_back(dynamic_af.accepted_var[i]);
		}
	} else if (dynamic_af.sem == SST || dynamic_af.sem == STG) {
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if (dynamic_af.arg_exists[i])
				objective_vars.push_back(dynamic_af.range_var[i]);
		}
	}
}

bool AFSolver::check_extensions(int32_t target, int32_t select)
{
	return false;
	/*for (auto it = extensions.rbegin(); it != extensions.rend(); ++it) {
		assume_current_structure(sat_solver);
		if (target) assume_lit(dynamic_af, sat_solver, target);
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if ((*it)[i]) assume_lit(dynamic_af, sat_solver, dynamic_af.accepted_var[i]);
			else assume_lit(dynamic_af, sat_solver, -dynamic_af.accepted_var[i]);
		}
		if (ipasir_solve(sat_solver) == 10) {
			get_assignment(sat_solver);
			if (dynamic_af.sem == CO || dynamic_af.sem == ST) {
				return true;
			} else if (dynamic_af.sem == PR || dynamic_af.sem == SST || dynamic_af.sem == STG) {
				assume_current_structure(sat_solver);
				if (select) assume_lit(dynamic_af, sat_solver, -select);
				vector<int> complement_clause;
				if (select) complement_clause.push_back(select);
				for (uint32_t i = 0; i < objective_vars.size(); i++) {
					if (assignment[objective_vars[i]] >= 0) {
						assume_lit(dynamic_af, sat_solver, objective_vars[i]);
					} else {
						complement_clause.push_back(objective_vars[i]);
					}
				}
				add_clause(dynamic_af, sat_solver, complement_clause);
				int superset_exists = ipasir_solve(sat_solver);
				if (superset_exists == 20) return true;
			}
		}
	}
	return false;*/
}

bool AFSolver::cegar(int32_t target, int32_t select)
{
	while (true) {
		assume_current_structure(sat_solver);
		if (select) assume_lit(dynamic_af, sat_solver, -select);
		if (target) assume_lit(dynamic_af, sat_solver, target);
		int abstraction = ipasir_solve(sat_solver);
		if (abstraction == 20) break;
		get_assignment(sat_solver);

		while (true) {
			assume_current_structure(sat_solver);
			if (select) assume_lit(dynamic_af, sat_solver, -select);
			if (target) assume_lit(dynamic_af, sat_solver, target);
			vector<int> complement_clause;
			if (select) complement_clause.push_back(select);
			for (uint32_t i = 0; i < objective_vars.size(); i++) {
				if (assignment[objective_vars[i]] >= 0) {
					assume_lit(dynamic_af, sat_solver, objective_vars[i]);
				} else {
					complement_clause.push_back(objective_vars[i]);
				}
			}
			add_clause(dynamic_af, sat_solver, complement_clause);
			int maximize = ipasir_solve(sat_solver);
			if (maximize == 20) {
				if (target) break;
				//extensions.push_back(assignment);
				return true;
			}
			get_assignment(sat_solver);
		}

		assume_current_structure(sat_solver);
		if (select) assume_lit(dynamic_af, sat_solver, -select);
		for (uint32_t i = 0; i < objective_vars.size(); i++) {
			if (assignment[objective_vars[i]] >= 0) {
				assume_lit(dynamic_af, sat_solver, objective_vars[i]);
			}
		}
		if (ipasir_solve(sat_solver) == 20) {
			//extensions.push_back(assignment);
			return true;
		}
	}
	return false;
}

bool AFSolver::ideal(int32_t target, int32_t select) {

	vector<int> accepted_clause;
	vector<int> rejected_clause;
	vector<uint8_t> union_of_accepted(dynamic_af.args, 0);
	vector<uint8_t> union_of_rejected(dynamic_af.args, 0);
	bool solved = false;

	while (true) {
		assume_current_structure(sat_solver);
		if (!solved && target) assume_lit(dynamic_af, sat_solver, target);
		if (solved && select) assume_lit(dynamic_af, sat_solver, -select);
		int sat = ipasir_solve(sat_solver);
		if (sat == 20) {
			if (!solved && target) return false;
			break;
		}
		get_assignment(sat_solver);
		solved = true;
		accepted_clause.clear();
		rejected_clause.clear();
		if (select) accepted_clause.push_back(select);
		if (select) rejected_clause.push_back(select);
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if (!dynamic_af.arg_exists[i]) continue;
			if (assignment[dynamic_af.accepted_var[i]] >= 0) {
				union_of_accepted[i] = 1;
			} else if (!union_of_accepted[i]) {
				accepted_clause.push_back(dynamic_af.accepted_var[i]);
			}
			if (assignment[dynamic_af.rejected_var[i]] >= 0) {
				union_of_rejected[i] = 1;
			} else if (!union_of_rejected[i]) {
				rejected_clause.push_back(dynamic_af.rejected_var[i]);
			}
		}
		add_clause(dynamic_af, sat_solver, accepted_clause);
		add_clause(dynamic_af, sat_solver, rejected_clause);
	}

	objective_vars.clear();
	objective_vars.reserve(dynamic_af.args);
	for (uint32_t i = 0; i < dynamic_af.args; i++) {
		if (!dynamic_af.arg_exists[i]) continue;
		if (!union_of_rejected[i]) {
			objective_vars.push_back(dynamic_af.accepted_var[i]);
		} else {
			vector<int> clause = { -dynamic_af.accepted_var[i] };
			if (!dynamic_af.static_mode) clause.push_back(dynamic_af.count+1);
			add_clause(dynamic_af, sat_solver, clause);
		}
	}
	return cegar(target, (dynamic_af.static_mode ? 0 : ++dynamic_af.count));
}

void AFSolver::solve(bool cred)
{
	if (current_state == ERROR) return;
	if ((dynamic_af.static_mode && !enumerate) || dynamic_af.sem != solver_encoding) {
		if (sat_solver) ipasir_release(sat_solver);
		dynamic_af.set_not_encoded();
		sat_solver = ipasir_init();
	}

	if (!dynamic_af.static_mode) {
		dynamic_af.attackers_var.resize(dynamic_af.args, 0);
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if (dynamic_af.arg_encoded[i] && !dynamic_af.arg_exists[i]) {
				vector<int> clause = { -dynamic_af.arg_exists_var[i] };
				add_clause(dynamic_af, sat_solver, clause);
				clause = { -dynamic_af.accepted_var[i] };
				add_clause(dynamic_af, sat_solver, clause);
				clause = { -dynamic_af.rejected_var[i] };
				add_clause(dynamic_af, sat_solver, clause);
				if (dynamic_af.attackers_var[i]) {
					clause = { -dynamic_af.attackers_var[i] };
					add_clause(dynamic_af, sat_solver, clause);
				}
				dynamic_af.arg_encoded[i] = false;
				continue;
			}
			if (dynamic_af.atts_encoded[i]) continue;
			if (dynamic_af.attackers_var[i]) {
				vector<int> clause = { -dynamic_af.attackers_var[i] };
				add_clause(dynamic_af, sat_solver, clause);
			}
			dynamic_af.attackers_var[i] = ++dynamic_af.count;
#ifdef DEBUG
			dynamic_af.var_to_str[dynamic_af.count] = "attackers(" + to_string(dynamic_af.int_to_arg[i]) + ")";
#endif
		}
	}

	if (!enumerate) {
		if (dynamic_af.sem == AD) {
			add_admissible(dynamic_af, sat_solver);
		} else if (dynamic_af.sem == CO) {
			add_complete(dynamic_af, sat_solver);
		} else if (dynamic_af.sem == PR) {
	#if defined(CO_IN_PR)
			add_complete(dynamic_af, sat_solver);
	#else
			add_admissible(dynamic_af, sat_solver);
	#endif
		} else if (dynamic_af.sem == ST) {
			add_stable(dynamic_af, sat_solver);
		} else if (dynamic_af.sem == SST) {
	#if defined(CO_IN_SST)
			add_complete(dynamic_af, sat_solver);
	#else
			add_admissible(dynamic_af, sat_solver);
	#endif
			add_range(dynamic_af, sat_solver);
		} else if (dynamic_af.sem == STG) {
			add_conflict_free(dynamic_af, sat_solver);
			add_range(dynamic_af, sat_solver);
		} else if (dynamic_af.sem == ID) {
	#if defined(CO_IN_ID)
			add_complete(dynamic_af, sat_solver);
	#else
			add_admissible(dynamic_af, sat_solver);
	#endif
		} else {
			current_state = ERROR;
			return;
		}
		solver_encoding = dynamic_af.sem;
		dynamic_af.set_encoded();
	}

	int32_t target = add_target(dynamic_af, status_assumptions, sat_solver);
	status_assumptions.clear();
	if (!cred) target = -target;
	int32_t select = 0;
	if (!dynamic_af.static_mode && (dynamic_af.sem == PR || dynamic_af.sem == SST || dynamic_af.sem == STG || dynamic_af.sem == ID)) {
		select = ++dynamic_af.count;
#ifdef DEBUG
		if (!dynamic_af.static_mode) dynamic_af.var_to_str[dynamic_af.count] = "select";
#endif
	}

	if (dynamic_af.sem == PR || dynamic_af.sem == SST || dynamic_af.sem == STG)
		set_objective_vars();

	/*if (extensions.size() > MAX_EXTENSIONS) extensions.erase(extensions.begin());
	if (dynamic_af.sem != ID && check_extensions(target, select)) {
		current_state = (cred ? ACCEPT : REJECT);
		return;
	}*/

	if (dynamic_af.sem == AD || dynamic_af.sem == CO) {
		assume_current_structure(sat_solver);
		if (target) assume_lit(dynamic_af, sat_solver, target);
		if (ipasir_solve(sat_solver) == 10) {
			get_assignment(sat_solver);
			current_state = (cred ? ACCEPT : REJECT);
			//extensions.push_back(assignment);
		} else current_state = (cred ? REJECT : ACCEPT);

	} else if (dynamic_af.sem == ST) {

		assume_current_structure(sat_solver);
#if defined(GR_IN_ST)
		assume_grounded(sat_solver);
#endif
		if (target) assume_lit(dynamic_af, sat_solver, target);
		if (ipasir_solve(sat_solver) == 10) {
			get_assignment(sat_solver);
			current_state = (cred ? ACCEPT : REJECT);
			//extensions.push_back(assignment);
		} else current_state = (cred ? REJECT : ACCEPT);

	} else if (dynamic_af.sem == PR) {

		if (cegar(target, select))
			current_state = (cred ? ACCEPT : REJECT);
		else current_state = (cred ? REJECT : ACCEPT);

	} else if (dynamic_af.sem == SST) {

#if defined(ST_EXISTS_SST)
		assume_current_structure(sat_solver);
		assume_range(sat_solver);
		if (ipasir_solve(sat_solver) == 10) {
			assume_current_structure(sat_solver);
#if defined(GR_IN_ST)
			assume_grounded(sat_solver);
#endif		
			assume_range(sat_solver);
			if (target) assume_lit(dynamic_af, sat_solver, target);
			if (ipasir_solve(sat_solver) == 10) {
				get_assignment(sat_solver);
				current_state = (cred ? ACCEPT : REJECT);
			} else current_state = (cred ? REJECT : ACCEPT);
			return;
		}
#endif

		if (cegar(target, select))
			current_state = (cred ? ACCEPT : REJECT);
		else current_state = (cred ? REJECT : ACCEPT);

	} else if (dynamic_af.sem == STG) {

#if defined(ST_EXISTS_STG)
		assume_current_structure(sat_solver);
		assume_range(sat_solver);
		if (ipasir_solve(sat_solver) == 10) {
			assume_current_structure(sat_solver);
#if defined(GR_IN_ST)
			assume_grounded(sat_solver);
#endif		
			assume_range(sat_solver);
			if (target) assume_lit(dynamic_af, sat_solver, target);
			if (ipasir_solve(sat_solver) == 10) {
				get_assignment(sat_solver);
				current_state = (cred ? ACCEPT : REJECT);
			} else current_state = (cred ? REJECT : ACCEPT);
			return;
		}
#endif

		if (cegar(target, select))
			current_state = (cred ? ACCEPT : REJECT);
		else current_state = (cred ? REJECT : ACCEPT);

	} else if (dynamic_af.sem == ID) {

		if (ideal(target, select ? select : ++dynamic_af.count))
			current_state = (cred ? ACCEPT : REJECT);
		else current_state = (cred ? REJECT : ACCEPT);

	} else {
		current_state = ERROR;
	}

	if (select) {
		vector<int> clause = { select };
		add_clause(dynamic_af, sat_solver, clause);
	}
}

// TODO: experimental
// note: ownership of range_solver is external
void AFSolver::solve_range(void * range_solver, bool initialized)
{
	if (current_state == ERROR) return;
	if (dynamic_af.sem != SST && dynamic_af.sem != STG) {
		current_state = ERROR;
		return;
	}
	if (!range_solver || !dynamic_af.static_mode) {
		current_state = ERROR;
		return;
	}

	if (!initialized) {
		if (dynamic_af.sem == SST) {
			#if defined(CO_IN_SST)
			add_complete(dynamic_af, range_solver);
			#else
			add_admissible(dynamic_af, range_solver);
			#endif
		} else if (dynamic_af.sem == STG) {
			add_conflict_free(dynamic_af, range_solver);
		}
		add_range(dynamic_af, range_solver);
		for (uint32_t i = 0; i < dynamic_af.args; i++) {
			if (dynamic_af.arg_exists[i] && assignment[dynamic_af.range_var[i]] >= 0) {
				vector<int> clause = { dynamic_af.range_var[i] };
				add_clause(dynamic_af, range_solver, clause);
			}
		}
	}

	vector<int> clause;
	for (uint32_t i = 0; i < dynamic_af.args; i++) {
		if (dynamic_af.arg_exists[i]) {
			clause.push_back(assignment[dynamic_af.accepted_var[i]] > 0 ? -dynamic_af.accepted_var[i] : dynamic_af.accepted_var[i]);
		}
	}
	add_clause(dynamic_af, range_solver, clause);

	if (ipasir_solve(range_solver) == 10) {
		get_assignment(range_solver);
		current_state = ACCEPT;
	} else current_state = REJECT;
}

int32_t AFSolver::get_val_in(int32_t arg)
{
	if (arg < 0 || !dynamic_af.arg_to_int.count(arg)) {
		current_state = ERROR;
		return 0;
	}
	int32_t arg_index = dynamic_af.arg_to_int.at(arg);
	return assignment[dynamic_af.accepted_var[arg_index]] > 0 ? arg : -arg;
}

int32_t AFSolver::get_val_out(int32_t arg)
{
	if (arg < 0 || !dynamic_af.arg_to_int.count(arg)) {
		current_state = ERROR;
		return 0;
	}
	int32_t arg_index = dynamic_af.arg_to_int.at(arg);
	return assignment[dynamic_af.rejected_var[arg_index]] > 0 ? arg : -arg;
}

int32_t AFSolver::get_val_range(int32_t arg)
{
	if (arg < 0 || !dynamic_af.arg_to_int.count(arg)) {
		current_state = ERROR;
		return 0;
	}
	int32_t arg_index = dynamic_af.arg_to_int.at(arg);
	return assignment[dynamic_af.range_var[arg_index]] > 0 ? arg : -arg;
}

void AFSolver::block_solution()
{
	vector<int> clause;
	for (uint32_t i = 0; i < dynamic_af.args; i++) {
		if (dynamic_af.arg_exists[i]) {
			clause.push_back(assignment[dynamic_af.accepted_var[i]] > 0 ? -dynamic_af.accepted_var[i] : dynamic_af.accepted_var[i]);
		}
	}
	add_clause(dynamic_af, sat_solver, clause);
}