import random
from copy import deepcopy
from itertools import chain
from operator import itemgetter

import fpdf
import pytest
from funcy import juxt, merge, rpartial

from image_prediction.image_extractor.extractor import ImageMetadataPair
from image_prediction.info import Info
from test.conftest import get_base_position_metadata, add_image, random_single_color_image_from_metadata


def test_image_stitcher(partial_image_metadata_pairs):
    pass


@pytest.mark.parametrize("width", [160])
@pytest.mark.parametrize("height", [90])
@pytest.mark.parametrize("page_width", [int(160 * 1.1)])
@pytest.mark.parametrize("page_height", [int(90 * 1.1)])
def test_partial_image_metadata_pairs(patches_metadata, page_width, page_height):
    pdf = fpdf.FPDF(unit="pt", format=(page_width, page_height))

    for patch in patches_metadata:
        image = random_single_color_image_from_metadata(patch)
        add_image(pdf, ImageMetadataPair(image, patch))

    pdf.output("/tmp/bla.pdf")


class BoxSplitter:
    def __init__(self):
        self.__steps = None
        self.__horiz_keymap = {"dim": Info.WIDTH, "c1": Info.X1, "c2": Info.X2}
        self.__vert_keymap = {"dim": Info.HEIGHT, "c1": Info.Y1, "c2": Info.Y2}

    def split_box(self, box, steps=5):
        self.__steps = steps
        return self.__split_recursively(box, 0)

    def __split_recursively(self, box, step):
        return self.__split_and_recurse(box, step) if self.__steps_left(step) else self.__base_case(box)

    def __steps_left(self, step):
        return step < self.__steps

    @staticmethod
    def __base_case(box):
        return [box]

    def __split_and_recurse(self, box, step):
        new_boxes = self.__random_split(box)
        return chain.from_iterable(self.__tree_recurse(new_boxes, step + 1))

    def __random_split(self, box):
        splitter = random.choice([self.__split_horizontal, self.__split_vertical])
        new_boxes = splitter(box)
        return new_boxes

    def __tree_recurse(self, boxes, step):
        return map(rpartial(self.__split_recursively, step + 1), boxes)

    def __split_horizontal(self, box):
        return self.__split_if_large_enough(box, self.__horiz_keymap)

    def __split_vertical(self, box):
        return self.__split_if_large_enough(box, self.__vert_keymap)

    def __split_if_large_enough(self, box, keymap):
        return self.__get_child_boxes(box, keymap) if self.__large_enough(box, keymap) else self.__base_case(box)

    @staticmethod
    def __large_enough(box, keymap):
        return box[keymap["dim"]] >= 10

    @staticmethod
    def __get_child_boxes(box, keymap):
        dim, c1, c2 = itemgetter("dim", "c1", "c2")(keymap)

        split_len = random.randint(5, box[dim] - 5)
        split_point = box[c1] + split_len

        box_left, box_right = juxt(deepcopy, deepcopy)(box)

        box_left[dim] = split_len
        box_right[dim] = box[dim] - split_len

        box_left[c2] = split_point
        box_right[c1] = split_point

        return box_left, box_right


@pytest.fixture()
def patches_metadata(width, height, page_width, page_height):
    box = get_base_position_metadata(width, height, page_width, page_height)
    box = merge(box, {Info.X1: 0, Info.Y1: 0, Info.X2: width, Info.Y2: height})
    boxes = BoxSplitter().split_box(box)
    return boxes