"""
   an attempt to unify a variable with a formula that contains that variable
   (for example X = f(X, k)) must fail, otherwise X would try to become the
   "infinite" structure f(f(f(f(f(..., k), k), k), k), k).
   Prolog doesn't bother with this "occurs check" so I didn't mention it in
   class, but for completentess I have added it here.
   Using a stack instead of recursion seems to have simplified the design.
"""

def unify(a, b):
  stack = [(a, b)]
  unif = {}
  while stack:
    (a, b) = stack.pop()
    if type(a) == int and not contains(b, a):
      stack = subst(stack, a, b)
      unif = subst(unif, a, b)
      unif[a] = b
    elif type(b) == int and not contains(a, b):
      stack = subst(stack, b, a)
      unif = subst(unif, b, a)
      unif[b] = a
    elif type(a) == str or type(b) == str:
      if a != b:
        return False
      else:
        pass
    elif type(a) == tuple and type(b) == tuple:
      if a[0] != b[0]:
        return False
      elif len(a) != len(b):
        return False
      else:
        for i in range(1, len(a)):
          stack.append((a[i], b[i]))
    else:
      print("(unify) bad pair", a, b)
      return False
  return unif

def subst(f, var, val):
  if not contains(f, var):
    return f
  if type(f) == int:
    if f == var:
      return val
    else:
      return f
  elif type(f) == str:
    return f
  elif type(f) == tuple:
    return tuple(subst(p, var, val) for p in f)
  elif type(f) == list:
    return [(subst(a, var, val), subst(b, var, val)) for (a, b) in f]
  elif type(f) == dict:
    return {k : subst(f[k], var, val) for k in f}
  else:
    print("(subst) bad type", f)
    return False

def contains(f, var):
  if type(f) == int:
    return f == var
  elif type(f) == str:
    return False
  elif type(f) == tuple or type(f) == list:
    for p in f:
      if contains(p, var):
        return True
    return False
  elif type(f) == dict:
    for k in f:
      if contains(f[k], var):
        return True
    return False
  else:
    print("(contains) bad type", f)
    return True

f1 = ("f", 1,   ("g", "z",   3),   4)
f2 = ("f", "x", ("g", 2,   "w"), "y")
r = unify(f1, f2)
if r == False:
  print("FAIL")
elif type(r) == dict:
  for k in r:
    print(k, "=", r[k])
else:
  print("bad result", r)