package prolog.logic

import org.junit.jupiter.api.Assertions.*
import org.junit.jupiter.api.Disabled
import org.junit.jupiter.api.Test
import prolog.Substitutions
import prolog.ast.arithmetic.Integer
import prolog.ast.terms.Atom
import prolog.ast.terms.Structure
import prolog.ast.terms.Variable
import prolog.builtins.Add

/*
 * Based on: https://en.wikipedia.org/wiki/Unification_%28computer_science%29#Examples_of_syntactic_unification_of_first-order_terms
 */
class UnifyTests {
    @Test
    fun identical_atoms_unify() {
        val atom1 = Atom("a")
        val atom2 = Atom("a")

        val result = unify(atom1, atom2)

        assertTrue(result.isSuccess, "Identical atoms should unify")
        assertEquals(0, result.getOrNull()!!.size, "No substitutions should be made")
    }

    @Test
    fun different_atoms_do_not_unify() {
        val atom1 = Atom("a")
        val atom2 = Atom("b")

        val result = unify(atom1, atom2)

        assertFalse(result.isSuccess, "Different atoms should not unify")
    }

    /**
     * ?- X = X.
     * true.
     */
    @Test
    fun identical_variables_unify() {
        val variable1 = Variable("X")
        val variable2 = Variable("X")

        val result = unify(variable1, variable2)

        assertTrue(result.isSuccess, "Identical variables should unify")
        assertEquals(0, result.getOrNull()!!.size, "No substitutions should be made")
    }

    @Test
    fun variable_unifies_with_atom() {
        val variable = Variable("X")
        val atom = Atom("a")

        val result = unify(atom, variable)

        assertTrue(result.isSuccess, "Variable should unify with atom")
        assertEquals(1, result.getOrNull()!!.size, "There should be one substitution")
        assertEquals(atom, result.getOrNull()!![variable], "Variable should be substituted with atom")
    }

    @Test
    fun variables_alias_when_unified() {
        val variable1 = Variable("X")
        val variable2 = Variable("Y")

        val result = unify(variable1, variable2)

        assertTrue(result.isSuccess)
        assertEquals(1, result.getOrNull()!!.size)
        assertEquals(variable2, result.getOrNull()!![variable1], "Variable 1 should alias to variable 2")
    }

    @Test
    fun identical_compound_terms_unify() {
        val structure1 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        val structure2 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))

        val result = unify(structure1, structure2)

        assertTrue(result.isSuccess, "Identical compound terms should unify")
        assertEquals(0, result.getOrNull()!!.size, "No substitutions should be made")
    }

    @Test
    fun compound_terms_with_different_arguments_do_not_unify() {
        val structure1 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        val structure2 = Structure(Atom("f"), listOf(Atom("a"), Atom("c")))

        val result = unify(structure1, structure2)

        assertFalse(result.isSuccess, "Different compound terms should not unify")
    }

    @Test
    fun compound_terms_with_different_functors_do_not_unify() {
        val structure1 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))
        val structure2 = Structure(Atom("g"), listOf(Atom("a"), Atom("b")))

        val result = unify(structure1, structure2)

        assertFalse(result.isSuccess, "Compound terms with different functors should not unify")
    }

    /**
     * ?- X = f(a, b).
     * X = f(a, b).
     */
    @Test
    fun variable_unifies_with_compound_term() {
        val variable = Variable("X")
        val structure = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))

        val result = unify(variable, structure)

        assertTrue(result.isSuccess, "Variable should unify with compound term")

        val subs = result.getOrNull()!!

        assertEquals(1, subs.size, "There should be one substitution")
        assertTrue(subs.containsKey(variable), "Variable should be in the substitution map")
        assertTrue(
            equivalent(Structure(Atom("f"), listOf(Atom("a"), Atom("b"))), subs[variable]!!, subs),
            "Variable should be substituted with compound term"
        )
    }

    @Test
    fun compound_term_with_variable_unifies_with_part() {
        val variable = Variable("X")
        val structure1 = Structure(Atom("f"), listOf(Atom("a"), variable))
        val structure2 = Structure(Atom("f"), listOf(Atom("a"), Atom("b")))

        val result = unify(structure1, structure2)

        assertTrue(result.isSuccess, "Compound term with variable should unify with part")

        val subs = result.getOrNull()!!

        assertEquals(1, subs.size, "There should be one substitution")
        assertTrue(subs.containsKey(variable), "Variable should be in the substitution map")
        val equivalence = equivalent(Atom("b"), subs[variable]!!, subs)
        assertTrue(equivalence, "Variable should be substituted with atom")
    }

    @Test
    fun compound_terms_with_variable_arguments_lists_alias_variables() {
        val variable1 = Variable("X")
        val variable2 = Variable("Y")

        val structure1 = Structure(Atom("f"), listOf(variable1))
        val structure2 = Structure(Atom("f"), listOf(variable2))

        val result = unify(structure1, structure2)

        assertTrue(result.isSuccess, "Compound terms with variable arguments should unify")

        val subs = result.getOrNull()!!

        assertEquals(1, subs.size, "There should be one substitution")
        assertTrue(subs.containsKey(variable1), "Variable 1 should be in the substitution map")
        assertEquals(variable2, subs[variable1], "Variable 1 should alias to variable 2")
    }

    /**
     * f(X) = f(Y, Z)
     */
    @Test
    fun compound_terms_with_different_arity_do_not_unify() {
        val structure1 = Structure(Atom("f"), listOf(Variable("X")))
        val structure2 = Structure(Atom("f"), listOf(Variable("Y"), Variable("Z")))

        val result = unify(structure1, structure2)

        assertFalse(result.isSuccess, "Compound terms with different arity should not unify")
    }

    /**
     * ?- f(g(X)) = f(Y).
     * Y = g(X).
     */
    @Test
    fun nested_compound_terms_with_variables_unify() {
        val variable2 = Variable("Y")

        val structure1 = Structure(Atom("f"), listOf(Structure(Atom("g"), listOf(Variable("X")))))
        val structure2 = Structure(Atom("f"), listOf(variable2))

        val result = unify(structure1, structure2)

        assertTrue(result.isSuccess, "Nested compound terms with variables should unify")

        val subs = result.getOrNull()!!

        assertEquals(1, subs.size, "There should be one substitution")
        assertTrue(subs.containsKey(variable2), "Variable 2 should be in the substitution map")
        assertTrue(
            equivalent(Structure(Atom("g"), listOf(Variable("X"))), subs[variable2]!!, subs),
            "Variable should be substituted with compound term"
        )
    }

    /**
     * ?- f(g(X), X) = f(Y, a).
     * X = a,
     * Y = g(a).
     */
    @Test
    fun compound_terms_with_more_variables() {
        val variable1 = Variable("X")
        val variable2 = Variable("Y")

        val structure1 = Structure(Atom("f"), listOf(Structure(Atom("g"), listOf(variable1)), variable1))
        val structure2 = Structure(Atom("f"), listOf(variable2, Atom("a")))

        val result = unify(structure1, structure2)

        assertTrue(result.isSuccess, "Compound terms with more variables should unify")

        val subs = result.getOrNull()!!

        assertEquals(2, subs.size, "There should be two substitutions")
        assertTrue(subs.containsKey(variable1), "Variable 1 should be in the substitution map")
        assertTrue(
            equivalent(Atom("a"), subs[variable1]!!, subs),
            "Variable 1 should be substituted with atom"
        )
        assertTrue(subs.containsKey(variable2), "Variable 2 should be in the substitution map")
        assertTrue(
            equivalent(Structure(Atom("g"), listOf(Atom("a"))), subs[variable2]!!, subs),
            "Variable 2 should be substituted with compound term"

        )
    }

    /**
     * ?- X = f(X).
     * X = f(f(X)).
     */
    @Test
    fun recursive_unification() {
        val variable1 = Variable("X")
        val structure2 = Structure(Atom("f"), listOf(Variable("X")))

        val result = unifyLazy(variable1, structure2, emptyMap()).toList()

        assertEquals(1, result.size, "There should be one result")
        assertTrue(result[0].isSuccess, "Recursive unification should succeed")

        val subs = result[0].getOrNull()!!

        assertEquals(1, subs.size, "There should be one substitution")
        assertTrue(subs.containsKey(variable1), "Variable should be in the substitution map")
        assertEquals(structure2, subs[variable1], "Variable should be substituted with compound term")
    }

    /**
     * ?- X = bar, Y = bar, X = Y.
     * X = Y, Y = bar.
     */
    @Disabled
    @Test
    fun multiple_unification() {
        val variable1 = Variable("X")
        val variable2 = Variable("Y")
        val atom = Atom("bar")

        val map: Substitutions = mapOf(
            variable1 to atom,
            variable2 to atom
        )
        val result = unifyLazy(variable1, variable2, map).toList()

        assertEquals(1, result.size, "There should be one substitution")
        assertTrue(result[0].isSuccess, "Multiple unification should succeed")
        assertEquals(0, result[0].getOrNull()!!.size, "No (additional) substitutions should be made")
    }

    /**
     * ?- a = a().
     * false.
     */
    @Test
    fun atom_with_different_arity() {
        val atom1 = Atom("a")
        val structure2 = Structure(Atom("a"), emptyList())

        val result = unify(atom1, structure2)

        assertFalse(result.isSuccess, "Atom with different arity should not unify")
    }

    @Test
    fun identical_integers_unify() {
        val int1 = Integer(1)
        val int2 = Integer(1)

        val result = unify(int1, int2)

        assertTrue(result.isSuccess, "Identical integers should unify")
        assertEquals(0, result.getOrNull()!!.size, "No substitutions should be made")
    }

    @Test
    fun different_integers_do_not_unify() {
        val int1 = Integer(1)
        val int2 = Integer(2)

        val result = unify(int1, int2)

        assertFalse(result.isSuccess, "Different integers should not unify")
    }


    @Test
    fun `1 + 2 does not unify with 3`() {
        val expr1 = Add(Integer(1), Integer(2))
        val expr2 = Integer(3)

        val result = unify(expr1, expr2)

        assertFalse(result.isSuccess, "1 + 2 should not unify with 3")
    }
}