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

from __future__ import annotations
from typing import Callable

import copy

from dataclasses import dataclass

# Known const values that can be substituted into other values.
_known_consts : dict[str, Term] = {}

# Known functions to normalize certain specific terms with given shapes.
_known_norms : list[Callable[[Term], Term]] = []

USABLE_VAR_NAMES = [
    'x', 'y', 'z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l',
    'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w',
    'α', 'β', 'γ', 'δ', 'ε', 'ζ', 'η', 'θ', 'ι', 'κ',
    # 'λ', # Avoid confusion
    'μ', 'ν', 'ξ', 'ο', 'π', 'ρ', 'σ', 'τ', 'υ', 'φ', 'χ', 'ψ', 'ω'
]

def normalize(term : Term) -> Term:
    """
        Create a normalization function. This function simplifies functions as
        much as possible.
    """
    def norm_iter(term : Term) -> Term:
        """
            Normalize a single term. Use both the standard normalization scheme
        """
        for func in _known_norms:
            term = func(term)
        
        return term

    return traverse(f=norm_iter, term=term)

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

def register_known_norm(func : Callable[[Term], Term]) -> None:
    _known_norms.append(func)

def std_norm_iter(term : Term) -> Term:
    """
        Normalize a single term. Trust that all children values have
        already been normalized, or that they will be normalized later on.

        :param term: The term to normalize.
        :type term: Term
        :return: A normalized version of the term.
        :rtype: Term
    """
    match term:
        case App():
            match term.f:
                case Lambda():
                    return term.f.substitute(value=term.x)
                
                case _:
                    return term
        
        case Const():
            return _known_consts.get(term.name, term)
        
        case Lambda():
            return term
        
        case Var():
            return term
    
        case Term():
            return term
_known_norms.append(std_norm_iter)

def traverse(f : Callable[[Term], Term], term : Term) -> Term:
    """
        Traverse through the tree back-and-forth, aiming to run the given
        function as thoroughly as possible to manipulate the term tree.

        The function quits once no more operations are created.
        **WARNING:** This means the tree might traverse forever.

        :param f: The traversal function that replaces terms.
        :type f: Callable[[Term], Term]
        :param term: The term to traverse.
        :type term: Term
        :return: The result of the traversed term.
        :rtype: Term
    """
    a, b = Term(), term

    while a != b:
        a, b = b, traverse_bottom_up(f=f, term=traverse_top_down(f=f, term=term))
    
    return b

def traverse_bottom_up(f : Callable[[Term], Term], term : Term) -> Term:
    """
        Traverse a term from the bottom up, replacing values with whatever the
        inserted function returns.

        :param f: The traversal function that replaces terms.
        :type f: Callable[[Term], Term]
        :param term: The term to traverse.
        :type term: Term
        :return: The result of the traversed term.
        :rtype: Term
    """
    match term:
        case App():
            return f(App(
                f=traverse_bottom_up(f=f, term=term.f),
                x=traverse_bottom_up(f=f, term=term.x)),
            )
        
        case Const():
            return f(term)
        
        case Lambda():
            return term.replace_content(
                traverse_bottom_up(f=f, term=term.out)
            )
        
        case Var():
            return f(term)
        
        case Term():
            return f(term)

def traverse_top_down(f : Callable[[Term], Term], term : Term) -> Term:
    """
        Traverse a term from the top down, replacing values with whatever the
        inserted function returns.

        :param f: The traversal function that replaces terms.
        :type f: Callable[[Term], Term]
        :param term: The term to traverse.
        :type term: Term
        :return: The result of the traversed term.
        :rtype: Term
    """
    t = f(term)

    match t:
        case App():
            return App(
                f=traverse_top_down(f=f, term=t.f),
                x=traverse_top_down(f=f, term=t.x),
            )
        
        case Const():
            return t
        
        case Lambda():
            return t.replace_content(traverse_top_down(f=f, term=t.out))
        
        case Var():
            return t

        case Term():
            return t

class Term:
    """
        Base class for all terms.
    """

    def __eq__(self, other) -> bool:
        """
            Determine whether two terms are equal.
        """
        return str(self) == str(other)

    def __repr__(self) -> str:
        """
            Create a representation of a term.
        """
        return self.to_str(vars={})
    
    def normalize(self) -> Term:
        return normalize(term=self)

    def to_str(self, vars : dict[str, Var]) -> str:
        return "<undefined>"

@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 __repr__(self) -> str:
        """
            Create a representation of a term.
        """
        return self.to_str(vars={})

    def to_str(self, vars: dict[str, Var]) -> str:
        f_str = self.f.to_str(vars=vars)
        x_str = self.x.to_str(vars=vars)

        return f_str + " " + (f"({x_str})" if " " in x_str else x_str)

A1 = lambda fn, a : App(f=fn, x=a)
A2 = lambda fn, a, b : App(f=A1(fn, a), x=b)
A3 = lambda fn, a, b, c : App(f=A2(fn, a, b), x=c)
A4 = lambda fn, a, b, c, d : App(f=A3(fn, a, b, c), x=d)
A5 = lambda fn, a, b, c, d, e : App(f=A4(fn, a, b, c, d), x=e)
A6 = lambda fn, a, b, c, d, e, f : App(f=A5(fn, a, b, c, d, e), x=f)
A7 = lambda fn, a, b, c, d, e, f, g : App(f=A6(fn, a, b, c, d, e, f), x=g)
A8 = lambda fn, a, b, c, d, e, f, g, h : App(f=A7(fn, a, b, c, d, e, f, g), x=h)

@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 __repr__(self) -> str:
        """
            Create a representation of a term.
        """
        return self.to_str(vars={})

    def to_str(self, vars: dict[str, Var]) -> str:
        return self.name

class Lambda(Term):
    """
        Lambdas a nameless 1-ary functions.
    """

    out : Term
    var : Var

    def __init__(self, f : Callable[[Var], Term]) -> None:
        self.var = Var()
        self.out = f(self.var)

    def __repr__(self) -> str:
        """
            Create a representation of a term.
        """
        return self.to_str(vars={})
    
    def replace_content(self, term : Term) -> Lambda:
        """
            Replace the content of the lambda function with the given term.
            The new term may still contain the lambda's var as an argument.

            :param term: The term to insert in a lambda function.
            :type term: Term
            :return: The new Lambda function.
            :rtype: Lambda
        """
        return Lambda(
            lambda x : traverse_bottom_up(
                f=lambda t : x if t is self.var else t,
                term=term,
            )
        )

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

            :param value: The value to replace the variable with.
            :type value: Term
            :return: A resolved lambda where the term has been substituted.
            :rtype: Term
        """
        return traverse_bottom_up(
            f=lambda t : value if t is self.var else t,
            term=self.out,
        )
    
    def to_str(self, vars: dict[str, Var]) -> str:
        char = ''

        for c in USABLE_VAR_NAMES:
            if c not in vars:
                char = c
                vars[c] = self.var
                break
        else:
            raise ValueError(
                f"Cannot represent a lambda function that's over {len(USABLE_VAR_NAMES)} layers deep!"
            )
        
        s = self.out.to_str(vars=vars)

        del vars[char]
        return "λ" + char + "." + (f"({s})" if " " in s else s)

L1 = lambda f : Lambda(lambda x : f(x))
L2 = lambda f : Lambda(lambda x : L1(lambda y : f(x, y)))
L3 = lambda f : Lambda(lambda x : L2(lambda y, z : f(x, y, z)))
L4 = lambda f : Lambda(lambda x : L3(lambda y, z, a : f(x, y, z, a)))
L5 = lambda f : Lambda(lambda x : L4(lambda y, z, a, b : f(x, y, z, a, b)))
L6 = lambda f : Lambda(lambda x : L5(lambda y, z, a, b, c : f(x, y, z, a, b, c)))
L7 = lambda f : Lambda(lambda x : L6(lambda y, z, a, b, c, d : f(x, y, z, a, b, c, d)))
L8 = lambda f : Lambda(lambda x : L7(lambda y, z, a, b, c, d, e : f(x, y, z, a, b, c, d, e)))

@dataclass(frozen=True)
class Var(Term):

    """
        A variable is a value that can be substituted by a lambda function.

        Each lambda function stores a unique var, which can then point to a new
        structure later on.
    """

    def __repr__(self) -> str:
        """
            Create a representation of a term.
        """
        return self.to_str(vars={})

    def to_str(self, vars: dict[str, Var]) -> str:
        for i, v in vars.items():
            if v is self:
                return i
        else:
            return "[ERROR:UNKNOWN_VAR]"

if __name__ == "__main__":
    print("REPRESENTATIONS")
    print(30 * "=")

    # x
    print(Const("x"))

    # λx.(f x)
    print(Lambda(lambda a : App(f=Const("f"), x=a)))

    # λx.(λy.(x y))
    print(Lambda(lambda b : Lambda(lambda c : App(f=b, x=c))))

    # Nat.Succ (Nat.Succ Nat.Zero)
    print(App(f=Const("Nat.Succ"), x=App(f=Const("Nat.Succ"), x=Const("Nat.Zero"))))

    
    print("\nNORMALIZATIONS")
    print(30 * "=")

    # foo
    a = App(f=Lambda(lambda x : x), x=Const("foo"))
    print(a)
    print(a.normalize())

    # f x
    b = App(
        f=App(
            f=Lambda(
                lambda f : Lambda(
                    lambda x : App(
                        f=f, x=x
                    )
                )
            ),
            x=Const("f"),
        ),
        x= Const("x"),
    )
    print(b)
    print(b.normalize())

    # a b c
    c = App(
        f=App(
            f=App(
                f=Lambda(
                    lambda x : Lambda(
                        lambda y : Lambda(
                            lambda z : App(
                                f=App(f=x, x=y), x=z
                            )
                        )
                    )
                ),
                x=Const("a"),
            ),
            x=Const("b"),
        ),
        x=Const("c"),
    )
    print(c)
    print(c.normalize())

    #
    d = L3(lambda x, y, z : A2(x, y, z))
    print(d)
    print(d.normalize())

    
    print("\nSHORTHAND NOTATION")
    print(30 * "=")

    # λx.(λy.(λz.(x y z))
    e = L3(lambda a, b, c : App(f=App(f=a, x=b), x=c))
    print(e)


    # λx.(λy.(λz.(x y z))
    f = L3(lambda a, b, c : A3(L3(lambda a, b, c : A2(a, b, c)), a, b, c))
    print(f)
    print(f.normalize())
    assert e == f.normalize()

    # not
    bool_t = L2(lambda x, y : x).normalize()
    bool_f = L2(lambda x, y : y).normalize()
    g = L1(lambda b : A2(b, bool_f, bool_t))
    print(g)
    print(g.normalize())