import copy
import glob
import os
from multiprocessing.dummy import Pool as ThreadPool

from PIL import Image
from torchvision.transforms.functional import to_tensor

from ..Models import *


class ImageSplitter:
    # key points:
    # Boarder padding and over-lapping img splitting to avoid the instability of edge value
    # Thanks Waifu2x's autorh nagadomi for suggestions (https://github.com/nagadomi/waifu2x/issues/238)

    def __init__(self, seg_size=48, scale_factor=2, boarder_pad_size=3):
        self.seg_size = seg_size
        self.scale_factor = scale_factor
        self.pad_size = boarder_pad_size
        self.height = 0
        self.width = 0
        self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="bilinear")

    def split_img_tensor(self, pil_img, scale_method=Image.BILINEAR, img_pad=0):
        # resize image and convert them into tensor
        img_tensor = to_tensor(pil_img).unsqueeze(0)
        img_tensor = nn.ReplicationPad2d(self.pad_size)(img_tensor)
        batch, channel, height, width = img_tensor.size()
        self.height = height
        self.width = width

        if scale_method is not None:
            img_up = pil_img.resize(
                (2 * pil_img.size[0], 2 * pil_img.size[1]), scale_method
            )
            img_up = to_tensor(img_up).unsqueeze(0)
            img_up = nn.ReplicationPad2d(self.pad_size * self.scale_factor)(img_up)

        patch_box = []
        # avoid the residual part is smaller than the padded size
        if (
            height % self.seg_size < self.pad_size
            or width % self.seg_size < self.pad_size
        ):
            self.seg_size += self.scale_factor * self.pad_size

        # split image into over-lapping pieces
        for i in range(self.pad_size, height, self.seg_size):
            for j in range(self.pad_size, width, self.seg_size):
                part = img_tensor[
                    :,
                    :,
                    (i - self.pad_size) : min(
                        i + self.pad_size + self.seg_size, height
                    ),
                    (j - self.pad_size) : min(j + self.pad_size + self.seg_size, width),
                ]
                if img_pad > 0:
                    part = nn.ZeroPad2d(img_pad)(part)
                if scale_method is not None:
                    # part_up = self.upsampler(part)
                    part_up = img_up[
                        :,
                        :,
                        self.scale_factor
                        * (i - self.pad_size) : min(
                            i + self.pad_size + self.seg_size, height
                        )
                        * self.scale_factor,
                        self.scale_factor
                        * (j - self.pad_size) : min(
                            j + self.pad_size + self.seg_size, width
                        )
                        * self.scale_factor,
                    ]

                    patch_box.append((part, part_up))
                else:
                    patch_box.append(part)
        return patch_box

    def merge_img_tensor(self, list_img_tensor):
        out = torch.zeros(
            (1, 3, self.height * self.scale_factor, self.width * self.scale_factor)
        )
        img_tensors = copy.copy(list_img_tensor)
        rem = self.pad_size * 2

        pad_size = self.scale_factor * self.pad_size
        seg_size = self.scale_factor * self.seg_size
        height = self.scale_factor * self.height
        width = self.scale_factor * self.width
        for i in range(pad_size, height, seg_size):
            for j in range(pad_size, width, seg_size):
                part = img_tensors.pop(0)
                part = part[:, :, rem:-rem, rem:-rem]
                # might have error
                if len(part.size()) > 3:
                    _, _, p_h, p_w = part.size()
                    out[:, :, i : i + p_h, j : j + p_w] = part
                # out[:,:,
                # self.scale_factor*i:self.scale_factor*i+p_h,
                # self.scale_factor*j:self.scale_factor*j+p_w] = part
        out = out[:, :, rem:-rem, rem:-rem]
        return out


def load_single_image(
    img_file,
    up_scale=False,
    up_scale_factor=2,
    up_scale_method=Image.BILINEAR,
    zero_padding=False,
):
    img = Image.open(img_file).convert("RGB")
    out = to_tensor(img).unsqueeze(0)
    if zero_padding:
        out = nn.ZeroPad2d(zero_padding)(out)
    if up_scale:
        size = tuple(map(lambda x: x * up_scale_factor, img.size))
        img_up = img.resize(size, up_scale_method)
        img_up = to_tensor(img_up).unsqueeze(0)
        out = (out, img_up)

    return out


def standardize_img_format(img_folder):
    def process(img_file):
        img_path = os.path.dirname(img_file)
        img_name, _ = os.path.basename(img_file).split(".")
        out = os.path.join(img_path, img_name + ".JPEG")
        os.rename(img_file, out)

    list_imgs = []
    for i in ["png", "jpeg", "jpg"]:
        list_imgs.extend(glob.glob(img_folder + "**/*." + i, recursive=True))
    print("Found {} images.".format(len(list_imgs)))
    pool = ThreadPool(4)
    pool.map(process, list_imgs)
    pool.close()
    pool.join()