from .border import image_border
from .compose import place_at
from .filters import torgb
from contracts import contract
from procgraph import Block, BadConfig
from procgraph.block_utils import input_check_convertible_to_rgb
from . import np

__all__ = ["ImageGrid", "make_images_grid"]


class ImageGrid(Block):
    """
        A block that creates a larger image by arranging them in a grid.

        The output is rgb, uint8.

        Inputs are passed through the "torgb" function.

    """

    Block.alias("grid")

    Block.config("cols", "Columns in the grid.", default=None)
    Block.config("bgcolor", "Background color.", default=[0, 0, 0])

    Block.config("pad", "Padding for each cell", default=0)
    Block.input_is_variable("Images to arrange in a grid.", min=1)
    Block.output("grid", "Images arranged in a grid.")

    def update(self):
        if not self.all_input_signals_ready():
            return

        n = self.num_input_signals()
        for i in range(n):
            input_check_convertible_to_rgb(self, i)

        cols = self.config.cols

        if cols is not None and not isinstance(cols, int):
            raise BadConfig("Expected an integer.", self, "cols")

        images = [self.get_input(i) for i in range(n)]

        images = list(map(torgb, images))
        canvas = make_images_grid(
            images, cols=self.config.cols, pad=self.config.pad, bgcolor=self.config.bgcolor
        )

        self.set_output(0, canvas)


@contract(images="list[>=1](array)")
def make_images_grid(images, cols=None, pad=0, bgcolor=[1, 1, 1]):
    n = len(images)
    if cols is None:
        cols = int(np.ceil(np.sqrt(n)))

    rows = int(np.ceil(n * 1.0 / cols))

    assert cols > 0 and rows > 0
    assert n <= cols * rows

    # find width and height for the grid
    col_width = np.zeros(cols, dtype="int32")
    row_height = np.zeros(rows, dtype="int32")
    for i in range(n):
        image = images[i]
        col = i % cols
        row = int((i - i % cols) / cols)
        assert 0 <= col < cols
        assert 0 <= row < rows

        if pad > 0:
            image = image_border(image, left=pad, right=pad, top=pad, bottom=pad, color=bgcolor)
        width = image.shape[1]
        height = image.shape[0]

        col_width[col] = max(width, col_width[col])
        row_height[row] = max(height, row_height[row])

    canvas_width = sum(col_width)
    canvas_height = sum(row_height)

    # find position for each col and row
    col_x = np.zeros(cols, dtype="int32")
    for col in range(1, cols):
        col_x[col] = col_x[col - 1] + col_width[col - 1]

    assert canvas_width == col_x[-1] + col_width[-1]

    row_y = np.zeros(rows, dtype="int32")
    for row in range(1, rows):
        row_y[row] = row_y[row - 1] + row_height[row - 1]
    assert canvas_height == row_y[-1] + row_height[-1]

    canvas = np.zeros((canvas_height, canvas_width, 3), dtype="uint8")
    for k in range(3):
        canvas[:, :, k] = bgcolor[k] * 255

    for i in range(n):
        col = i % cols
        row = int((i - i % cols) / cols)
        assert 0 <= col < cols
        assert 0 <= row < rows
        image = images[i]
        x = col_x[col]
        y = row_y[row]

        # Pad if not right shape
        extra_hor = col_width[col] - image.shape[1]
        extra_ver = row_height[row] - image.shape[0]
        eleft = int(extra_hor / 2)
        eright = extra_hor - eleft
        etop = int(extra_ver / 2)
        ebottom = extra_ver - etop
        image = image_border(image, left=eleft, right=eright, top=etop, bottom=ebottom, color=bgcolor)

        # TODO: align here
        place_at(canvas, image, x, y)

    return canvas
