import os
import sys
from pathlib import Path
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from gfpgan import GFPGANer
from realesrgan import RealESRGANer
import numpy as np
import cv2
from PIL import Image
from rembg import remove

# DeOldify
os.system("hub install deoldify==1.2.0")
import paddlehub as hub
hub.server_check()
colorize_model = hub.Module(name='deoldify')

highres_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
bg_upsampler = RealESRGANer(
    scale=4,
    model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth',
    model=highres_model,
    tile=400,
    tile_pad=10,
    pre_pad=0,
    half=True
    )

upsampler =  GFPGANer(
    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
    upscale=4,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=bg_upsampler,
    device="cuda" if torch.cuda.is_available() else "cpu",
    )


os.makedirs("deoldify", exist_ok=True)
os.makedirs("gfpganOutput", exist_ok=True)
os.makedirs("greyscale", exist_ok=True)
os.makedirs("rembg", exist_ok=True)

def restore_image(image):
     _, _, output = upsampler.enhance(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR), has_aligned=False, only_center_face=False, paste_back=True)
     image = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
     return image

def edit_image(image, option):
    tools = ["High Res", "Colorize", "Greyscale", "Remove Background"]
    if option == tools[0]:
        restore_image(image).save("gfpganOutput/output.png")
        return './gfpganOutput/output.png', './gfpganOutput/output.png'
    elif option == tools[1]:
        image.convert("L").save("deoldify/input.png", "PNG", quality=80, optimize=True)
        colorize_model.predict("deoldify/input.png")
        return './output/DeOldify/'+Path('deoldify/input.png').stem+".png", './output/DeOldify/'+Path('deoldify/input.png').stem+".png"
    
    elif option == tools[2]:
        image.convert('L').save("greyscale/output.png")
        return './greyscale/output.png', './greyscale/output.png'
    elif option == tools[3]:
        remove(image).save("rembg/output.png")
        return './rembg/output.png', './rembg/output.png'