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

Automatically generated by Colaboratory.

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

import os
from data import create_dataset
from models import create_model
from util.visualizer import save_images
from PIL import Image
from torchvision import transforms
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(model_epoch):
    opt = easydict.EasyDict({
        'isTrain' : False,
        'use_wandb' : False,
        'gpu_ids' : [],
        'checkpoints_dir' : 'experiments',
        'batch_size' : 1,
        
        
        'model' : 'pix2pix',
        'input_nc' : 3,
        'output_nc' : 3,
        'ngf' : 64,
        'ndf' : 64,
        'netD' : 'basic',
        'netG' : 'unet_256',
        'n_layers_D' : 3,
        'norm' : 'batch',
        'init_type' : 'normal',
        'init_gain': 0.02,
        'no_dropout' : False,

        'direction' : 'AtoB',
        'serial_batches' : True,
        'num_threads' : 0,
        'load_size' : 512,
        'crop_size' : 256,
        'max_dataset_size' : 50000,
        'preprocess' : [],
        'no_flip' : True,
        'display_winsize' : 512,
        'verbose' : False,
        'suffix' : '',
        'load_iter' : 0,

        #test_arguments
        'aspect_ratio' : 1.0,
        'phase' : 'test',
        'eval' : True,
        'num_test' : 1,
        'model' : 'test',
        'load_size' : 512,
        'dataset_mode' : 'single',
        'model_suffix' : '',
        'epoch' : 110,   #latest
        'name' : 'pix2pix',
         })
    opt.epoch = model_epoch
    model = create_model(opt)      # create a model given opt.model and other options
    model.setup(opt)               # regular setup: load and print networks; create schedulers
    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