import torch
import numpy as np

from PIL import Image, ImageFilter
from modules.util import resample_image, set_image_shape_ceil, get_image_shape_ceil
from modules.upscaler import perform_upscale
import cv2


inpaint_head_model = None


class InpaintHead(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.head = torch.nn.Parameter(torch.empty(size=(320, 5, 3, 3), device='cpu'))

    def __call__(self, x):
        x = torch.nn.functional.pad(x, (1, 1, 1, 1), "replicate")
        return torch.nn.functional.conv2d(input=x, weight=self.head)


current_task = None


def box_blur(x, k):
    x = Image.fromarray(x)
    x = x.filter(ImageFilter.BoxBlur(k))
    return np.array(x)


def max_filter_opencv(x, ksize=3):
    # Use OpenCV maximum filter
    # Make sure the input type is int16
    return cv2.dilate(x, np.ones((ksize, ksize), dtype=np.int16))


def morphological_open(x):
    # Convert array to int16 type via threshold operation
    x_int16 = np.zeros_like(x, dtype=np.int16)
    x_int16[x > 127] = 256

    for i in range(32):
        # Use int16 type to avoid overflow
        maxed = max_filter_opencv(x_int16, ksize=3) - 8
        x_int16 = np.maximum(maxed, x_int16)

    # Clip negative values to 0 and convert back to uint8 type
    x_uint8 = np.clip(x_int16, 0, 255).astype(np.uint8)
    return x_uint8


def up255(x, t=0):
    y = np.zeros_like(x).astype(np.uint8)
    y[x > t] = 255
    return y


def imsave(x, path):
    x = Image.fromarray(x)
    x.save(path)


def regulate_abcd(x, a, b, c, d):
    H, W = x.shape[:2]
    if a < 0:
        a = 0
    if a > H:
        a = H
    if b < 0:
        b = 0
    if b > H:
        b = H
    if c < 0:
        c = 0
    if c > W:
        c = W
    if d < 0:
        d = 0
    if d > W:
        d = W
    return int(a), int(b), int(c), int(d)


def compute_initial_abcd(x):
    indices = np.where(x)
    a = np.min(indices[0])
    b = np.max(indices[0])
    c = np.min(indices[1])
    d = np.max(indices[1])
    abp = (b + a) // 2
    abm = (b - a) // 2
    cdp = (d + c) // 2
    cdm = (d - c) // 2
    l = int(max(abm, cdm) * 1.15)
    a = abp - l
    b = abp + l + 1
    c = cdp - l
    d = cdp + l + 1
    a, b, c, d = regulate_abcd(x, a, b, c, d)
    return a, b, c, d


def solve_abcd(x, a, b, c, d, k):
    k = float(k)
    assert 0.0 <= k <= 1.0

    H, W = x.shape[:2]
    if k == 1.0:
        return 0, H, 0, W
    while True:
        if b - a >= H * k and d - c >= W * k:
            break

        add_h = (b - a) < (d - c)
        add_w = not add_h

        if b - a == H:
            add_w = True

        if d - c == W:
            add_h = True

        if add_h:
            a -= 1
            b += 1

        if add_w:
            c -= 1
            d += 1

        a, b, c, d = regulate_abcd(x, a, b, c, d)
    return a, b, c, d


def fooocus_fill(image, mask):
    current_image = image.copy()
    raw_image = image.copy()
    area = np.where(mask < 127)
    store = raw_image[area]

    for k, repeats in [(512, 2), (256, 2), (128, 4), (64, 4), (33, 8), (15, 8), (5, 16), (3, 16)]:
        for _ in range(repeats):
            current_image = box_blur(current_image, k)
            current_image[area] = store

    return current_image


class InpaintWorker:
    def __init__(self, image, mask, use_fill=True, k=0.618):
        a, b, c, d = compute_initial_abcd(mask > 0)
        a, b, c, d = solve_abcd(mask, a, b, c, d, k=k)

        # interested area
        self.interested_area = (a, b, c, d)
        self.interested_mask = mask[a:b, c:d]
        self.interested_image = image[a:b, c:d]

        # super resolution
        if get_image_shape_ceil(self.interested_image) < 1024:
            self.interested_image = perform_upscale(self.interested_image)

        # resize to make images ready for diffusion
        self.interested_image = set_image_shape_ceil(self.interested_image, 1024)
        self.interested_fill = self.interested_image.copy()
        H, W, C = self.interested_image.shape

        # process mask
        self.interested_mask = up255(resample_image(self.interested_mask, W, H), t=127)

        # compute filling
        if use_fill:
            self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask)

        # soft pixels
        self.mask = morphological_open(mask)
        self.image = image

        # ending
        self.latent = None
        self.latent_after_swap = None
        self.swapped = False
        self.latent_mask = None
        self.inpaint_head_feature = None
        return

    def load_latent(self, latent_fill, latent_mask, latent_swap=None):
        self.latent = latent_fill
        self.latent_mask = latent_mask
        self.latent_after_swap = latent_swap
        return

    def patch(self, inpaint_head_model_path, inpaint_latent, inpaint_latent_mask, model):
        global inpaint_head_model

        if inpaint_head_model is None:
            inpaint_head_model = InpaintHead()
            sd = torch.load(inpaint_head_model_path, map_location='cpu')
            inpaint_head_model.load_state_dict(sd)

        feed = torch.cat([
            inpaint_latent_mask,
            model.model.process_latent_in(inpaint_latent)
        ], dim=1)

        inpaint_head_model.to(device=feed.device, dtype=feed.dtype)
        inpaint_head_feature = inpaint_head_model(feed)

        def input_block_patch(h, transformer_options):
            if transformer_options["block"][1] == 0:
                h = h + inpaint_head_feature.to(h)
            return h

        m = model.clone()
        m.set_model_input_block_patch(input_block_patch)
        return m

    def swap(self):
        if self.swapped:
            return

        if self.latent is None:
            return

        if self.latent_after_swap is None:
            return

        self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
        self.swapped = True
        return

    def unswap(self):
        if not self.swapped:
            return

        if self.latent is None:
            return

        if self.latent_after_swap is None:
            return

        self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
        self.swapped = False
        return

    def color_correction(self, img):
        fg = img.astype(np.float32)
        bg = self.image.copy().astype(np.float32)
        w = self.mask[:, :, None].astype(np.float32) / 255.0
        y = fg * w + bg * (1 - w)
        return y.clip(0, 255).astype(np.uint8)

    def post_process(self, img):
        a, b, c, d = self.interested_area
        content = resample_image(img, d - c, b - a)
        result = self.image.copy()
        result[a:b, c:d] = content
        result = self.color_correction(result)
        return result

    def visualize_mask_processing(self):
        return [self.interested_fill, self.interested_mask, self.interested_image]