# -*- coding: utf-8 -*-
"""train.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1nXacyY7r1lbMC9m9aZvuSOLc343bPtrV
"""

import os
from collections import OrderedDict
from torch.autograd import Variable
from models.models import create_model
from PIL import Image
from torchvision import transforms
import util.util as util
import easydict
import torch
import numpy as np
import cv2
from basicsr.archs.rrdbnet_arch import RRDBNet

from realesrgan import RealESRGANer
from realesrgan.archs.srvgg_arch import SRVGGNetCompact

def build_esrgan(
    model_name = 'RealESRGAN_x4plus_anime_6B',
    outscale = 4,
    suffix = 'out',
    tile = 0,
    tile_pad = 10,
    pre_pad = 0,
    face_enhance = False,
    half = False,
    alpha_upsampler = 'realesrgan',
    ext = 'png'

    ):
    """Inference demo for Real-ESRGAN.
    """
    args = easydict.EasyDict({
        'model_name' : model_name,
        'outscale' : outscale,
        'suffix' : suffix,
        'tile' : tile,
        'tile_pad' : tile_pad,
        'pre_pad' : pre_pad,
        'face_enhance' : face_enhance,
        'half' : half,
        'alpha_upsampler' : alpha_upsampler,
        'ext' : ext
    })


    # determine models according to model names
    args.model_name = args.model_name.split('.')[0]
    if args.model_name in ['RealESRGAN_x4plus', 'RealESRNet_x4plus']:  # x4 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
    elif args.model_name in ['RealESRGAN_x4plus_anime_6B']:  # x4 RRDBNet model with 6 blocks
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
        netscale = 4
    elif args.model_name in ['RealESRGAN_x2plus']:  # x2 RRDBNet model
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
        netscale = 2
    elif args.model_name in [
            'RealESRGANv2-anime-xsx2', 'RealESRGANv2-animevideo-xsx2-nousm', 'RealESRGANv2-animevideo-xsx2'
    ]:  # x2 VGG-style model (XS size)
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
        netscale = 2
    elif args.model_name in [
            'RealESRGANv2-anime-xsx4', 'RealESRGANv2-animevideo-xsx4-nousm', 'RealESRGANv2-animevideo-xsx4'
    ]:  # x4 VGG-style model (XS size)
        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
        netscale = 4

    # determine model paths
    model_path = os.path.join('experiments/pretrained_models', args.model_name + '.pth')
    if not os.path.isfile(model_path):
        model_path = os.path.join('realesrgan/weights', args.model_name + '.pth')
    if not os.path.isfile(model_path):
        raise ValueError(f'Model {args.model_name} does not exist.')

    # restorer
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        model=model,
        tile=args.tile,
        tile_pad=args.tile_pad,
        pre_pad=args.pre_pad,
        half=args.half)



    return upsampler

def build_pix2pix():
    opt = easydict.EasyDict({
    'isTrain' : False,
    'name' : 'anime2cheek',
    'gpu_ids' : [],
    'checkpoints_dir' : 'experiments',
    'model' : 'pix2pixHD',
    'norm' : 'instance',
    'use_dropout' : False,
    'data_type' : 32,
    'verbose' : False,
    'fp16' : False,
    'local_rank' : 0,

    'batchSize' : 1,
    'loadSize' : 512,
    'fineSize' : 512,
    'label_nc' : 0,
    'input_nc' : 3,
    'output_nc' : 3,

    'resize_or_crop' : [],
    'serial_batches' : True,
    'no_flip' : True,
    'nThreads' : 1,
    'max_dataset_size' : 50000,

    'display_winsize' : 512,
    'tf_log' : False,

    'netG' : 'global',
    'ngf' : 64,
    'n_downsample_global' : 4,
    'n_blocks_global' : 9,
    'n_blocks_local' : 3,
    'n_local_enhancers' : 1,
    'niter_fix_global' : 0,

    'no_instance' : True,
    'instance_feat' : False,
    'label_feat' : False,
    'feat_num' : 3,
    'load_features' : False,
    'n_downsample_E' : 4,
    'nef' : 16,
    'n_clusters' : 10,

    'initialized' : True,

    'ntest' : float('inf'),
    'aspect_ratio' : 1.0,
    'phase' : 'test',
    'which_epoch' : 'latest',
    'cluster_path' : None,
    'use_encoded_image' : False,
    'export_onnx' : None,
    'engine' : None,
    'onnx' : None,
})
    model = create_model(opt)      # create a model given opt.model and other options
    model.eval()
    return model

def image_preprosses(img, vivid):
    if (img.mode == 'RGBA') or (img.mode == 'P'):
        img.load()
        background = Image.new("RGB", img.size, (255, 255, 255))
        background.paste(img, mask=img.split()[3]) # 3 is the alpha channel
        img = background

    assert (img.mode == 'RGB')
    width, height = img.size

    if not (width == height):
        minsize = min(width, height)
        left = (width - minsize)/2
        top = (height - minsize)/2
        right = (width + minsize)/2
        bottom = (height + minsize)/2
        img = img.crop((left, top, right, bottom))

    assert img.width == img.height

    if (img.width < 400) or (vivid == True):
        img = np.array(img.resize((128,128)))
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    else : img = img.resize((512,512))

    return img

def test_pix2pix(img, pix2pix):
    pretransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    img = pretransform(img)
    img = img.unsqueeze(dim=0)
    with torch.no_grad():
        img = pix2pix.netG(img)
    img = img.data[0].float().numpy()
    img = (np.transpose(img, (1, 2, 0)) + 1) / 2.0 * 255.0
    img = img.astype(np.uint8)
    #img = Image.fromarray(img)

    return img