import torch
import gradio as gr
from random import randint
from pathlib import Path
from super_image import ImageLoader, EdsrModel, MsrnModel, MdsrModel, AwsrnModel, A2nModel, CarnModel, PanModel, \
    HanModel, DrlnModel, RcanModel

title = "super-image"
description = "State of the Art Image Super-Resolution Models."
article = "<p style='text-align: center'><a href='https://github.com/eugenesiow/super-image'>Github Repo</a>" \
          "| <a href='https://eugenesiow.github.io/super-image/'>Documentation</a> " \
          "| <a href='https://github.com/eugenesiow/super-image#scale-x2'>Models</a></p>"


def inference(img, scale_str, model_name):
    _id = randint(1, 1000)
    output_dir = Path('./tmp/')
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / ('output_image' + str(_id) + '.jpg')
    scale = int(scale_str.replace('x', ''))
    if model_name == 'EDSR':
        model = EdsrModel.from_pretrained('eugenesiow/edsr', scale=scale)
    elif model_name == 'MSRN':
        model = MsrnModel.from_pretrained('eugenesiow/msrn', scale=scale)
    elif model_name == 'MDSR':
        model = MdsrModel.from_pretrained('eugenesiow/mdsr', scale=scale)
    elif model_name == 'AWSRN-BAM':
        model = AwsrnModel.from_pretrained('eugenesiow/awsrn-bam', scale=scale)
    elif model_name == 'A2N':
        model = A2nModel.from_pretrained('eugenesiow/a2n', scale=scale)
    elif model_name == 'CARN':
        model = CarnModel.from_pretrained('eugenesiow/carn', scale=scale)
    elif model_name == 'PAN':
        model = PanModel.from_pretrained('eugenesiow/pan', scale=scale)
    elif model_name == 'HAN':
        model = HanModel.from_pretrained('eugenesiow/han', scale=scale)
    elif model_name == 'DRLN':
        model = DrlnModel.from_pretrained('eugenesiow/drln', scale=scale)
    elif model_name == 'RCAN':
        model = RcanModel.from_pretrained('eugenesiow/rcan', scale=scale)
    else:
        model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=scale)
    inputs = ImageLoader.load_image(img)
    preds = model(inputs)
    output_file_str = str(output_file.resolve())
    ImageLoader.save_image(preds, output_file_str)
    return output_file_str


torch.hub.download_url_to_file('http://people.rennes.inria.fr/Aline.Roumy/results/images_SR_BMVC12/input_groundtruth/baby_mini_d3_gaussian.bmp',
                               'baby.bmp')
torch.hub.download_url_to_file('http://people.rennes.inria.fr/Aline.Roumy/results/images_SR_BMVC12/input_groundtruth/woman_mini_d3_gaussian.bmp',
                               'woman.bmp')

gr.Interface(
    inference,
    [
        gr.inputs.Image(type="pil", label="Input"),
        gr.inputs.Radio(["x2", "x3", "x4"], label='scale'),
        gr.inputs.Dropdown(choices=['EDSR-base', 'EDSR', 'MSRN', 'MDSR', 'AWSRN-BAM', 'A2N', 'CARN', 'PAN', 'HAN',
                                    'DRLN', 'RCAN'],
                           label='Model')
    ],
    gr.outputs.Image(type="file", label="Output"),
    title=title,
    description=description,
    article=article,
    examples=[
        ['baby.bmp', 'x2', 'EDSR-base'],
        ['woman.bmp', 'x3', 'MSRN']
    ],
    enable_queue=True
    ).launch(debug=True)