# %%

# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

from networks.mat import Generator
import gradio as gr
import gradio.components as gc
import base64
import glob
import os
import random
import re
from http import HTTPStatus
from io import BytesIO
from typing import Dict, List, NamedTuple, Optional, Tuple

import click
import cv2
import numpy as np
import PIL.Image
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageOps
from pydantic import BaseModel

import dnnlib
import legacy


pyspng = None


def num_range(s: str) -> List[int]:
    '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.'''

    range_re = re.compile(r'^(\d+)-(\d+)$')
    m = range_re.match(s)
    if m:
        return list(range(int(m.group(1)), int(m.group(2))+1))
    vals = s.split(',')
    return [int(x) for x in vals]


def copy_params_and_buffers(src_module, dst_module, require_all=False):
    assert isinstance(src_module, torch.nn.Module)
    assert isinstance(dst_module, torch.nn.Module)
    src_tensors = {name: tensor for name,
                   tensor in named_params_and_buffers(src_module)}
    for name, tensor in named_params_and_buffers(dst_module):
        assert (name in src_tensors) or (not require_all)
        if name in src_tensors:
            tensor.copy_(src_tensors[name].detach()).requires_grad_(
                tensor.requires_grad)


def params_and_buffers(module):
    assert isinstance(module, torch.nn.Module)
    return list(module.parameters()) + list(module.buffers())


def named_params_and_buffers(module):
    assert isinstance(module, torch.nn.Module)
    return list(module.named_parameters()) + list(module.named_buffers())


class Inpainter:
    def __init__(self,
                 network_pkl,
                 resolution=512,
                 truncation_psi=1,
                 noise_mode='const',
                 sdevice='cpu'
                 ):
        self.resolution = resolution
        self.truncation_psi = truncation_psi
        self.noise_mode = noise_mode
        print(f'Loading networks from: {network_pkl}')
        self.device = torch.device(sdevice)
        with dnnlib.util.open_url(network_pkl) as f:
            G_saved = (
                legacy.load_network_pkl(f)
                ['G_ema']
                .to(self.device)
                .eval()
                .requires_grad_(False))  # type: ignore
        net_res = 512 if resolution > 512 else resolution
        self.G = (
            Generator(
                z_dim=512,
                c_dim=0,
                w_dim=512,
                img_resolution=net_res,
                img_channels=3
            )
            .to(self.device)
            .eval()
            .requires_grad_(False)
        )
        copy_params_and_buffers(G_saved,  self.G, require_all=True)

    def generate_images2(
        self,
        dpath: List[PIL.Image.Image],
        mpath: List[Optional[PIL.Image.Image]],
        seed: int = 42,
    ):
        """
        Generate images using pretrained network pickle.
        """
        resolution = self.resolution
        truncation_psi = self.truncation_psi
        noise_mode = self.noise_mode
        # seed = 240  # pick up a random number

        def seed_all(seed):
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            torch.cuda.manual_seed(seed)
        if seed is not None:
            seed_all(seed)

        # no Labels.
        label = torch.zeros([1,  self.G.c_dim], device=self.device)

        def read_image(image):
            image = np.array(image)
            if image.ndim == 2:
                image = image[:, :, np.newaxis]  # HW => HWC
                image = np.repeat(image, 3, axis=2)
            image = image.transpose(2, 0, 1)  # HWC => CHW
            image = image[:3]
            return image
        if resolution != 512:
            noise_mode = 'random'
        results = []
        with torch.no_grad():
            for i, (ipath, m) in enumerate(zip(dpath, mpath)):
                if seed is None:
                    seed_all(i)

                image = read_image(ipath)
                image = (torch.from_numpy(image).float().to(
                    self. device) / 127.5 - 1).unsqueeze(0)

                mask = np.array(m).astype(np.float32) / 255.0
                mask = torch.from_numpy(mask).float().to(
                    self. device).unsqueeze(0).unsqueeze(0)

                z = torch.from_numpy(np.random.randn(
                    1,  self.G.z_dim)).to(self.device)
                output = self.G(image, mask, z, label,
                                truncation_psi=truncation_psi, noise_mode=noise_mode)
                output = (output.permute(0, 2, 3, 1) * 127.5 +
                          127.5).round().clamp(0, 255).to(torch.uint8)
                output = output[0].cpu().numpy()
                results.append(PIL.Image.fromarray(output, 'RGB'))

        return results


# if __name__ == "__main__":
#     generate_images()  # pylint: disable=no-value-for-parameter

# ----------------------------------------------------------------------------
def mask_to_alpha(img, mask):
    img = img.copy()
    img.putalpha(mask)
    return img


def blend(src, target, mask):
    mask = np.expand_dims(mask, axis=-1)
    result = (1-mask) * src + mask * target
    return Image.fromarray(result.astype(np.uint8))


def pad(img, size=(128, 128), tosize=(512, 512), border=1):
    if isinstance(size, float):
        size = (int(img.size[0] * size), int(img.size[1] * size))
    # remove border
    w, h = tosize

    new_img = Image.new('RGBA', (w, h))

    rimg = img.resize(size, resample=Image.Resampling.NEAREST)
    rimg = ImageOps.crop(rimg, border=border)
    tw, th = size
    tw, th = tw - border*2, th - border*2
    tc = ((w-tw)//2, (h-th)//2)

    new_img.paste(rimg, tc)
    mask = Image.new('L', (w, h))
    white = Image.new('L', (tw, th), 255)
    mask.paste(white, tc)

    if 'A' in rimg.getbands():
        mask.paste(rimg.getchannel('A'), tc)
    return new_img, mask


def b64_to_img(b64):
    return Image.open(BytesIO(base64.b64decode(b64)))


def img_to_b64(img):
    with BytesIO() as f:
        img.save(f, format='PNG')
        return base64.b64encode(f.getvalue()).decode('utf-8')


class Predictor:
    def __init__(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.models = {
            "places2": Inpainter(
                network_pkl='models/Places_512_FullData.pkl',
                resolution=512,
                truncation_psi=1.,
                noise_mode='const',
            ),
            "places2+laion300k": Inpainter(
                network_pkl='models/Places_512_FullData+LAION300k.pkl',
                resolution=512,
                truncation_psi=1.,
                noise_mode='const',
            ),
            "places2+laion300k+laion300k(opmasked)": Inpainter(
                network_pkl='models/Places_512_FullData+LAION300k+OPM300k.pkl',
                resolution=512,
                truncation_psi=1.,
                noise_mode='const',
            ),
            "places2+laion300k+laion1200k(opmasked)": Inpainter(
                network_pkl='models/Places_512_FullData+LAION300k+OPM1200k.pkl',
                resolution=512,
                truncation_psi=1.,
                noise_mode='const',
            ),

        }

    # The arguments and types the model takes as input

    def predict(
        self,
        img: Image.Image,
        tosize=(512, 512),
        border=5,
        seed=42,
        size=0.5,
        model='places2',
    ) -> Image:
        i, m = pad(
            img,
            size=size,  # (328, 328),
            tosize=tosize,
            border=border
        )
        """Run a single prediction on the model"""
        imgs = self.models[model].generate_images2(
            dpath=[i.resize((512, 512), resample=Image.Resampling.NEAREST)],
            mpath=[m.resize((512, 512), resample=Image.Resampling.NEAREST)],
            seed=seed,
        )
        img_op_raw = imgs[0].convert('RGBA')
        img_op_raw = img_op_raw.resize(
            tosize, resample=Image.Resampling.NEAREST)
        inpainted = img_op_raw.copy()

        # paste original image to remove inpainting/scaling artifacts
        inpainted = blend(
            i,
            inpainted,
            1-(np.array(m) / 255)
        )
        minpainted = mask_to_alpha(inpainted, m)
        return inpainted, minpainted,  ImageOps.invert(m)

    def predict_tiled(
        self,
        img: Image.Image,
        tosize=(512, 512),
        border=5,
        seed=42,
        size=0.5,
        model='places2',
    ) -> Image:

        i, morig = pad(
            img,
            size=size,  # (328, 328),
            tosize=tosize,
            border=border
        )
        i.putalpha(morig)
        img = i
        img.save('0.png')
        assert img.width == img.height
        assert img.width > 512 and img.width < 512*2

        def tile_coords(image, n=2, tile_size=512):
            assert image.width == image.height
            offsets = np.linspace(0, image.width - tile_size, n).astype(int)
            for i in range(n):
                for j in range(n):
                    left = offsets[j]
                    upper = offsets[i]
                    right = left + tile_size
                    lower = upper + tile_size
                    # tile = image.crop((left, upper, right, lower))
                    yield [left, upper, right, lower]

        for ix, tc in enumerate(tile_coords(img, n=2)):
            i = img.crop(tc)
            i.save(f't{ix}.png')
            m = i.getchannel('A')

            """Run a single prediction on the model"""
            imgs = self.models[model].generate_images2(
                dpath=[i.resize((512, 512), resample=Image.Resampling.NEAREST)],
                mpath=[m.resize((512, 512), resample=Image.Resampling.NEAREST)],
                seed=seed,
            )
            img_op_raw = imgs[0].convert('RGBA')
            # img_op_raw = img_op_raw.resize(tosize, resample=Image.Resampling.NEAREST)
            inpainted = img_op_raw.copy()

            # paste original image to remove inpainting/scaling artifacts
            inpainted = blend(
                i,
                inpainted,
                1-(np.array(m) / 255)
            )
            inpainted.save(f't{ix}_op.png')
            minpainted = mask_to_alpha(inpainted, m)
            # continue with partially inpainted image
            # since the tiles overlap, the next tile will contain (possibly inpainted) parts of the previous tile
            img.paste(inpainted, tc)

        # restore original alpha channel
        img.putalpha(morig)
        return img.convert('RGB'), img,  ImageOps.invert(img.getchannel('A'))
predictor = Predictor()

# %%


def _outpaint(img, tosize, border, seed, size, model, tiled):
    if tiled:
        img_op = predictor.predict_tiled(
            img,
            border=border,
            seed=seed,
            tosize=(tosize, tosize),
            size=float(size),
            model=model,
        )
    else:
        img_op = predictor.predict(
            img,
            border=border,
            seed=seed,
            tosize=(tosize, tosize),
            size=float(size),
            model=model,
        )
    return img_op
# %%


with gr.Blocks() as demo:
    maturl = 'https://github.com/fenglinglwb/MAT'
    gr.Markdown(f'''
        # MAT Primer for Stable Diffusion
        ## based on MAT: Mask-Aware Transformer for Large Hole Image Inpainting
        ### create a primer for use in stable diffusion outpainting

        i have added 2 example scripts to the repo:
        - outpainting_example1.py  using the inpainting pipeline
        - outpainting_example2.py  using the img2img pipeline. this is basically what i used for the examples below
        ''')

    gr.HTML(f'''<a href="{maturl}">{maturl}</a>''')
    with gr.Box():
        with gr.Row():
            gr.Markdown(f"""example with strength 0.5""")
        with gr.Row():
            gr.HTML("<img src='file/hild.gif'> ")
            gr.HTML("<img src='file/process.gif'>")
            gr.HTML("<img src='file/flagscapes.gif'>")
    btn = gr.Button("Run", variant="primary")
    with gr.Row():
        with gr.Column():
            searchimage = gc.Image(label="image", type='pil', image_mode='RGBA')
            to_size = gc.Slider(1, 1920, 512, step=1, label='output size')
            border = gc.Slider(1, 50, 0, step=1, label='border to crop from the image before outpainting')
            seed = gc.Slider(1, 65536, 10, step=1, label='seed')
            size = gc.Slider(0, 1, .5, step=0.01,label='scale of the image before outpainting')
            tiled = gc.Checkbox(label='tiled: run the network with 4 tiles of size 512x512 . only usable if output size >512 and <1024', value=False)

            model = gc.Dropdown(
                choices=['places2',
                         'places2+laion300k',
                         'places2+laion300k+laion300k(opmasked)',
                         'places2+laion300k+laion1200k(opmasked)'],
                value='places2+laion300k+laion1200k(opmasked)',
                label='model',
            )
        with gr.Column():
            outwithoutalpha = gc.Image(label="primed image without alpha channel", type='pil', image_mode='RGBA')
            mask = gc.Image(label="outpainting mask", type='pil')
            out = gc.Image(label="primed image with alpha channel",type='pil', image_mode='RGBA')

    btn.click(
        fn=_outpaint,
        inputs=[searchimage, to_size, border, seed, size, model,tiled],
        outputs=[outwithoutalpha, out,  mask])


# %% launch
demo.launch()