jQuery(document).ready(function($) {

    // false means 'collapsed', set true to 'open'
    var open = {
        'level2': false,
        'level3': false
    };

    /**
     * Calculate depth of menu element.
     **/
    function depth(element, root) {
        var depth = 1;
        while (true) {
            if (element[0] == root[0]) {
                break;
            }
            depth++;
            element = $(element).parent().parent();
        }
        return depth;
    }

    function setButtonClose() {
        $(this).attr('class', 'closed');
        $(this).html('▶');
    }

    function setButtonOpen() {
        $(this).attr('class', 'open');
        $(this).html('▼');
    }

    // Add toggle button to all menu elements that have children pages.
    $("#menu > li > ul > li").each(function() {
        var root = $(this).parent();
        $("ul.children", this).each(function() {
            var parentListItem = $(this).parent(),
                currentDepth = depth(this, root);

            parentListItem.addClass("has-children");
            parentListItem.prepend('<button></button>');
            var buttons = $("button", parentListItem);

            // by default all elements are open
            buttons.each(setButtonOpen);

            // closes elements by config
            if (typeof open['level' + currentDepth] !== 'undefined' && open['level' + currentDepth] === false) {
                setButtonClose.call($("button", parentListItem));
                $(this).toggle();
            }

            // prevent collapsing the submenu if it contains the current page item
            var has_active_child = $('.current_page_item', this).length > 0;
            if (has_active_child === true || parentListItem.hasClass('current_page_item')) {
                if ($("button", parentListItem).hasClass('closed')) {
                    setButtonOpen.call($("button", parentListItem));
                    $(this).toggle();
                }
            }

            // toggle elements on button click
            var currentMenu = this;
            buttons.click(function(evt) {
                if ($(this).hasClass('open')) {
                    setButtonClose.call(this);
                } else {
                    setButtonOpen.call(this);
                }
                $(currentMenu).toggle();
            });
        });
    });
});