Мега меню Walker

Я пытаюсь создать walker для мега меню. К сожалению, walkers полностью выходят за пределы моих знаний в программировании. Мне очень нужна помощь в его реализации. Вот необходимые функции:

• Обернуть <ul> второго уровня в <section>. [ГОТОВО]
• Когда пользователь устанавливает класс "break" на <li> второго уровня <ul>, сделать так, чтобы этот <li> начинал новый <ul>. Если это первый <li> в списке, ничего не делать, чтобы избежать создания пустых ненумерованных списков. [ГОТОВО]
• Когда пользователь устанавливает класс "widget" на <li> первого уровня, который имеет подменю <ul>, добавить виджет в конец этого <ul>. [ГОТОВО]
• Добавить класс mega-menu-columns-# к элементам <li> первого уровня, которые содержат выпадающие меню с несколькими колонками и/или виджет. # представляет количество элементов <ul>, +1 для виджета, если он существует. [ГОТОВО]

У меня есть часть кода, который выполняет некоторые из этих функций, но не все. Ниже приведены вырезанные секции:

Обертка <ul> второго уровня в <section>:

function start_lvl(&$output, $depth = 0, $args = array()) {
    if ($depth == 0) {
        $output .= "<section>";
    }
    $output .= "<ul class=\"sub-menu\">";
}
function end_lvl(&$output, $depth = 0, $args = array()) {
    $output .= "</ul>";
    if ($depth == 0) {
        $output .= "</section>\n";
    }
}

Генерация HTML виджета:

ob_start();
dynamic_sidebar("Navigation Callout");
$widget = ob_get_contents();
ob_end_clean();

Пример выходного HTML:

<ul>
    <li id="menu-item-1" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-has-children menu-item-1 mega-menu-columns-2">
        <a href="http://www.example.com/about/">
            О компании
        </a>
        <section>
            <ul class="sub-menu">
                <li id="menu-item-2" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-2">
                    <a href="http://www.example.com/about/company-profile/">
                        Профиль компании
                    </a>
                </li>
                <li id="menu-item-3" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-3">
                    <a href="http://www.example.com/about/leadership-team/">
                        Руководство
                    </a>
                </li>
                <li id="menu-item-4" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-4">
                    <a href="http://www.example.com/about/professional-affiliations/">
                        Профессиональные аффилиации
                    </a>
                </li>
            </ul>
            <ul class="sub-menu">
                <li id="menu-item-5" class="break menu-item menu-item-type-post_type menu-item-object-page menu-item-5">
                    <a href="http://www.example.com/about/clients/">
                        Клиенты
                    </a>
                </li>
                <li id="menu-item-6" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-6">
                    <a href="http://www.example.com/about/partnerships/">
                        Партнерства
                    </a>
                </li>
            </ul>
        </section>
    </li>
    <li id="menu-item-7" class="widget menu-item menu-item-type-post_type menu-item-object-page menu-item-has-children menu-item-7 mega-menu-columns-3">
        <a href="http://www.example.com/services/">
            Услуги
        </a>
        <section>
            <ul class="sub-menu">
                <li id="menu-item-8" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-8">
                    <a href="http://www.example.com/services/civil-engineering/">
                        Гражданское строительство
                    </a>
                </li>
                <li id="menu-item-9" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-9">
                    <a href="http://www.example.com/services/land-planning/">
                        Планирование землепользования
                    </a>
                </li>
                <li id="menu-item-10" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-10">
                    <a href="http://www.example.com/services/surveying/">
                        Геодезия
                    </a>
                </li>
            </ul>
            <ul class="sub-menu">
                <li id="menu-item-11" class="break menu-item menu-item-type-post_type menu-item-object-page menu-item-11">
                    <a href="http://www.example.com/services/information-technology/">
                        Информационные технологии
                    </a>
                </li>
                <li id="menu-item-12" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-12">
                    <a href="http://www.example.com/services/subsurface-utility-engineering/">
                        Подземные коммуникации
                    </a>
                </li>
            </ul>
            <aside>
                <h6>Заголовок виджета</h6>
                <p>Maecenas quis semper arcu. Quisque consequat risus nisi. Sed venenatis urna porta eros malesuada euismod. Nulla sollicitudin fringilla posuere. Nulla et tellus eu nisi sodales convallis non vel tellus.</p>
            </aside>
        </section>
    </li>
    <li id="menu-item-13" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-15">
        <a href="http://www.example.com/contact/">
            Контакты
        </a>
    </li>
</ul>

ОБНОВЛЕНИЕ: У меня проблемы со счетчиками. Они считают только после того, как подменю было сгенерировано, что мне не помогает. Посмотрите этот скриншот, чтобы понять, что я имею в виду:

Верхние числа берутся в start_el. Нижние числа берутся в end_el. Как видите, верхние числа не считают мои .breaks как должны. Они считают класс widget, потому что те считаются в $depth = 0. Пожалуйста, спасите меня от этого ужаса!

// walker для мега меню
/*
    ОСТАВШАЯСЯ ОШИБКА:
    - Нужно добавить класс к LI содержащему mega-menu-columns-#
*/
class megaMenuWalker extends Walker_Nav_Menu {
    private $column_limit = 3; /* нужно установить для каждого сайта */
    private $show_widget = false;
    private $column_count = 0;
    static $li_count = 0;
    function start_el(&$output, $item, $depth = 0, $args = array(), $id = 0) {
        $classes = empty($item->classes) ? array() : (array) $item->classes;
        $item_id = $item->ID;
        if ($depth == 0) self::$li_count = 0;
        if ($depth == 0 && in_array("widget", $classes)) {
            $this->show_widget = true;
            $this->column_count++;
        }
        if ($depth == 1 && self::$li_count == 1) {
            $this->column_count++;
        }
        if ($depth == 1 && in_array("break", $classes) && self::$li_count != 1 && $this->column_count < $this->column_limit) {
            $output .= "</ul><ul class=\"sub-menu\">";
            $this->column_count++;
        }
        if ($depth == 0 && $this->column_count > 0) {
            $mega_menu_class = " mega-menu-columns-" . $this->column_count;
        }
        $class_names = join(" ", apply_filters("nav_menu_css_class", array_filter($classes), $item));
        $class_names = " class=\"" . esc_attr($class_names . $mega_menu_class) . "\"";
        $output .= sprintf(
            "<li id=\"menu-item-%s\"%s><a href=\"%s\">%s</a>",
            $item_id,
            $class_names,
            $item->url,
            $item->title
        );
        self::$li_count++;
    }
    function start_lvl(&$output, $depth = 0, $args = array()) {
        if ($depth == 0) {
            $output .= "<section>";
        }
        $output .= "<ul class=\"sub-menu\">";
    }
    function end_lvl(&$output, $depth = 0, $args = array()) {
        $output .= "</ul>";
        if ($depth == 0) {
            if ($this->show_widget) {
                ob_start();
                dynamic_sidebar("Navigation Callout");
                $widget = ob_get_contents();
                ob_end_clean();
                $output .= $widget;
                $this->show_widget = false;
            }
            $output .= "</section>";
        }
    }
    function end_el(&$output, $item, $depth = 0, $args = array(), $id = 0) {
        if ($depth == 0 && $this->column_count > 0) {
            /* нужно добавить к открывающему li уровня 0 */
            $column_count_class = " mega-menu-columns-" . $this->column_count;
            $output .= $column_count_class;
            /* конец */
            $this->column_count = 0;
        }
        $output .= "</li>";
    }
}   

ОБНОВЛЕНИЕ 2: Вот пример вывода с комментариями, описывающими как должен считать класс mega-menu-columns-:

<ul>
    <!-- +1 потому что имеет класс "widget" -->
    <li id="menu-item-1" class="widget menu-item menu-item-type-post_type menu-item-object-page menu-item-has-children menu-item-1 mega-menu-columns-3">
        <a href="http://www.example.com/about/">
            О компании
        </a>
        <!-- +1 потому что существует выпадающее меню -->
        <!-- добавляется моим walker'ом -->
        <section>
        <!-- конец добавляется моим walker'ом -->
            <ul class="sub-menu">
                <!-- +0 потому что этот "break" первый дочерний элемент -->
                <li id="menu-item-2" class="break menu-item menu-item-type-post_type menu-item-object-page menu-item-2">
                    <a href="http://www.example.com/about/company-profile/">
                        Профиль компании
                    </a>
                    <ul>
                        <!-- +0 потому что этот "break" на уровне 2 -->
                        <li id="menu-item-3" class="break menu-item menu-item-type-post_type menu-item-object-page menu-item-3">
                            <a href="http://www.example.com/about/our-team/">
                                Наша команда
                            </a>
                        </li>
                    </ul>
                </li>
                <li id="menu-item-4" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-4">
                    <a href="http://www.example.com/about/leadership-team/">
                        Руководство
                    </a>
                </li>
                <li id="menu-item-5" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-5">
                    <a href="http://www.example.com/about/professional-affiliations/">
                        Профессиональные аффилиации
                    </a>
                </li>
            <!-- добавляется моим walker'ом -->
            </ul>
            <ul class="sub-menu">
            <!-- конец добавляется моим walker'ом -->
                <!-- +1 потому что этот "break" на уровне 1 и не первый дочерний элемент -->
                <li id="menu-item-6" class="break menu-item menu-item-type-post_type menu-item-object-page menu-item-6">
                    <a href="http://www.example.com/about/clients/">
                        Клиенты
                    </a>
                </li>
                <li id="menu-item-7" class="menu-item menu-item-type-post_type menu-item-object-page menu-item-7">
                    <a href="http://www.example.com/about/partnerships/">
                        Партнерства
                    </a>
                </li>
            </ul>
            <!-- добавляется моим walker'ом для .widget -->
            <section>
                <header>
                    <h1>Заголовок виджета</h1>
                </header>
                <p>Это виджет. Было сложно заставить его появиться!</p>
            </section>
            <!-- конец добавляется моим walker'ом для .widget -->
        <!-- добавляется моим walker'ом -->
        </section>
        <!-- конец добавляется моим walker'ом -->
    </li>
</ul>

ОБНОВЛЕНИЕ: Вот мой финальный Walker и функции. Он делает точно то, что я хотел. Спасибо за помощь!

// walker для мега меню
class megaMenuWalker extends Walker_Nav_Menu {
    private $column_limit = 3;
    private $show_widget = false;
    private $column_count = 0;
    static $li_count = 0;
    function start_el(&$output, $item, $depth = 0, $args = array(), $id = 0) {
        $classes = empty($item->classes) ? array() : (array) $item->classes;
        $item_id = $item->ID;
        if ($depth == 0) {
            self::$li_count = 0;
        }
        if ($depth == 0 && in_array("widget", $classes)) {
            $this->show_widget = true;
            $this->column_count++;
        }
        if ($depth == 1 && self::$li_count == 1) {
            $this->column_count++;
        }
        if ($depth == 1 && in_array("break", $classes) && self::$li_count != 1 && $this->column_count < $this->column_limit) {
            $output .= "</ul><ul class=\"sub-menu\">";
            $this->column_count++;
        }
        $class_names = join(" ", apply_filters("nav_menu_css_class", array_filter($classes), $item)); // настройка массива классов для добавления к каждому li
        $class_names = " class=\"" . esc_attr($class_names) . "\"";
        $output .= sprintf(
            "<li id=\"menu-item-%s\"%s><a href=\"%s\">%s</a>",
            $item_id,
            $class_names,
            $item->url,
            $item->title
        );
        self::$li_count++;
    }
    function start_lvl(&$output, $depth = 0, $args = array()) {
        if ($depth == 0) {
            $output .= "<section>";
        }
        $output .= "<ul class=\"sub-menu\">";
    }
    function end_lvl(&$output, $depth = 0, $args = array()) {
        $output .= "</ul>";
        if ($depth == 0) {
            if ($this->show_widget) {
                ob_start();
                dynamic_sidebar("Navigation Callout");
                $widget = ob_get_contents();
                ob_end_clean();
                $output .= $widget;
                $this->show_widget = false;
            }
            $output .= "</section>";
        }
    }
    function end_el(&$output, $item, $depth = 0, $args = array(), $id = 0) {
        if ($depth == 0 && $this->column_count > 0) {
            $this->column_count = 0;
        }
        $output .= "</li>";
    }
}

// добавление классов mega-menu-columns-#
function add_column_number($items, $args) {
    static $column_limit = 3;
    static $post_id = 0;
    static $x_key = 0;
    static $column_count = 0;
    static $li_count = 0;
    $tmp = array();
    foreach($items as $key => $item) {
        if (0 == $item->menu_item_parent) {
            $x_key = $key;
            $post_id = $item->ID;
            $column_count = 0;
            $li_count = 0;
            if (in_array("widget", $item->classes, 1)) {
                $column_count++;
            }
        }
        if ($post_id == $item->menu_item_parent) {
            $li_count++;
            if ($column_count < $column_limit && $li_count == 1) {
                $column_count++;
            }
            if (in_array("break", $item->classes, 1) && $li_count > 1 && $column_count < $column_limit) {
                $column_count++;
            }
            $tmp[$x_key] = $column_count;
        }
    }
    foreach($tmp as $key => $value) {
        $items[$key]->classes[] = sprintf("mega-menu-columns-%d", $value);
    }
    unset($tmp);
    return $items;
};

// добавление классов колонок
add_filter("wp_nav_menu_args", function($args) {
    if ($args["walker"] instanceof megaMenuWalker) {
        add_filter("wp_nav_menu_objects", "add_column_number");
    }
    return $args;
});

// остановка функции добавления классов
add_filter("wp_nav_menu", function( $nav_menu ) {
    remove_filter("wp_nav_menu_objects", "add_column_number");
    return $nav_menu;
});