from functools import partial

import cv2
import numpy as np
import pdf2image
from iteration_utilities import starfilter, first

from vidocp.utils.display import show_mpl
from vidocp.utils.draw import draw_contours
from vidocp.utils.filters import is_large_enough, is_filled, is_boxy


def is_likely_redaction(contour, hierarchy, min_area):
    return is_filled(hierarchy) and is_boxy(contour) and is_large_enough(contour, min_area)


def find_redactions(image: np.array, min_normalized_area=200000):

    min_normalized_area /= 200  # Assumes 200 DPI PDF -> image conversion resolution

    if len(image.shape) > 2:
        gray = ~cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = ~image
    blurred = cv2.GaussianBlur(gray, (5, 5), 1)
    thresh = cv2.threshold(blurred, 252, 255, cv2.THRESH_BINARY)[1]

    contours, hierarchies = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    try:
        contours = map(
            first, starfilter(partial(is_likely_redaction, min_area=min_normalized_area), zip(contours, hierarchies[0]))
        )
        return list(contours)
    except:
        return []


def annotate_redactions_in_pdf(pdf_path, page_index=1, show=True):

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

    redaction_contours = find_redactions(page)
    page = draw_contours(page, redaction_contours)

    if show:
        show_mpl(page)
    else:
        return page