import itertools
from typing import List

from PIL import Image
from funcy import lsplit, lfilter

from cv_analysis.utils import every_nth, zipmap
from cv_analysis.utils.geometric import is_square_like
from cv_analysis.utils.merging import merge_related_rectangles
from cv_analysis.utils.postprocessing import remove_included, remove_overlapping
from cv_analysis.utils.rectangle import Rectangle
from synthesis.random import rnd
from synthesis.segment.segments import (
    generate_random_text_block,
    generate_recursive_random_table_with_caption,
    generate_random_plot_with_caption,
)


class ContentGenerator:
    def __init__(self):
        self.constrain_layouts = True

    def __call__(self, boxes: List[Rectangle]) -> Image:
        rnd.shuffle(boxes)

        figure_boxes, text_boxes = lsplit(is_square_like, boxes)

        if self.constrain_layouts:
            figure_boxes = merge_related_rectangles(figure_boxes)
            figure_boxes = lfilter(is_square_like, figure_boxes)
            text_boxes = merge_related_rectangles(text_boxes)

        boxes = list(
            itertools.chain(
                map(generate_random_text_block, every_nth(2, text_boxes)),
                *zipmap(generate_recursive_random_table_with_caption, every_nth(2, text_boxes[1:])),
                *zipmap(generate_recursive_random_table_with_caption, every_nth(2, figure_boxes)),
                *zipmap(generate_random_plot_with_caption, every_nth(2, figure_boxes[1:])),
            )
        )

        if self.constrain_layouts:
            boxes = remove_included(boxes)
            boxes = remove_overlapping(boxes)

        return boxes