from itertools import count

import cv2
import numpy as np
import pdf2image
from matplotlib import pyplot as plt
import imutils


def find_layout_boxes(image: np.array):

    gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray_scale, (5, 5), 1)
    thresh = cv2.threshold(blurred, 253, 255, cv2.THRESH_BINARY)[1]
    img_bin = ~thresh

    line_min_width = 10
    kernel_h = np.ones((10, line_min_width), np.uint8)
    kernel_v = np.ones((line_min_width, 10), np.uint8)

    img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h)
    img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v)

    img_bin_final = img_bin_h | img_bin_v

    contours = cv2.findContours(img_bin_final, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = imutils.grab_contours(contours)
    for c in contours:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.04 * peri, True)
        yield cv2.boundingRect(approx)


def annotate_layout_boxes(image, rects):
    for rect in rects:
        (x, y, w, h) = rect
        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

    return image


def annotate_layout_in_pdf(pdf_path, page_index=1):

    page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0]
    page = np.array(page)

    layout_boxes = find_layout_boxes(page)
    page = annotate_layout_boxes(page, layout_boxes)

    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(20, 20)
    ax.imshow(page)
    plt.show()