#!/usr/bin/env python

from __future__ import annotations

import os
import random

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageOps
import uuid

import argparse
import cv2
import glob
import os
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url

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

DESCRIPTION = '''<center><h1>☝️ Bigmojis ☝️</h1></span>  
<span font-size:16px;">An emoji upscaler, for when you <i>really</i> mean it</span>  
</center>

Space by [ZachNagengast](https://huggingface.co/ZachNagengast) 

[Follow me on Twitter!](https://twitter.com/ZachNagengast)

Upscaler models provided by the [BasicSR](https://github.com/XPixelGroup/BasicSR) and [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) python packages.

Emojis by [Apple](https://www.apple.com).
'''

upsampler = None
netscale = 4

# load model
def load_model(model_type='RealESRGAN_x4plus'):    
    if model_type == 'RealESRGAN_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
        file_url = f'models/{model_type}.pth'
    elif model_type == '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
        file_url = f'models/{model_type}.pth'
    elif model_type == '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
        file_url = f'models/{model_type}.pth'
    elif model_type == '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
        file_url = f'models/{model_type}.pth'
    elif model_type == 'realesr-animevideov3':  # 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
        file_url = f'models/{model_type}.pth'

    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    model_path = os.path.join(ROOT_DIR, file_url)

    upsampler = RealESRGANer(
                    scale=netscale,
                    model_path=model_path,
                    model=model,
                    
                )
    
    return upsampler

def switch_model(model_type):
    global upsampler
    upsampler = load_model(model_type)

def save_image(img):
    unique_name = str(uuid.uuid4()) + '.png'
    img.save(unique_name)
    return unique_name

def generate_emoji_image(
        prompt: str,
        font_type: str = "Apple",
        background: str = "white",
        size: int = 160,
        padding: int = 0,
):
    font_name = "AppleColorEmoji.ttc"
    if font_type == "Google":
        font_name = "NotoColorEmoji.ttf"
    elif font_type == "Twitter":
        font_name = "TwitterColorEmoji.ttf"
    
    font = ImageFont.truetype(
        font_name, size=160
    )

    if background == "transparent":
        background = (0, 0, 0, 0)

    im = Image.new("RGBA", (size, size), background)
    d = ImageDraw.Draw(im)

    d.text((padding, padding), prompt, fill='white', embedded_color=True, font=font)
    return im

def generate_preview(
    prompt: str,
    font_type: str = "Apple",
    background: str = "white",
):
    im = generate_emoji_image(prompt, font_type, background, 180, 10)
    return im


def upscale_image(img):
    print(f"Upscaling...")
    cv2_im = np.array(img)
    output, _ = upsampler.enhance(cv2_im, outscale=netscale)
    return Image.fromarray(output)

def generate_upscaled_emoji(
    prompt: str,
    font_type: str = "Apple",
    background: str = "white",
    model: str = "RealESRGAN_x4plus",
):
    padding = 10 # prevents border artifacts
    im = generate_emoji_image(prompt, font_type, background, 160+padding*2, padding)

    result = upscale_image(im)

    # crop padding
    cropped = ImageOps.crop(result, border=padding)

    return cropped



examples = ['🤗', '😂', '🦙', '👍', '💸', '✨', '🚀', '🫱🏼‍🫲🏾', '🎉', '😎', '🛸', '🍩', '🦜', '🗿', '🧌', '🦋', '🆙']

with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            with gr.Group():
                with gr.Row():
                    prompt = gr.Text(
                        label="Emoji",
                        elem_id="prompt-text",
                        show_label=True,
                        max_lines=1,
                        placeholder="Enter your emoji",
                        container=False,
                    )
                    run_button = gr.Button("Upscale", scale=0)
            with gr.Row():
                preview = gr.Image(label="Glyph Bitmap", show_label=True, scale=1)
                background = gr.Radio(choices=["transparent", "white", "black", "red", "green", "blue"], value="transparent", label="Background Color")
                font_type = gr.Radio(choices=["Apple", "Google", "Twitter"], value="Apple", label="Type", visible=False)
            with gr.Row():
                gr.Examples(
                    examples=examples,
                    inputs=prompt,
                    outputs=[preview],
                    fn=generate_emoji_image,
                    cache_examples=False,
                )
            with gr.Row():
                model = gr.Radio(choices=['RealESRNet_x4plus','RealESRGAN_x4plus','RealESRGAN_x4plus_anime_6B','RealESRGAN_x2plus', 'realesr-animevideov3'], value="RealESRNet_x4plus", label="Model")
        with gr.Column():
            upscaled = gr.Image(label="Upscaled", value="huggingface-big.png")
            gr.DuplicateButton(
                value="Duplicate Space for private use",
                elem_id="duplicate-button",
                visible=True,
            )

    prompt.change(generate_preview, inputs=[prompt, font_type, background], outputs=[preview], api_name="run")
    background.change(generate_preview, inputs=[prompt, font_type, background], outputs=[preview], api_name="run")
    font_type.change(generate_preview, inputs=[prompt, font_type, background], outputs=[preview], api_name="run")
    model.change(switch_model, inputs=[model], outputs=[])
    prompt.submit(generate_upscaled_emoji, inputs=[prompt, font_type, background, model], outputs=[upscaled], api_name="run")
    run_button.click(generate_upscaled_emoji, inputs=[prompt, font_type, background, model], outputs=[upscaled], api_name="run")

if __name__ == "__main__":
    upsampler = load_model('RealESRNet_x4plus')

    demo.queue(max_size=20).launch()