from functools import reduce, partial
from typing import Iterable

import cv2
import numpy as np
from funcy import compose, rcompose, first, lkeep

from cv_analysis.utils.connect_rects import connect_related_rectangles
from cv_analysis.utils.conversion import box_to_rectangle
from cv_analysis.utils.postprocessing import remove_included, has_no_parent
from cv_analysis.utils.rectangle import Rectangle


def parse_layout(image: np.array):

    rectangles = find_segments(image)
    rectangles = remove_included(rectangles)
    rectangles = connect_related_rectangles(rectangles)
    rectangles = remove_included(rectangles)

    return rectangles


def find_segments(image):
    rectangles = rcompose(
        prepare_for_initial_detection,
        __find_segments,
        partial(prepare_for_meta_detection, image.copy()),
        __find_segments,
    )(image)

    return rectangles


def prepare_for_initial_detection(image: np.ndarray):
    return compose(dilate_page_components, normalize_to_gray_scale)(image)


def __find_segments(image):
    def to_rectangle_if_valid(contour, hierarchy):
        return (
            box_to_rectangle(cv2.boundingRect(contour))
            if is_likely_segment(contour) and has_no_parent(hierarchy)
            else None
        )

    rectangles = lkeep(map(to_rectangle_if_valid, *find_contours(image)))

    return rectangles


def find_contours(image):
    contours, hierarchies = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return contours, first(hierarchies)


def is_likely_segment(rect, min_area=100):
    # FIXME: Parameterize via factory
    return cv2.contourArea(rect, False) > min_area


def dilate_page_components(image):
    # FIXME: Parameterize via factory
    image = cv2.GaussianBlur(image, (7, 7), 0)
    # FIXME: Parameterize via factory
    thresh = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
    # FIXME: Parameterize via factory
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    # FIXME: Parameterize via factory
    return cv2.dilate(thresh, kernel, iterations=4)


def prepare_for_meta_detection(image: np.ndarray, rectangles: Iterable[Rectangle]):

    image = fill_rectangles(image, rectangles)
    image = threshold_image(image)
    image = invert_image(image)
    image = normalize_to_gray_scale(image)

    return image


def normalize_to_gray_scale(image):
    if len(image.shape) > 2:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return image


def threshold_image(image):
    # FIXME: Parameterize via factory
    _, image = cv2.threshold(image, 254, 255, cv2.THRESH_BINARY)
    return image


def invert_image(image):
    return ~image


def fill_rectangles(image, rectangles):
    image = reduce(fill_in_component_area, rectangles, image)
    return image


def fill_in_component_area(image, rect):
    x, y, w, h = rect
    cv2.rectangle(image, (x, y), (x + w, y + h), (0, 0, 0), -1)
    cv2.rectangle(image, (x, y), (x + w, y + h), (255, 255, 255), 7)

    return image