import math
import warnings

from PIL import Image, ImageFilter


class Outcrop(object):
    def __init__(
        self,
        image,
        generate,  # current generate object
    ):
        self.image = image
        self.generate = generate

    def process(
        self,
        extents: dict,
        opt,  # current options
        orig_opt,  # ones originally used to generate the image
        image_callback=None,
        prefix=None,
    ):
        # grow and mask the image
        extended_image = self._extend_all(extents)

        # switch samplers temporarily
        curr_sampler = self.generate.sampler
        self.generate.sampler_name = opt.sampler_name
        self.generate._set_sampler()

        def wrapped_callback(img, seed, **kwargs):
            preferred_seed = (
                orig_opt.seed
                if orig_opt.seed is not None and orig_opt.seed >= 0
                else seed
            )
            image_callback(img, preferred_seed, use_prefix=prefix, **kwargs)

        result = self.generate.prompt2image(
            opt.prompt,
            seed=opt.seed or orig_opt.seed,
            sampler=self.generate.sampler,
            steps=opt.steps,
            cfg_scale=opt.cfg_scale,
            ddim_eta=self.generate.ddim_eta,
            width=extended_image.width,
            height=extended_image.height,
            init_img=extended_image,
            strength=0.90,
            image_callback=wrapped_callback if image_callback else None,
            seam_size=opt.seam_size or 96,
            seam_blur=opt.seam_blur or 16,
            seam_strength=opt.seam_strength or 0.7,
            seam_steps=20,
            tile_size=32,
            color_match=True,
            force_outpaint=True,  # this just stops the warning about erased regions
        )

        # swap sampler back
        self.generate.sampler = curr_sampler
        return result

    def _extend_all(
        self,
        extents: dict,
    ) -> Image:
        """
        Extend the image in direction ('top','bottom','left','right') by
        the indicated value. The image canvas is extended, and the empty
        rectangular section will be filled with a blurred copy of the
        adjacent image.
        """
        image = self.image
        for direction in extents:
            assert direction in [
                "top",
                "left",
                "bottom",
                "right",
            ], 'Direction must be one of "top", "left", "bottom", "right"'
            pixels = extents[direction]
            # round pixels up to the nearest 64
            pixels = math.ceil(pixels / 64) * 64
            print(f">> extending image {direction}ward by {pixels} pixels")
            image = self._rotate(image, direction)
            image = self._extend(image, pixels)
            image = self._rotate(image, direction, reverse=True)
        return image

    def _rotate(self, image: Image, direction: str, reverse=False) -> Image:
        """
        Rotates image so that the area to extend is always at the top top.
        Simplifies logic later. The reverse argument, if true, will undo the
        previous transpose.
        """
        transposes = {
            "right": ["ROTATE_90", "ROTATE_270"],
            "bottom": ["ROTATE_180", "ROTATE_180"],
            "left": ["ROTATE_270", "ROTATE_90"],
        }
        if direction not in transposes:
            return image
        transpose = transposes[direction][1 if reverse else 0]
        return image.transpose(Image.Transpose.__dict__[transpose])

    def _extend(self, image: Image, pixels: int) -> Image:
        extended_img = Image.new("RGBA", (image.width, image.height + pixels))

        extended_img.paste((0, 0, 0), [0, 0, image.width, image.height + pixels])
        extended_img.paste(image, box=(0, pixels))

        # now make the top part transparent to use as a mask
        alpha = extended_img.getchannel("A")
        alpha.paste(0, (0, 0, extended_img.width, pixels))
        extended_img.putalpha(alpha)

        return extended_img