'Split an optree to stacks.'

@memoize_stacks
def _split_to_stacks(self, optree):

    whole_stack = optree.traverse(list())
    stages = []
    main_stage = []
    main_stage_axis = []
    axis_count = [0, 0]
    for s in whole_stack:
        if (isinstance(s, dict) and (s['op'] in OpCollection.reduction_ops)):
            assert ((s['axis'] == 0) or (s['axis'] == 1))
            axis_count[s['axis']] += 1
    minority_axis = (0 if (axis_count[0] <= axis_count[1]) else 1)
    for s in whole_stack:
        if isinstance(s, dict):
            if (s['op'] == 'dot'):
                right = main_stage.pop()
                main_stage_axis.pop()
                left = main_stage.pop()
                main_stage_axis.pop()
                if isinstance(left, OpTreeNode):
                    left_buf = self._buf_malloc(left.shape)
                    stages.append(OpTreeNode({'op': 'assign'}, left_buf, left))
                    left = left_buf
                if isinstance(right, OpTreeNode):
                    right_buf = self._buf_malloc(right.shape)
                    stages.append(OpTreeNode({'op': 'assign'}, right_buf, right))
                    right = right_buf
                buf = self._buf_malloc((left.shape[0], right.shape[1]))
                stages.append(OpTreeNode({'op': 'assign'}, buf, OpTreeNode(s, left, right)))
                main_stage.append(buf)
                main_stage_axis.append(None)
            elif (s['op'] == 'transpose'):
                operand = main_stage.pop()
                main_stage_axis.pop()
                buf = self._buf_malloc(operand.shape)
                stages.append(OpTreeNode({'op': 'assign'}, buf, operand))
                main_stage.append(buf.T)
                main_stage_axis.append(None)
            elif (s['op'] in OpCollection.reduction_ops):
                assert (s['axis'] is not None)
                operand = main_stage.pop()
                prev_axis = main_stage_axis.pop()
                if ((prev_axis is not None) and (prev_axis != s['axis'])):
                    buf = self._buf_malloc(operand.shape)
                    stages.append(OpTreeNode({'op': 'assign'}, buf, operand))
                    main_stage.append(OpTreeNode(s, buf, None))
                    main_stage_axis.append(s['axis'])
                else:
                    main_stage.append(OpTreeNode(s, operand, None))
                    main_stage_axis.append(s['axis'])
            elif (s['op'] in OpCollection.unary_ops):
                operand = main_stage.pop()
                axis = main_stage_axis.pop()
                main_stage.append(OpTreeNode(s, operand, None))
                main_stage_axis.append(axis)
            elif (s['op'] in OpCollection.binary_ops):
                right = main_stage.pop()
                prev_axis_right = main_stage_axis.pop()
                left = main_stage.pop()
                prev_axis_left = main_stage_axis.pop()
                if ((prev_axis_right is not None) and (prev_axis_left is not None) and (prev_axis_left != prev_axis_right)):
                    if (prev_axis_left == minority_axis):
                        buf = self._buf_malloc(left.shape)
                        stages.append(OpTreeNode({'op': 'assign'}, buf, left))
                        left = buf
                        axis = prev_axis_right
                    else:
                        buf = self._buf_malloc(right.shape)
                        stages.append(OpTreeNode({'op': 'assign'}, buf, right))
                        right = buf
                        axis = prev_axis_left
                    main_stage.append(OpTreeNode(s, left, right))
                    main_stage_axis.append(axis)
                else:
                    main_stage.append(OpTreeNode(s, left, right))
                    axis = None
                    if (prev_axis_left is not None):
                        axis = prev_axis_left
                    else:
                        axis = prev_axis_right
                    main_stage_axis.append(axis)
            else:
                return NotImplemented
        else:
            main_stage.append(s)
            main_stage_axis.append(None)
    stages.append(main_stage[0])
    stacks = []
    for stage in stages:
        assert isinstance(stage, OpTreeNode)
        stacks.append(stage.traverse(list()))
    self._buf_free()
    return stacks
