"""
    Terms are values in an equation. They can be constants, variables,
    functions, operators, and many more.
"""

from __future__ import annotations

import copy

from dataclasses import dataclass

__known_consts : dict[str, Term] = {
}

def register_known_const(name : str, value : Term) -> None:
    __known_consts[name] = value

class Term:
    """
        Base class for all terms.
    """
    
    def normalize(self) -> Term:
        """
            Function to return a simplified version of the term.
        """
        return self
    
    def reduce_on_equality(self, other : Term) -> tuple[Term, Term]:
        """
            Reduce parts that the two items are equal on.
        """
        return self, other
    
    def substitute(self, index : int, value : Term) -> Term:
        """
            Substitute all variables of a given index number.

            :param index: The variable's index.
            :type index: int
            :param value: The value to replace the variable with.
            :type value: 
        """
        return self

@dataclass(frozen=True)
class App(Term):
    """
        An App is the application of a function with an argument.
        You could consider it the inverse of the lambda function.
    """

    f : Term
    x : Term

    def normalize(self) -> Term:
        f = self.f.normalize()
        x = self.x.normalize()

        default = App(f=f, x=x)

        match f:
            case App():
                return default

            case Const():
                return default

            case Lambda():
                return f.resolve(x)

            case Term():
                return default
            
            case Var():
                return default
    
    def reduce_on_equality(self, other: Term) -> tuple[Term, Term]:
        if not isinstance(other, App):
            return self, other
        
        f1, f2 = self.f.normalize(), other.f.normalize()
        x1, x2 = self.x.normalize(), other.x.normalize()

        match ( f1 == f2, x1 == x2 ):
            case ( True, True ):
                return Term(), Term()
            
            case ( True, False ):
                return x1.reduce_on_equality(x2)
            
            case ( False, True ):
                return f1.reduce_on_equality(f2)
            
            case ( False, False ):
                return self, other

    def substitute(self, index : int, value : Term) -> Term:
        """
            Substitute all variables of a given index number.

            :param index: The variable's index.
            :type index: int
            :param value: The value to replace the variable with.
            :type value: 
        """
        return App(
            f=self.f.substitute(index=index, value=value),
            x=self.x.substitute(index=index, value=value),
        )

@dataclass(frozen=True)
class Const(Term):
    """
        Constants are 0-ary functions. They can represent variables, constants
        or other values that generally take no input.
    """

    name : str

    def normalize(self) -> Term:
        if self.name in __known_consts:
            return copy.deepcopy(__known_consts.get(self.name, self))
        else:
            return self

@dataclass(frozen=True)
class Lambda(Term):
    """
        Lambdas a nameless 1-ary functions.
    """

    out : Term

    def normalize(self) -> Term:
        return Lambda(out=self.out.normalize())
    
    def reduce_on_equality(self, other: Term) -> tuple[Term, Term]:
        match other:
            case Lambda():
                return self.out, other.out
            case _:
                return self, other

    def resolve(self, value : Term) -> Term:
        """
            Resolve this lambda function by inserting a value.

            :param value: The value to substitute.
        """
        return self.substitute(index=-1, value=value)

    def substitute(self, index : int, value : Term) -> Term:
        """
            Substitute all variables of a given index number.

            :param index: The variable's index.
            :type index: int
            :param value: The value to replace the variable with.
            :type value: 
        """
        if index == -1:
            # We're resolving this lambda function!
            return self.out.substitute(index=index + 1, value=value)
        else:
            return Lambda(out=self.out.substitute(index=index + 1, value=value))

@dataclass(frozen=True)
class Var(Term):
    """
        A variable is a value that can be substituted by a lambda function.
    """

    index : int

    def substitute(self, index : int, value : Term) -> Term:
        """
            Substitute all variables of a given index number.

            :param index: The variable's index.
            :type index: int
            :param value: The value to replace the variable with.
            :type value: 
        """
        if self.index == index:
            return copy.deepcopy(value)
        else:
            return self