# =====================================================
# Apckeyl_RealESRGAN
# Version 7.0
# tile_processor.py
# =====================================================

from PIL import Image, ImageOps

from config import (
    TILE_SIZE,
    TILE_PAD,
    PRE_PAD,
    MODEL_SCALE,
)


class TileProcessor:

    def __init__(self):

        self.tile = TILE_SIZE
        self.tile_pad = TILE_PAD
        self.pre_pad = PRE_PAD
        self.scale = MODEL_SCALE

    # -------------------------------------------------

    def preprocess(self, image):

        if self.pre_pad > 0:

            image = ImageOps.expand(
                image,
                border=self.pre_pad,
                fill=0
            )

        return image

    # -------------------------------------------------

    def split(self, image):

        image = self.preprocess(image)

        width, height = image.size

        tiles = []

        for top in range(0, height, self.tile):

            for left in range(0, width, self.tile):

                tile_h = min(self.tile, height - top)
                tile_w = min(self.tile, width - left)

                input_left = max(0, left - self.tile_pad)
                input_top = max(0, top - self.tile_pad)

                input_right = min(
                    width,
                    left + tile_w + self.tile_pad
                )

                input_bottom = min(
                    height,
                    top + tile_h + self.tile_pad
                )

                tile = image.crop(
                    (
                        input_left,
                        input_top,
                        input_right,
                        input_bottom
                    )
                )

                tiles.append({

                    "tile": tile,

                    "x": left,
                    "y": top,

                    "w": tile_w,
                    "h": tile_h,

                    "input_left": input_left,
                    "input_top": input_top,
                    "input_right": input_right,
                    "input_bottom": input_bottom,

                })

        return tiles

    # -------------------------------------------------

    def create_output_canvas(
        self,
        image
    ):

        image = self.preprocess(image)

        width, height = image.size

        return Image.new(

            "RGBA",

            (
                width * self.scale,
                height * self.scale
            ),

            (0, 0, 0, 0)

        )

    # -------------------------------------------------

    def paste_tile(

        self,
        canvas,
        tile_info,
        upscaled_tile

    ):

        pad = self.tile_pad * self.scale

        crop_left = 0 if tile_info["input_left"] == 0 else pad
        crop_top = 0 if tile_info["input_top"] == 0 else pad

        crop_right = crop_left + tile_info["w"] * self.scale
        crop_bottom = crop_top + tile_info["h"] * self.scale

        upscaled_tile = upscaled_tile.crop(

            (
                crop_left,
                crop_top,
                crop_right,
                crop_bottom
            )

        )

        canvas.alpha_composite(

            upscaled_tile.convert("RGBA"),

            (
                tile_info["x"] * self.scale,
                tile_info["y"] * self.scale
            )

        )

    # -------------------------------------------------

    def finish(
        self,
        canvas
    ):

        if self.pre_pad > 0:

            pad = self.pre_pad * self.scale

            w, h = canvas.size

            canvas = canvas.crop(

                (
                    pad,
                    pad,
                    w - pad,
                    h - pad
                )

            )

        return canvas.convert("RGB")