2025LogProg-project-GhentProlog/src/prolog/builtins/arithmeticOperators.kt

package prolog.builtins

import prolog.Answers
import prolog.Substitutions
import prolog.ast.arithmetic.ArithmeticOperator
import prolog.ast.arithmetic.Expression
import prolog.ast.arithmetic.Integer
import prolog.ast.arithmetic.Simplification
import prolog.ast.logic.Satisfiable
import prolog.ast.terms.*
import prolog.logic.*

// TODO >

// TODO <

// TODO =<

// TODO >=

/**
 * True if expression Expr1 evaluates to a number non-equal to Expr2.
 */
class EvaluatesToDifferent(private val left: Expression, private val right: Expression) :
    Operator("=\\=", left, right), Satisfiable {
    override fun satisfy(subs: Substitutions): Answers {
        val t1 = left.simplify(subs)
        val t2 = right.simplify(subs)

        // Should both be instantiated
        if (!atomic(t1.to, subs) || !atomic(t2.to, subs)) {
            return sequenceOf(Result.failure(IllegalArgumentException("Both operands must be instantiated")))
        }

        return if (equivalent(t1.to, t2.to, subs)) {
            emptySequence()
        } else {
            sequenceOf(Result.success(emptyMap()))
        }
    }

    override fun applySubstitution(subs: Substitutions): EvaluatesToDifferent = EvaluatesToDifferent(
        applySubstitution(left, subs),
        applySubstitution(right, subs)
    )
}

/**
 * True if Expr1 evaluates to a number equal to Expr2.
 */
class EvaluatesTo(private val left: Expression, private val right: Expression) :
    Operator("=:=", left, right) {
    override fun satisfy(subs: Substitutions): Answers {
        val t1 = left.simplify(subs)
        val t2 = right.simplify(subs)

        // Should both be instantiated
        if (!atomic(t1.to, subs) || !atomic(t2.to, subs)) {
            return sequenceOf(Result.failure(IllegalArgumentException("Both operands must be instantiated")))
        }

        return if (equivalent(t1.to, t2.to, subs)) sequenceOf(Result.success(emptyMap())) else emptySequence()
    }

    override fun applySubstitution(subs: Substitutions): EvaluatesTo = EvaluatesTo(
        applySubstitution(left, subs),
        applySubstitution(right, subs)
    )
}

/**
 * True when Number is the value to which Expr evaluates.
 */
class Is(val number: Expression, val expr: Expression) :
    Operator("is", number, expr), Satisfiable {
    override fun satisfy(subs: Substitutions): Answers {
        val t1 = number.simplify(subs)
        val t2 = expr.simplify(subs)

        if (!atomic(t2.to, subs)) {
            return sequenceOf(Result.failure(IllegalArgumentException("Right operand must be instantiated")))
        }

        if (!variable(t1.to, subs) && equivalent(t1.to, t2.to, subs + listOfNotNull(t1.mapped, t2.mapped))) {
            return sequenceOf(Result.success(emptyMap()))
        }

        return unifyLazy(t1.to, t2.to, subs)
    }

    override fun applySubstitution(subs: Substitutions): Is = Is(
        applySubstitution(number, subs),
        applySubstitution(expr, subs)
    )
}

/**
 * Result = -Expr
 */
class Negate(operand: Expression) : Subtract(Integer(0), operand)

/**
 * Result = Expr
 */
class Positive(operand: Expression) : Add(Integer(0), operand)

/**
 * Result = Expr1 + Expr2
 */
open class Add(private val expr1: Expression, private val expr2: Expression) :
    ArithmeticOperator("+", expr1, expr2) {
    override fun simplify(subs: Substitutions): Simplification {
        val result = Variable("Result")
        val map = plus(expr1, expr2, result, subs)
        val simplification = result.simplify(map.first().getOrThrow())
        return Simplification(this, simplification.to)
    }

    override fun applySubstitution(subs: Substitutions): Add = Add(
        applySubstitution(expr1, subs),
        applySubstitution(expr2, subs)
    )
}

class Plus(private val expr1: Expression, private val expr2: Expression, private val expr3: Expression) :
    CompoundTerm("plus", expr1, expr2, expr3) {
    override fun satisfy(subs: Substitutions): Answers = plus(expr1, expr2, expr3, subs)
}

/**
 * Result = Expr1 - Expr2
 */
open class Subtract(private val expr1: Expression, private val expr2: Expression) :
    ArithmeticOperator("-", expr1, expr2) {
    override fun simplify(subs: Substitutions): Simplification {
        val result = Variable("Result")
        val map = plus(expr2, result, expr1, subs)
        val simplification = result.simplify(map.first().getOrThrow())
        return Simplification(this, simplification.to)
    }

    override fun applySubstitution(subs: Substitutions): Subtract = Subtract(
        applySubstitution(expr1, subs),
        applySubstitution(expr2, subs)
    )
}

/**
 * Result = Expr1 * Expr2
 */
class Multiply(val expr1: Expression, val expr2: Expression) :
    ArithmeticOperator("*", expr1, expr2) {
    override fun simplify(subs: Substitutions): Simplification {
        val result = Variable("Result")
        val map = mul(expr1, expr2, result, subs)
        val simplification = result.simplify(map.first().getOrThrow())
        return Simplification(this, simplification.to)
    }

    override fun applySubstitution(subs: Substitutions): Multiply = Multiply(
        applySubstitution(expr1, subs),
        applySubstitution(expr2, subs)
    )
}

class Divide(private val expr1: Expression, private val expr2: Expression) :
    ArithmeticOperator("/", expr1, expr2) {
    override fun simplify(subs: Substitutions): Simplification {
        val result = Variable("Result")
        val map = div(expr1, expr2, result, subs)
        val simplification = result.simplify(map.first().getOrThrow())
        return Simplification(this, simplification.to)
    }

    override fun applySubstitution(subs: Substitutions): Divide = Divide(
        applySubstitution(expr1, subs),
        applySubstitution(expr2, subs)
    )
}

// TODO Expr mod Expr

// TODO Expr rem Expr

class Between(private val expr1: Expression, private val expr2: Expression, private val expr3: Expression) :
    CompoundTerm("between", expr1, expr2, expr3), Satisfiable {
    override fun satisfy(subs: Substitutions): Answers {
        val e1 = expr1.simplify(subs)
        val e2 = expr2.simplify(subs)
        val e3 = expr3.simplify(subs)

        require(e1.to is Integer && e2.to is Integer) { "Arguments must be integers" }

        val v1 = e1.to
        val v2 = e2.to

        return if (variable(e3.to, subs)) {
            between(v1, v2, e3.to as Variable).map { answer ->
                answer.mapCatching { it + listOfNotNull(e1.mapped, e2.mapped) }
            }
        } else {
            between(v1, v2, e3.to as Integer).map { answer ->
                answer.mapCatching { it + listOfNotNull(e1.mapped, e2.mapped) }
            }
        }
    }

    override fun applySubstitution(subs: Substitutions): Between = Between(
        applySubstitution(expr1, subs),
        applySubstitution(expr2, subs),
        applySubstitution(expr3, subs)
    )

    override fun toString(): String = "$expr1..$expr3..$expr2"
}

class Successor(private val expr1: Expression, private val expr2: Expression) :
    CompoundTerm("succ", expr1, expr2), Satisfiable {
    override fun satisfy(subs: Substitutions): Answers = succ(expr1, expr2, subs)
}