from cv_analysis.utils.rectangle import Rectangle
from synthesis.partitioner.page_partitioner import PagePartitioner
from synthesis.random import rnd


class TwoColumnPagePartitioner(PagePartitioner):
    def __init__(self):
        super().__init__()
        self.max_recursion_depth = 3

    def generate_content_boxes(self, box: Rectangle, depth=0):
        if depth >= self.max_recursion_depth:
            yield box

        else:
            if depth == 0:
                axis = "x"
                split_percentage = 0.5
            else:
                axis = "y"
                split_percentage = rnd.choice([0.3, 0.7])

            child_boxes = self.generate_child_boxes(box, axis=axis, split_percentage=split_percentage)

            yield from (self.generate_content_boxes(b, depth + 1) for b in child_boxes)