from emis_funky_funktions import *

from dataclasses import dataclass
from functools import reduce
from typing import Collection, FrozenSet, Sequence, TypeAlias

@dataclass(frozen=True)
class Subst:
	"""
	A substitution which may be performed on a term.

	Only variables may be substituted, but they may be substituted with any term.  That
	is, `x1/Chris` is valid, but `Chris/x1` is not.

	It is strongly recommended to avoid recursive substitutions, that is, substitutions
	which contain the variable they are replacing
	"""
	variable: str
	replacement: 'IRTerm'
	def __str__(self) -> str:
		return repr(self)
	def __repr__(self) -> str:
		return f'{self.replacement}/{self.variable}'

Substitutions: TypeAlias = Sequence[Subst]

@dataclass(frozen=True)
class UnificationMismatch:
	"""
	Indicates that two terms failed to unify

	Contains the two terms, each of which is a valid subterm of one of the two original
	terms, and which do not superficially unify.
	"""
	term1: 'IRTerm'
	term2: 'IRTerm'

@dataclass(frozen=True)
class LengthMismatch:
	"""
	Indicates that two clauses/argument lists failed to unify due to a length mismatch

    Contains the first element of the two lists which didn't have a corresponding term on
    the other side
	"""
	term: 'IRTerm'

UnificationError = UnificationMismatch | LengthMismatch

@dataclass(frozen=True)
class IRProp:
	"""
	Represents a proposition or object for resolution

	Can have any number of arguments, each of which should be a `IRTerm`.  Note that no
	distinction is made between predicates (n-arity, logical statements), functions
	(positive-arity functions over objects), and constants (stand-ins for objects)
	"""

	name: str
	"""
	The identifier of this thing, including its location in the source
	"""

	arguments: 'Tuple[IRTerm, ...]' = tuple()

	def subst(self, subst: Subst) -> 'IRTerm':
		"""
		Perform substitution on a proposition

		Returns the same proposition, but with any instances of the variable named in the
		substitution replaced with the contents of the substitution.

		>>> original = IRProp('angry', (IRVar('x1'),))
		>>> original
		angry(*x1)

		>>> original.subst(Subst('x1', IRProp('Alex')))
		angry(Alex())
		"""
		return IRProp(self.name, tuple(arg.subst(subst) for arg in self.arguments))
	def __str__(self) -> str:
		return repr(self)
	def __repr__(self) -> str:
		return f'{self.name}({",".join(str(arg) for arg in self.arguments)})'
	def __contains__(self, var: str) -> bool:
		"""
		Test if a variable with a given name exists in this term

		>>> 'x1' in IRProp('friends', [IRProp('John'), IRProp('mother_of', [IRVar('x1')])])
		True
		"""
		return any(var in term for term in self.arguments)

@dataclass(frozen=True)
class IRVar:
	"""
	A variable which may be substituted for any other term
	"""
	name: str
	def subst(self, subst: Subst) -> 'IRTerm':
		"""
		Perform substitution on a proposition

		Returns the same proposition, but with any instances of the variable named in the
		substitution replaced with the contents of the substitution.

		>>> IRVar('x1').subst(Subst('x1', IRProp('Alex')))
		Alex()

		>>> IRVar('x1').subst(Subst('x2', IRProp('Alex')))
		*x1
		"""
		if self.name == subst.variable:
			return subst.replacement
		else:
			return self
	def __str__(self) -> str:
		return repr(self)
	def __repr__(self) -> str:
		return f'*{self.name}'
	def __contains__(self, var: str) -> bool:
		"""
		Test if a variable with a given name exists in this term

		>>> 'x1' in IRVar('x1')
		True

		>>> 'x1' in IRVar('x2')
		False
		"""
		return var == self.name

@dataclass(frozen=True)
class IRNeg:
	"""
	A negated proposition
	"""
	inner: 'IRTerm'
	def subst(self, subst: Subst) -> 'IRTerm':
		"""
		Perform substitution on a proposition

		Returns the same proposition, but with any instances of the variable named in the
		substitution replaced with the contents of the substitution.

		>>> original = IRNeg(IRProp('happy', [IRVar('x1')]))
		>>> original
		¬happy(*x1)

		>>> original.subst(Subst('x1', IRProp('parent', [IRProp('Susie')])))
		¬happy(parent(Susie()))
		"""
		return IRNeg(self.inner.subst(subst))
	def __str__(self) -> str:
		return repr(self)
	def __repr__(self) -> str:
		return f'¬{self.inner}'
	def __contains__(self, var: str) -> bool:
		"""
		Test if a variable with a given name exists in this term

		>>> 'x1' in IRNeg(IRProp('round', [IRVar('x1')]))
		True
		"""
		return var in self.inner

IRTerm: TypeAlias = IRVar | IRProp | IRNeg
Clause: TypeAlias = FrozenSet[IRTerm]
Clause_: TypeAlias = Collection[IRTerm]
"""
A more general definition of `Clause` which uses a collection rather than a frozen set

Every `Clause` is a `Clause_`, but not vice versa.  In other words, true `Clause` is a
subclass of `Clause_`.

Due to this generalization, `Clause_` does not necessarily benefit from hashability or
deduplication.  It exists mostly to make inputing things easier, e.g. in doctests.
"""

sub_all: Callable[[Substitutions, IRTerm], IRTerm] = p(reduce, lambda t, s: t.subst(s)) #type:ignore
"""
Perform a series of substitutions on a term

Applies every substitution to the term in order

>>> sub_all(
...     [Subst('x1', IRVar('x2')), Subst('x2', IRProp('Karkat'))],
...     IRProp('kismesis', [IRVar('x1'), IRVar('x2')]),
... )
kismesis(Karkat(), Karkat())
"""

def unify(t1: IRTerm, t2: IRTerm) -> Result[Substitutions, UnificationError]:
	"""
	Attempt to find a substitution that unifies two terms

	If successful, the returned substitutions will cause both term to be equal, when
	applied to both.

	If this method fails, then the pair of subterms which caused the unification to fail
	are returned.

	>>> unify(
	...		IRProp('imaginary', [IRProp('Rufio')]),
	...		IRProp('imaginary', [IRVar('x1')])
	... )
	Ok((Rufio()/x1,))

	>>> unify(
	...		IRProp('dating', [IRProp('Jade'), IRVar('x1')]),
	...		IRProp('dating', [IRVar('x1'), IRProp('John')])
	... )
	Err(UnificationMismatch(term1=Jade(), term2=John()))

	>>> unify(
	...		IRProp('mother_of', [IRVar('x1')]),
	...		IRVar('x1')
	... )
	Err(UnificationMismatch(term1=mother_of(*x1), term2=*x1))
	"""
	match (t1, t2):
		case (IRVar(v1), IRVar(v2)) if v1 == v2:
			return Ok(tuple())
		case (IRVar(v), t_other) | (t_other, IRVar(v)) if v not in t_other: #type: ignore
			return Ok((Subst(v, t_other),))
		case (IRProp(n1, a1), IRProp(n2, a2)) if n1 == n2 and len(a1) == len(a2):
			return unify_lists(a1, a2)
		case (IRNeg(i1), IRNeg(i2)):
			return unify(i1, i2)
	return Err(UnificationMismatch(t1, t2))

def unify_lists(c1: Sequence[IRTerm], c2: Sequence[IRTerm]) -> Result[Substitutions, UnificationError]:
    """
    Attempt to perform unification on two term/argument lists

    See `unify()` for the details of how this works.  When working with lists, the same
    rules apply.  The substitutions, when applied to every term of both lists, will
    cause the lists to become exactly the same.

    Lists which are not the same length cannot be unified, and will always fail.

    Notice the difference between a list of `IRTerm`s and a `Clause`:  Namely, that a
    `Clause` is unordered, while a list of `IRTerm`s is ordered.

    >>> unify_lists(
    ...     [ IRProp('imaginary', [IRProp('Rufio')]), IRProp('friend', [IRVar('x1'), IRVar('x3')])      ],
    ...     [ IRProp('imaginary', [IRVar('x1')]),     IRProp('friend', [IRVar('x2'), IRProp('Tavros')]) ]
    ... )
    Ok((Rufio()/x1, Rufio()/x2, Tavros()/x3))

    >>> unify_lists(
    ...     [ IRProp('imaginary', [IRProp('Rufio')]), IRProp('friend', [IRVar('x1'), IRVar('x3')])      ],
    ...     [ IRProp('imaginary', [IRVar('x1')])                                                        ]
    ... )
    Err(LengthMismatch(term=friend(Rufio(),*x3)))
    """
    match (c1, c2):
        case ([], []):
            return Ok(tuple())
        case ([h1, *t1], [h2, *t2]):
            return unify(h1, h2) << (lambda subs:
                unify_lists((*map(p(sub_all,subs),t1),), (*map(p(sub_all,subs),t2),)) <= (
                    lambda final_subs: (*subs, *final_subs)))
        case ([h, *t], []) | ([], [h, *t]):
            return Err(LengthMismatch(h))
    raise Exception('Unreachable')

if __name__ == '__main__':
	import doctest
	doctest.testmod()