import os
import cv2
import numpy as np
import gradio as gr
from inference import Predictor
from utils.image_processing import resize_image

os.makedirs('output', exist_ok=True)


def inference(
    image: np.ndarray,
    style,
    imgsz=None,
):
    retain_color = False

    weight = {
        "AnimeGAN_Hayao": "hayao",
        "AnimeGAN_Shinkai": "shinkai",
        "AnimeGANv2_Hayao": "hayao:v2",
        "AnimeGANv2_Shinkai": "shinkai:v2",
        "AnimeGANv2_Arcane": "arcane:v2",
    }[style]
    predictor = Predictor(
        weight,
        device='cpu',
        retain_color=retain_color,
        imgsz=imgsz,
    )

    save_path = f"output/out.jpg"    
    image = resize_image(image, width=imgsz)
    anime_image = predictor.transform(image)[0]
    cv2.imwrite(save_path, anime_image[..., ::-1])
    return anime_image, save_path


title = "AnimeGANv2: To produce your own animation."
description = r"""Turn your photo into anime style 😊"""
article = r"""
[![GitHub Stars](https://img.shields.io/github/stars/ptran1203/pytorch-animeGAN?style=social)](https://github.com/ptran1203/pytorch-animeGAN)    
### 🗻 Demo

"""

gr.Interface(
    fn=inference,
    inputs=[
        gr.components.Image(label="Input"),
        gr.Dropdown(
            [
                'AnimeGAN_Hayao',
                'AnimeGAN_Shinkai',
                'AnimeGANv2_Hayao',
                'AnimeGANv2_Shinkai',
                'AnimeGANv2_Arcane',
            ],
            type="value",
            value='AnimeGANv2_Hayao',
            label='Style'
        ),
        gr.Dropdown(
            [
                None,
                416,
                512,
                768,
                1024,
                1536,
            ],
            type="value",
            value=None,
            label='Image size'
        )
    ],
    outputs=[
        gr.components.Image(type="numpy", label="Output (The whole image)"),
        gr.components.File(label="Download the output image")
    ],
    title=title,
    description=description,
    article=article,
    allow_flagging="never",
    examples=[
        ['example/arcane/girl4.jpg', 'AnimeGANv2_Arcane', "Yes"],
        ['example/arcane/leo.jpg', 'AnimeGANv2_Arcane', "Yes"],
        ['example/arcane/girl.jpg', 'AnimeGANv2_Arcane', "Yes"],
        ['example/arcane/anne.jpg', 'AnimeGANv2_Arcane', "Yes"],
        # ['example/boy2.jpg', 'AnimeGANv3_Arcane', "No"],
        # ['example/cap.jpg', 'AnimeGANv3_Arcane', "No"],
        ['example/more/hayao_v2/pexels-camilacarneiro-6318793.jpg', 'AnimeGANv2_Hayao', "Yes"],
        ['example/more/hayao_v2/pexels-nandhukumar-450441.jpg', 'AnimeGANv2_Hayao', "Yes"],
    ]
).launch()