package prolog.logic

import org.junit.jupiter.api.Assertions.*
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
import org.junit.jupiter.api.Disabled
import org.junit.jupiter.api.Nested
import prolog.ast.lists.List
import prolog.ast.lists.List.Empty
import prolog.ast.lists.List.Cons

/*
 * Based on: https://en.wikipedia.org/wiki/Unification_%28computer_science%29#Examples_of_syntactic_unification_of_first-order_terms
 */
class UnificationTests {
    @Nested
    class `unify logic` {
        @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
        @Disabled("If the occurs check is applied, this should fail")
        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.
         */
        @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")
        }
    }

    @Nested
    class `applySubstitution logic` {
        @Test
        fun `apply substitution without sub`() {
            val term = Variable("X")
            val subs: Substitutions = emptyMap()

            val result = applySubstitution(term, subs)

            assertEquals(term, result)
        }

        @Test
        fun `apply single substitution`() {
            val sub = Variable("X") to Integer(5)
            val subs: Substitutions = mapOf(sub)

            val term = Variable("X")

            val result = applySubstitution(term, subs)

            assertEquals(Integer(5), result)
        }

        @Test
        fun `apply chained substitution`() {
            val sub1 = Variable("HP") to Variable("HP(19)")
            val sub2 = Variable("HP(19)") to Integer(35)

            val subs: Substitutions = mapOf(sub1, sub2)

            val term = Variable("HP")

            val result = applySubstitution(term, subs)

            assertEquals(Integer(35), result)
        }
    }

    @Nested
    class `equivalent logic` {
        @Test
        fun `empty lists are equivalent`() {
            var eq = equivalent(Empty, Empty, emptyMap())
            assertTrue(eq, "Empty lists should be equivalent")

            val variable = Variable("X")

            eq = equivalent(Empty, variable, mapOf(variable to Empty))
            assertTrue(eq, "Empty list and variable should be equivalent")

            eq = equivalent(variable, Empty, mapOf(variable to Empty))
            assertTrue(eq, "Variable and empty list should be equivalent")
        }

        @Test
        fun `singletons are equivalent`() {
            val atom = Atom("a")
            val list1 = Cons(atom, Empty)
            val list2 = Cons(atom, Empty)

            var eq = equivalent(list1, list2, emptyMap())
            assertTrue(eq, "Singleton lists should be equivalent")

            val variable = Variable("X")

            eq = equivalent(list1, variable, mapOf(variable to list2))
            assertTrue(eq, "Singleton list and variable should be equivalent")

            eq = equivalent(variable, list1, mapOf(variable to list2))
            assertTrue(eq, "Variable and singleton list should be equivalent")
        }

        @Test
        fun `singleton and empty list are not equivalent`() {
            val atom = Atom("a")
            val list1 = Cons(atom, Empty)
            val list2 = Empty

            var eq = equivalent(list1, list2, emptyMap())
            assertFalse(eq, "Singleton and empty lists should not be equivalent")

            eq = equivalent(list2, list1, emptyMap())
            assertFalse(eq, "Empty and singleton lists should not be equivalent")

            val variable = Variable("X")

            eq = equivalent(list1, variable, mapOf(variable to list2))
            assertFalse(eq, "Singleton list and variable should not be equivalent")

            eq = equivalent(variable, list1, mapOf(variable to list2))
            assertFalse(eq, "Variable and singleton list should not be equivalent")
        }

        @Test
        fun `identical lists are equivalent`() {
            val list1 = Cons(Atom("a"), Cons(Atom("b"), Empty))
            val list2 = Cons(Atom("a"), Cons(Atom("b"), Empty))

            var eq = equivalent(list1, list2, emptyMap())
            assertTrue(eq, "Identical lists should be equivalent")

            eq = equivalent(list2, list1, emptyMap())
            assertTrue(eq, "Identical lists should be equivalent")

            val variable = Variable("X")

            eq = equivalent(list1, variable, mapOf(variable to list2))
            assertTrue(eq, "Identical lists should be equivalent")

            eq = equivalent(variable, list2, mapOf(variable to list1))
            assertTrue(eq, "Identical lists should be equivalent")
        }

        @Test
        fun `identical nested lists are equivalent`() {
            val list1 = Cons(Atom("foo"), Cons(Atom("bar"), Cons(Cons(Atom("baz"), Cons(Atom("bar"), Empty)), Empty)))
            val list2 = Cons(Atom("foo"), Cons(Atom("bar"), Cons(Cons(Atom("baz"), Cons(Atom("bar"), Empty)), Empty)))

            var eq = equivalent(list1, list2, emptyMap())
            assertTrue(eq, "Identical nested lists should be equivalent")

            eq = equivalent(list2, list1, emptyMap())
            assertTrue(eq, "Identical nested lists should be equivalent")

            val variable = Variable("X")

            eq = equivalent(list1, variable, mapOf(variable to list2))
            assertTrue(eq, "Identical nested lists should be equivalent")

            eq = equivalent(variable, list2, mapOf(variable to list1))
            assertTrue(eq, "Identical nested lists should be equivalent")
        }

        @Test
        fun `lists with different nests are not equivalent`() {
            val list1 = Cons(Atom("foo"), Cons(Atom("bar"), Cons(Cons(Atom("bar"), Cons(Atom("baz"), Empty)), Empty)))
            val list2 = Cons(Atom("foo"), Cons(Atom("bar"), Cons(Cons(Atom("baz"), Cons(Atom("bar"), Empty)), Empty)))

            var eq = equivalent(list1, list2, emptyMap())
            assertFalse(eq, "Lists with different nests should not be equivalent")

            eq = equivalent(list2, list1, emptyMap())
            assertFalse(eq, "Lists with different nests should not be equivalent")

            val variable = Variable("X")

            eq = equivalent(list1, variable, mapOf(variable to list2))
            assertFalse(eq, "Lists with different nests should not be equivalent")

            eq = equivalent(variable, list2, mapOf(variable to list1))
            assertFalse(eq, "Lists with different nests should not be equivalent")
        }
    }
}