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)

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.draw(g, with_labels=True)
plt.show()