'''Pyxelate Demo

- Based on https://huggingface.co/spaces/akhaliq/Pyxelate
- Credits to akhaliq

2021-12-16 first created for testing
'''

import os
from glob import glob
from loguru import logger

import gradio as gr
from skimage import io as sio
from pyxelate import Pyx, Pal


# ----------- Settings -----------
examples = sorted(glob(os.path.join('examples', '*.jpg')))
examples = [[image_file, 5, 5, 'none', 'none'] for image_file in examples]

# ----------- Logging -----------
logger.add('app.log', mode='a')
logger.info('===== APP RESTARTED =====')

# ----------- Params -----------
DOWNSAMPLE_MIN = 1
DOWNSAMPLE_MAX = 10
COLOR_MIN = 1
COLOR_MAX = 20
PALETTE_CHOICES = [
    'none', # if not chosen
    'TELETEXT', 'BBC_MICRO', 'CGA_MODE4_PAL1', 'CGA_MODE5_PAL1', 
    'CGA_MODE4_PAL2', 'ZX_SPECTRUM', 'APPLE_II_LO', 'APPLE_II_HI', 
    'COMMODORE_64', 'GAMEBOY_COMBO_UP', 'GAMEBOY_COMBO_DOWN', 'GAMEBOY_COMBO_LEFT', 
    'GAMEBOY_COMBO_RIGHT', 'GAMEBOY_A_UP', 'GAMEBOY_A_DOWN', 'GAMEBOY_A_LEFT', 
    'GAMEBOY_A_RIGHT', 'GAMEBOY_B_UP', 'GAMEBOY_B_DOWN', 'GAMEBOY_B_LEFT', 
    'GAMEBOY_B_RIGHT', 'GAMEBOY_ORIGINAL', 'GAMEBOY_POCKET', 'GAMEBOY_VIRTUALBOY', 
    'MICROSOFT_WINDOWS_16', 'MICROSOFT_WINDOWS_20', 'MICROSOFT_WINDOWS_PAINT', 
    'PICO_8', 'MSX', 'MONO_OBRADINN_IBM', 'MONO_OBRADINN_MAC', 'MONO_BJG', 'MONO_BW', 
    'MONO_PHOSPHOR_AMBER', 'MONO_PHOSPHOR_LTAMBER', 'MONO_PHOSPHOR_GREEN1', 
    'MONO_PHOSPHOR_GREEN2', 'MONO_PHOSPHOR_GREEN3', 'MONO_PHOSPHOR_APPLE', 'APPLE_II_MONO', 
    'MONO_PHOSPHOR_APPLEC', 'APPLE_II_MONOC'
]
DITHER_CHOICES = ['none', 'naive', 'bayer', 'floyd', 'atkinson']


def predict(image_obj, sampling_param, color_param, palette_param, dither_param):
    img = sio.imread(image_obj.name)
    logger.info('--- image loaded')

    if palette_param != 'none':
        color_param = 'none'
        palette = Pal[palette_param]
    else:
        palette = color_param

    pyx = Pyx(factor=sampling_param, palette=palette, dither=dither_param)
    pyx.fit(img)
    img_out = pyx.transform(img)
    logger.info('--- output generated')
    return img_out

iface = gr.Interface(
    predict,
    title='이미지를 픽셀화 시키는 데모입니다.',
    description='이미지가 주어졌을 때 이미지를 스케일링하고 색감을 바꾸어 픽셀화 시킬 수 있습니다.',
    inputs=[
        gr.inputs.Image(label='인풋 이미지를 준비해주세요', type='file'),
        gr.inputs.Slider(label='다운샘플할 크기를 입력해주세요 (클수록 픽셀크기가 커짐 = 저화질)',
                         minimum=DOWNSAMPLE_MIN, maximum=DOWNSAMPLE_MAX, step=1, default=5),
        gr.inputs.Slider(label='사용할 색깔의 개수를 입력해주세요 (선택 시 아래 팔레트 설정은 "none"으로 해주세요)',
                         minimum=DOWNSAMPLE_MIN, maximum=DOWNSAMPLE_MAX, step=1, default=5),
        gr.inputs.Dropdown(label='색깔 팔레트를 선택하세요 (선택 시 위의 색깔 개수는 무시됩니다)',
                           choices=PALETTE_CHOICES, default='none', type='value'),
        gr.inputs.Dropdown(label='결과 이미지의 부드러움 정도를 나타냅니다 ("dithering")', 
                           choices=DITHER_CHOICES, default='none', type='value'),
    ],
    outputs=[
        gr.outputs.Image(label='결과 이미지 입니다')
    ],
    examples=examples,
    enable_queue=True,
    article='<p style="text-align:center">Credits to <a href="https://github.com/sedthh/pyxelate">GitHub</a></p>',
)

iface.launch(debug=True)