#include <anthem/Completion.h>

#include <anthem/ASTUtils.h>
#include <anthem/ASTVisitors.h>

namespace anthem
{

////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Completion
//
////////////////////////////////////////////////////////////////////////////////////////////////////

// Checks that two matching predicates (same name, same arity) have the same arguments
void checkMatchingPredicates(const ast::Term &lhs, const ast::Term &rhs)
{
	if (!lhs.is<ast::Variable>() || !rhs.is<ast::Variable>())
		throw std::runtime_error("cannot preform completion, only variables supported in predicates currently");

	if (lhs.get<ast::Variable>().name != rhs.get<ast::Variable>().name)
		throw std::runtime_error("cannot perform completion, inconsistent predicate argument naming");
}

////////////////////////////////////////////////////////////////////////////////////////////////////

void completePredicate(const ast::Predicate &predicate, std::vector<ast::Formula> &formulas, std::size_t &formulaIndex)
{
	// Check that predicate is in normal form
	for (auto i = formulaIndex; i < formulas.size(); i++)
	{
		auto &formula = formulas[i];
		assert(formula.is<ast::Implies>());
		auto &implies = formula.get<ast::Implies>();

		if (!implies.consequent.is<ast::Predicate>())
			continue;

		auto &otherPredicate = implies.consequent.get<ast::Predicate>();

		if (predicate.arity() != otherPredicate.arity() || predicate.name != otherPredicate.name)
			continue;

		for (std::size_t i = 0; i < predicate.arguments.size(); i++)
			checkMatchingPredicates(predicate.arguments[i], otherPredicate.arguments[i]);
	}

	// Copy the predicate’s arguments for the completed formula
	std::vector<ast::Variable> variables;
	variables.reserve(predicate.arguments.size());

	for (const auto &argument : predicate.arguments)
	{
		assert(argument.is<ast::Variable>());
		variables.emplace_back(ast::deepCopy(argument.get<ast::Variable>()));
	}

	auto or_ = ast::Formula::make<ast::Or>();

	// Build the conjunction of all formulas with the predicate as consequent
	for (auto i = formulaIndex; i < formulas.size();)
	{
		auto &formula = formulas[i];
		assert(formula.is<ast::Implies>());
		auto &implies = formula.get<ast::Implies>();

		if (!implies.consequent.is<ast::Predicate>())
		{
			i++;
			continue;
		}

		auto &otherPredicate = implies.consequent.get<ast::Predicate>();

		if (predicate.arity() != otherPredicate.arity() || predicate.name != otherPredicate.name)
		{
			i++;
			continue;
		}

		ast::VariableStack variableStack;
		variableStack.push(&variables);

		auto variables = ast::collectFreeVariables(implies.antecedent, variableStack);

		if (variables.empty())
			or_.get<ast::Or>().arguments.emplace_back(std::move(implies.antecedent));
		else
		{
			auto exists = ast::Formula::make<ast::Exists>(std::move(variables), std::move(implies.antecedent));
			or_.get<ast::Or>().arguments.emplace_back(std::move(exists));
		}

		if (i > formulaIndex)
			formulas.erase(formulas.begin() + i);
		else
			i++;
	}

	auto biconditionalRight = std::move(or_);

	// If the disjunction contains only one element, drop the enclosing disjunction
	if (biconditionalRight.get<ast::Or>().arguments.size() == 1)
		biconditionalRight = biconditionalRight.get<ast::Or>().arguments.front();

	// If the biconditional would be of the form “F <-> true” or “F <-> false,” simplify the output
	if (biconditionalRight.is<ast::Boolean>())
	{
		const auto &boolean = biconditionalRight.get<ast::Boolean>();

		if (boolean.value == true)
			formulas[formulaIndex] = ast::deepCopy(predicate);
		else
			formulas[formulaIndex] = ast::Formula::make<ast::Not>(ast::deepCopy(predicate));

		formulaIndex++;
		return;
	}

	// Build the biconditional within the completed formula
	auto biconditional = ast::Formula::make<ast::Biconditional>(ast::deepCopy(predicate), std::move(biconditionalRight));

	if (predicate.arguments.empty())
	{
		formulas[formulaIndex] = std::move(biconditional);
		formulaIndex++;
		return;
	}

	auto completedFormula = ast::Formula::make<ast::ForAll>(std::move(variables), std::move(biconditional));

	formulas[formulaIndex] = std::move(completedFormula);
	formulaIndex++;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

void completeBoolean(std::vector<ast::Formula> &formulas, std::size_t &formulaIndex)
{
	auto &formula = formulas[formulaIndex];
	assert(formula.is<ast::Implies>());
	auto &implies = formula.get<ast::Implies>();
	assert(implies.consequent.is<ast::Boolean>());
	auto &boolean = implies.consequent.get<ast::Boolean>();

	auto variables = ast::collectFreeVariables(implies.antecedent);

	// Implications of the form “T -> true” are useless
	if (boolean.value == true)
	{
		formulas.erase(formulas.begin() + formulaIndex);
		return;
	}

	auto argument = ast::Formula::make<ast::Not>(std::move(implies.antecedent));

	if (variables.empty())
	{
		formula = std::move(argument);
		formulaIndex++;
		return;
	}

	formula = ast::Formula::make<ast::ForAll>(std::move(variables), std::move(argument));
	formulaIndex++;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

void complete(std::vector<ast::Formula> &formulas)
{
	for (const auto &formula : formulas)
	{
		if (!formula.is<ast::Implies>())
			throw std::runtime_error("cannot perform completion, formula not in normal form");

		auto &implies = formula.get<ast::Implies>();

		if (!implies.consequent.is<ast::Predicate>() && !implies.consequent.is<ast::Boolean>())
			throw std::runtime_error("cannot perform completion, only single predicates and Booleans supported as formula consequent currently");
	}

	for (std::size_t i = 0; i < formulas.size();)
	{
		auto &formula = formulas[i];
		auto &implies = formula.get<ast::Implies>();

		if (implies.consequent.is<ast::Predicate>())
		{
			auto &predicate = implies.consequent.get<ast::Predicate>();
			completePredicate(predicate, formulas, i);
		}
		else if (implies.consequent.is<ast::Boolean>())
			completeBoolean(formulas, i);
	}
}

////////////////////////////////////////////////////////////////////////////////////////////////////

}