package at.ac.tuwien.kr.alpha.core.programs.transformation;

import java.util.ArrayList;
import java.util.List;

import at.ac.tuwien.kr.alpha.api.programs.NormalProgram;
import at.ac.tuwien.kr.alpha.api.programs.atoms.Atom;
import at.ac.tuwien.kr.alpha.api.programs.atoms.BasicAtom;
import at.ac.tuwien.kr.alpha.api.programs.atoms.ComparisonAtom;
import at.ac.tuwien.kr.alpha.api.programs.literals.Literal;
import at.ac.tuwien.kr.alpha.api.programs.rules.NormalRule;
import at.ac.tuwien.kr.alpha.api.programs.rules.heads.NormalHead;
import at.ac.tuwien.kr.alpha.api.programs.terms.ArithmeticTerm;
import at.ac.tuwien.kr.alpha.api.programs.terms.ConstantTerm;
import at.ac.tuwien.kr.alpha.api.programs.terms.FunctionTerm;
import at.ac.tuwien.kr.alpha.api.programs.terms.IntervalTerm;
import at.ac.tuwien.kr.alpha.api.programs.terms.Term;
import at.ac.tuwien.kr.alpha.api.programs.terms.VariableTerm;
import at.ac.tuwien.kr.alpha.commons.comparisons.ComparisonOperators;
import at.ac.tuwien.kr.alpha.commons.programs.Programs;
import at.ac.tuwien.kr.alpha.commons.programs.atoms.Atoms;
import at.ac.tuwien.kr.alpha.commons.programs.rules.Rules;
import at.ac.tuwien.kr.alpha.commons.programs.rules.heads.Heads;
import at.ac.tuwien.kr.alpha.commons.programs.terms.Terms;
import at.ac.tuwien.kr.alpha.commons.util.Util;

/**
 * Transforms rules such that arithmetic terms only occur in comparison predicates.
 * For example p(X+1) :- q(Y/2), r(f(X*2),Y), X-2 = Y*3, X = 0..9. is transformed into
 * p(_A1) :- q(_A2), r(f(_A3),Y), X-2 = Y*3, _A1 = X+1, _A2 = Y/2, _A3 = X*2, X = 0..9.
 *
 * Copyright (c) 2020-2021, the Alpha Team.
 */
public class ArithmeticTermsRewriting extends ProgramTransformation<NormalProgram, NormalProgram> {
	private static final String ARITHMETIC_VARIABLES_PREFIX = "_A";
	private int numArithmeticVariables;

	@Override
	public NormalProgram apply(NormalProgram inputProgram) {
		List<NormalRule> rewrittenRules = new ArrayList<>();
		boolean didRewrite = false;
		for (NormalRule inputProgramRule : inputProgram.getRules()) {
			if (containsArithmeticTermsToRewrite(inputProgramRule)) {
				rewrittenRules.add(rewriteRule(inputProgramRule));
				didRewrite = true;
			} else {
				// Keep rule as-is if no ArithmeticTerm occurs.
				rewrittenRules.add(inputProgramRule);
			}
		}
		if (!didRewrite) {
			return inputProgram;
		}
		// Create new program with rewritten rules.
		return Programs.newNormalProgram(rewrittenRules, inputProgram.getFacts(), inputProgram.getInlineDirectives());
	}

	/**
	 * Takes a normal rule and rewrites it such that {@link ArithmeticTerm}s only appear inside
	 * {@link at.ac.tuwien.kr.alpha.common.atoms.ComparisonLiteral}s.
	 *
	 * @param inputProgramRule the rule to rewrite.
	 * @return the rewritten rule. Note that a new {@link NormalRule} is returned for every call of this method.
	 */
	private NormalRule rewriteRule(NormalRule inputProgramRule) {
		numArithmeticVariables = 0; // Reset number of introduced variables for each rule.
		NormalHead rewrittenHead = null;
		List<Literal> rewrittenBodyLiterals = new ArrayList<>();
		// Rewrite head.
		if (!inputProgramRule.isConstraint()) {
			BasicAtom headAtom = inputProgramRule.getHeadAtom();
			if (containsArithmeticTermsToRewrite(headAtom)) {
				rewrittenHead = Heads.newNormalHead((BasicAtom) rewriteAtom(headAtom, rewrittenBodyLiterals));
			} else {
				rewrittenHead = inputProgramRule.getHead();
			}
		}
		// Rewrite body.
		for (Literal literal : inputProgramRule.getBody()) {
			if (!containsArithmeticTermsToRewrite(literal.getAtom())) {
				// Keep body literal as-is if no ArithmeticTerm occurs.
				rewrittenBodyLiterals.add(literal);
				continue;
			}
			rewrittenBodyLiterals.add(rewriteAtom(literal.getAtom(), rewrittenBodyLiterals).toLiteral(!literal.isNegated()));
		}
		return Rules.newNormalRule(rewrittenHead, rewrittenBodyLiterals);
	}

	/**
	 * Checks whether a normal rule contains an {@link ArithmeticTerm} outside of a
	 * {@link at.ac.tuwien.kr.alpha.common.atoms.ComparisonLiteral}.
	 *
	 * @param inputProgramRule the rule to check for presence of arithmetic terms outside comparison literals.
	 * @return true if the inputProgramRule contains an {@link ArithmeticTerm} outside of a
	 *         {@link at.ac.tuwien.kr.alpha.common.atoms.ComparisonLiteral}.
	 */
	private boolean containsArithmeticTermsToRewrite(NormalRule inputProgramRule) {
		if (!inputProgramRule.isConstraint()) {
			Atom headAtom = inputProgramRule.getHeadAtom();
			if (containsArithmeticTermsToRewrite(headAtom)) {
				return true;
			}
		}
		// Check whether body contains an ArithmeticTerm.
		for (Literal literal : inputProgramRule.getBody()) {
			if (containsArithmeticTermsToRewrite(literal.getAtom())) {
				return true;
			}
		}
		return false;
	}

	private Term rewriteArithmeticSubterms(Term term, List<Literal> bodyLiterals) {
		// Keep term as-is if it contains no ArithmeticTerm.
		if (!containsArithmeticTerm(term)) {
			return term;
		}
		// Switch on term type.
		if (term instanceof ArithmeticTerm) {
			VariableTerm replacementVariable = Terms.newVariable(ARITHMETIC_VARIABLES_PREFIX + numArithmeticVariables++);
			bodyLiterals.add(Atoms.newComparisonAtom(replacementVariable, term, ComparisonOperators.EQ).toLiteral());
			return replacementVariable;
		} else if (term instanceof VariableTerm || term instanceof ConstantTerm) {
			return term;
		} else if (term instanceof FunctionTerm) {
			List<Term> termList = ((FunctionTerm) term).getTerms();
			List<Term> rewrittenTermList = new ArrayList<>();
			for (Term subterm : termList) {
				rewrittenTermList.add(rewriteArithmeticSubterms(subterm, bodyLiterals));
			}
			return Terms.newFunctionTerm(((FunctionTerm) term).getSymbol(), rewrittenTermList);
		} else {
			throw Util.oops("Rewriting unknown Term type: " + term.getClass());
		}
	}

	private Atom rewriteAtom(Atom atomToRewrite, List<Literal> bodyLiterals) {
		if (atomToRewrite instanceof ComparisonAtom) {
			throw Util.oops("Trying to rewrite ComparisonAtom.");
		}
		List<Term> rewrittenTerms = new ArrayList<>();
		for (Term atomTerm : atomToRewrite.getTerms()) {
			// Rewrite arithmetic term.
			rewrittenTerms.add(rewriteArithmeticSubterms(atomTerm, bodyLiterals));
		}
		return atomToRewrite.withTerms(rewrittenTerms);
	}

	private boolean containsArithmeticTermsToRewrite(Atom atom) {
		// ComparisonAtom needs no rewriting.
		if (atom instanceof ComparisonAtom) {
			return false;
		}
		for (Term term : atom.getTerms()) {
			if (containsArithmeticTerm(term)) {
				return true;
			}
		}
		return false;
	}

	private boolean containsArithmeticTerm(Term term) {
		// Note: this check probably should be part of the Term interface and done by subtype polymorphism.
		if (term instanceof ArithmeticTerm) {
			return true;
		} else if (term instanceof ConstantTerm || term instanceof VariableTerm) {
			return false;
		} else if (term instanceof FunctionTerm) {
			for (Term subterm : ((FunctionTerm) term).getTerms()) {
				if (containsArithmeticTerm(subterm)) {
					return true;
				}
			}
			return false;
		} else if (term instanceof IntervalTerm) {
			return containsArithmeticTerm(((IntervalTerm) term).getLowerBound()) || containsArithmeticTerm(((IntervalTerm) term).getUpperBound());
		} else {
			throw Util.oops("Unexpected term type: " + term.getClass());
		}
	}
}
