import itertools
from typing import List, Iterable

from PIL import Image
from funcy import lsplit, lfilter

from cv_analysis.logging import logger
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.randomization import rnd
from synthesis.segment.content_rectangle import ContentRectangle
from synthesis.segment.recursive_content_rectangle import RecursiveContentRectangle
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)

        boxes = list(unpack_boxes(boxes))
        for b in boxes:
            logger.trace(f"Generated {b}")

        return boxes


class BoxChildrenVisitor:
    def visit_content_rectangle(self, _box: ContentRectangle):
        return []

    def visit_recursive_content_rectangle(self, box: RecursiveContentRectangle):
        return box.children


def unpack_boxes(boxes: Iterable[ContentRectangle], recursed=False) -> Iterable[ContentRectangle]:
    yield from itertools.chain.from_iterable(map(lambda b: unpack_box(b, recursed), boxes))


def unpack_box(box: ContentRectangle, recursed=False) -> Iterable[ContentRectangle]:
    children = box.accept(BoxChildrenVisitor())

    def is_a_leaf():
        return not children

    def is_an_internal_node():
        return recursed and children

    def is_a_root_node():
        return not recursed and children

    if is_a_root_node():
        yield box
        yield from unpack_boxes(children, True)

    elif is_an_internal_node():
        yield from unpack_boxes(children, True)

    elif is_a_leaf():
        yield box

    else:
        raise ValueError("This should not happen")