from cv_analysis.layout_parsing import annotate_layout_in_pdf
from cv_analysis.figure_detection import figures_in_image, detect_figures
from cv_analysis.table_parsing import tables_in_image
from cv_analysis.utils.text import find_primary_text_regions, remove_primary_text_regions
from cv_analysis.utils.draw import draw_rectangles
from cv_analysis.utils.display import show_mpl


def detect_parting_line(image):
    pass


def cut_out_content_structures(layout_rects, page):
    large_enough_rects = []
    too_small_rects = []
    for x, y, w, h in layout_rects:
        rect = (x, y, w, h)
        if w * h >= 100000:
            cropped_page = page[y:(y + h), x:(x + w)]
            large_enough_rects.append([rect, cropped_page])
        else:
            cropped_page = page[y:(y + h), x:(x + w)]
            too_small_rects.append([rect, cropped_page])
    return large_enough_rects, too_small_rects


def parse_and_label_content_structures(page, large_enough_rects, too_small_rects):
    for coordinates, cropped_image in large_enough_rects:
        non_text_rects = detect_figures(cropped_image)
        print(len(non_text_rects), len(list(non_text_rects)))
        if len(non_text_rects) == 0:
            page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True)
        elif tables_in_image(cropped_image)[0]:
            page = draw_rectangles(page, [coordinates], color=(255, 0, 0), annotate=True)
        else:
            page = draw_rectangles(page, [coordinates], color=(0, 0, 255), annotate=True)

    # for coordinates, cropped_image in too_small_rects:
    #     non_text_rects = detect_figures(cropped_image)
    #     if len(non_text_rects) == 0 and len(list(find_primary_text_regions(cropped_image))) > 0:
    #         page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True)
    #     else:
    #         page = draw_rectangles(page, [coordinates], color=(0, 255, 255), annotate=True)
    return page


def detect_figures_over_layout():
    # pdf_path = "/home/lillian/PycharmProjects/ner_address/data/pdfs/syngenta/026c917f04660aaea4bb57d180f9598b.pdf"
    # pdf_path = "/home/lillian/ocr_docs/ocr1.pdf"
    pdf_path = "/home/lillian/ocr_docs/Report on spectra.pdf"
    #pdf_path = "/home/lillian/ocr_docs/VV-857853.pdf"
    page_index = 13
    layout_rects, page = annotate_layout_in_pdf(pdf_path, page_index, return_rects=True)
    big_structures, small_structures = cut_out_content_structures(layout_rects, page)
    page = parse_and_label_content_structures(page, big_structures, small_structures)
    show_mpl(page)


detect_figures_over_layout()