from dataclasses import asdict
from operator import truth

from funcy import lmap, flatten

from cv_analysis.figure_detection.figure_detection import detect_figures
from cv_analysis.table_parsing import parse_tables
from cv_analysis.utils.structures import Rectangle
from pdf2img.conversion import convert_pages_to_images
from pdf2img.default_objects.image import ImagePlus, ImageInfo
from pdf2img.default_objects.rectangle import RectanglePlus


def get_analysis_pipeline(operation):
    if operation == "table":
        return make_analysis_pipeline(parse_tables, table_parsing_formatter, dpi=200)
    elif operation == "figure":
        return make_analysis_pipeline(detect_figures, figure_detection_formatter, dpi=200)
    else:
        raise


def make_analysis_pipeline(analysis_fn, formatter, dpi):
    def analyse_pipeline(pdf: bytes, index=None):
        def parse_page(page: ImagePlus):
            image = page.asarray()
            rects = analysis_fn(image)
            if not rects:
                return
            infos = formatter(rects, page, dpi)
            return infos

        pages = convert_pages_to_images(pdf, index=index, dpi=dpi)
        results = map(parse_page, pages)

        yield from flatten(filter(truth, results))

    return analyse_pipeline


def table_parsing_formatter(rects, page, dpi):
    def format_rect(rect: Rectangle):
        rect_plus = RectanglePlus.from_pixels(*rect.xyxy(), page.info, alpha=False, dpi=dpi)
        return rect_plus.asdict(derotate=True)

    bboxes = lmap(format_rect, rects)

    return {"pageInfo": page.asdict(), "tableCells": bboxes}


def figure_detection_formatter(rects, page, dpi):
    def format_rect(rect: Rectangle):
        rect_plus = RectanglePlus.from_pixels(*rect.xyxy(), page.info, alpha=False, dpi=dpi)
        return asdict(ImageInfo(page.info, rect_plus.asbbox(derotate=False), rect_plus.alpha))

    return lmap(format_rect, rects)