from lambda_calculus import Variable, Abstraction, Application
from lambda_calculus.visitors.normalisation import BetaNormalisingVisitor
import networkx as nx
import sys
import itertools

S = Abstraction('x', 
        Abstraction('y', 
            Abstraction('z', 
                Application(Application(Variable('x'), Variable('z')),
                     Application(Variable('y'), Variable('z'))))))

K = Abstraction('x', 
        Abstraction('y',
            Variable('x')))


class Node(object):
  def __init__(self, left, right):
    self.left = left
    self.right = right

  def __repr__(self):
    return '(%s %s)' % (self.left if isinstance(self.left, Node) else self.left[1], self.right if isinstance(self.right, Node) else self.right[1])

  def app(self):
    return Application(
        self.left.app() if isinstance(self.left, Node) else self.left[0],
        self.right.app() if isinstance(self.right, Node) else self.right[0])

# Given a tree and a label, yields every possible augmentation of the tree by
# adding a new node with the label as a child "above" some existing Node or Leaf.
def add_leaf(tree, label):
  yield Node(label, tree)
  if isinstance(tree, Node):
    for left in add_leaf(tree.left, label):
      yield Node(left, tree.right)
    for right in add_leaf(tree.right, label):
      yield Node(tree.left, right)

# Given a list of labels, yield each rooted, unordered full binary tree with
# the specified labels.
def enum_unordered(labels):
  if len(labels) == 1:
    yield labels[0]
  else:
    for tree in enum_unordered(labels[1:]):
      for new_tree in add_leaf(tree, labels[0]):
        yield new_tree

def enum_ordered(labels):
  if len(labels) == 1:
    yield labels[0]
  else:
    for i in range(1, len(labels)):
      for left in enum_ordered(labels[:i]):
        for right in enum_ordered(labels[i:]):
          yield Node(left, right)

# def count(n,elements):
#     l = len(elements)
#     def Convert(n,base):
#        string = elements
#        if n < base:
#           return [string[n]]
#        else:
#           return Convert(n//base,base) + [string[n%base]]
#     return (Convert(i,l) for i in range(0,l**n+1))

sys.setrecursionlimit(1000)

all_products = lambda *args, repeat:itertools.chain(*(itertools.product(*args, repeat=n) for n in range(1,repeat+1)))

g = nx.Graph()
for term in all_products([(S,'S'), (K,'K')], repeat=5):
    for parensed in enum_ordered(term):
        try:
            beta = BetaNormalisingVisitor().skip_intermediate(
                parensed.app() if isinstance(parensed, Node) else parensed[0]
            )
            g.add_edge(repr(parensed) if isinstance(parensed, Node) else parensed[1], str(beta))
        except:
          print(str(beta), "cant be reduced")
          pass

import matplotlib.pyplot as plt
# # nx.nx_pydot.write_dot(g, 'out.dot')
nx.draw(g, with_labels=True)
plt.show()

# for n1 in [S,K]:
#     for n2 in [S,K]:
#         for n3 in [S,K]:
#             # for n4 in [S,K]:
#             for term in enum_ordered([n1,n2,n3]):
#                 beta = BetaNormalisingVisitor().skip_intermediate(term) 
#                 if beta.body == Variable(beta.bound):
#                     print(str(term), 'reduces to identity')