""" Strategies to Traverse a Tree """
from __future__ import print_function, division

from .util import basic_fns
from sympy.strategies.core import chain, do_one

def top_down(rule, fns=basic_fns):
    """ Apply a rule down a tree running it on the top nodes first """
    return chain(rule, lambda expr: sall(top_down(rule, fns), fns)(expr))

def bottom_up(rule, fns=basic_fns):
    """ Apply a rule down a tree running it on the bottom nodes first """
    return chain(lambda expr: sall(bottom_up(rule, fns), fns)(expr), rule)

def top_down_once(rule, fns=basic_fns):
    """ Apply a rule down a tree - stop on success """
    return do_one(rule, lambda expr: sall(top_down(rule, fns), fns)(expr))

def bottom_up_once(rule, fns=basic_fns):
    """ Apply a rule up a tree - stop on success """
    return do_one(lambda expr: sall(bottom_up(rule, fns), fns)(expr), rule)

def sall(rule, fns=basic_fns):
    """ Strategic all - apply rule to args """
    op, new, children, leaf = map(fns.get, ('op', 'new', 'children', 'leaf'))
    def all_rl(expr):
        if leaf(expr):
            return expr
        else:
            args = map(rule, children(expr))
            return new(op(expr), *args)
    return all_rl