import argparse

import gradio as gr
import torch
from PIL import Image
from torchvision.transforms import Compose, ConvertImageDtype, Normalize, PILToTensor, Resize
from torchvision.transforms.functional import InterpolationMode

from holocron import models


def main(args):

    model = models.rexnet1_3x(pretrained=True).eval()

    preprocessor = Compose([
        Resize(model.default_cfg['input_shape'][1:], interpolation=InterpolationMode.BILINEAR),
        PILToTensor(),
        ConvertImageDtype(torch.float32),
        Normalize(model.default_cfg['mean'], model.default_cfg['std'])
    ])

    def predict(input):
        input = Image.fromarray(input.astype('uint8'), 'RGB')
        input = preprocessor(input)
        with torch.inference_mode():
            prediction = torch.nn.functional.softmax(model(input.unsqueeze(0))[0], dim=0)
        return {class_name: float(conf) for class_name, conf in zip(model.default_cfg['classes'], prediction)}

    image = gr.inputs.Image()
    outputs = gr.outputs.Label(num_top_classes=3)

    interface = gr.Interface(
        fn=predict,
        inputs=[image],
        outputs=outputs,
        title="Holocron: image classification demo",
        article=("<p style='text-align: center'><a href='https://github.com/frgfm/Holocron'>" "Github Repo</a> | "
                 "<a href='https://frgfm.github.io/Holocron/'>Documentation</a></p>"),
        live=True,
        theme="huggingface",
        layout="horizontal",
    )

    interface.launch(server_port=args.port, show_error=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Holocron image classification demo',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--port", type=int, default=8001, help="Port on which the webserver will be run")
    args = parser.parse_args()

    main(args)