'''Artist Classifier

prototype

---
- 2022-01-18 jkang first created
'''
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import seaborn as sns

import io
import json
import numpy as np
import skimage
import skimage.io
from skimage.transform import resize
from loguru import logger
from huggingface_hub import from_pretrained_keras
import gradio as gr
import tensorflow as tf
tfk = tf.keras

from gradcam_utils import get_img_4d_array, make_gradcam_heatmap, align_image_with_heatmap

# ---------- Settings ----------
ARTIST_META = 'artist.json'
TREND_META = 'trend.json'
EXAMPLES = ['monet.jpg']
ALPHA = 0.9

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

# ---------- Model ----------
logger.info('loading models...')
artist_model = from_pretrained_keras("jkang/drawing-artist-classifier")
trend_model = from_pretrained_keras("jkang/drawing-artistic-trend-classifier")
logger.info('both models loaded')

def load_json_as_dict(json_file):
    with open(json_file, 'r') as f:
        out = json.load(f)
    return dict(out)

def load_image_as_array(image_file):
    img = skimage.io.imread(image_file, as_gray=False, plugin='matplotlib')
    if (img.shape[-1] > 3):  # if RGBA
        img = img[..., :-1]
    return img

def resize_image(img_array, width, height):
    img_resized = resize(img_array, (height, width), 
                         anti_aliasing=True, 
                         preserve_range=False)
    return skimage.img_as_ubyte(img_resized)

def predict(input_image):
    img_3d_array = load_image_as_array(input_image)
    img_3d_array = resize_image(img_3d_array)
    img_4d_array = img_3d_array[np.newaxis,...]
    logger.info(f'--- {input_image} loaded')
    
    artist2id = load_json_as_dict(ARTIST_META)
    trend2id = load_json_as_dict(TREND_META)
    id2artist = {artist2id[artist]:artist for artist in artist2id}
    id2trend = {trend2id[trend]:trend for trend in trend2id}
    
    # Artist model
    a_heatmap, a_pred_id, a_pred_out = make_gradcam_heatmap(artist_model,
                                                            img_4d_array,
                                                            pred_idx=None)
    a_img_pil = align_image_with_heatmap(
        img_4d_array, a_heatmap, alpha=ALPHA, cmap='jet')
    a_img = np.asarray(a_img_pil).astype('float32')/255
    a_label = id2artist[a_pred_id]
    a_prob = a_pred_out[a_pred_id]

    # Trend model
    t_heatmap, t_pred_id, t_pred_out = make_gradcam_heatmap(trend_model,
                                                            img_4d_array,
                                                            pred_idx=None)
    
    t_img_pil = align_image_with_heatmap(
        img_4d_array, t_heatmap, alpha=ALPHA, cmap='jet')
    t_img = np.asarray(t_img_pil).astype('float32')/255
    t_label = id2trend[t_pred_id]
    t_prob = t_pred_out[t_pred_id]
    
    with sns.plotting_context('poster', font_scale=0.7):
        fig, (ax1, ax2, ax3) = plt.subplots(
            1, 3, figsize=(12, 6), facecolor='white')
        for ax in (ax1, ax2, ax3):
            ax.set_xticks([])
            ax.set_yticks([])

        ax1.imshow(img_3d_array)
        ax2.imshow(a_img)
        ax3.imshow(t_img)

        ax1.set_title(f'Input Image', ha='left', x=0, y=1.05)
        ax2.set_title(f'Artist Prediction:\n => {a_label} ({a_prob:.2f})', ha='left', x=0, y=1.05)
        ax3.set_title(f'Style Prediction:\n => {t_label} ({t_prob:.2f})', ha='left', x=0, y=1.05)
        fig.tight_layout()
        
    buf = io.BytesIO()
    fig.savefig(buf, bbox_inces='tight', fotmat='jpg')
    buf.seek(0)
    pil_img = Image.open(buf)
    plt.close()
    logger.info('--- image generated')

    a_labels = {id2artist[i]: float(pred) for i, pred in enumerate(a_pred_out)}
    t_labels = {id2trend[i]: float(pred) for i, pred in enumerate(t_pred_out)}
    return a_labels, t_labels, pil_img

iface = gr.Interface(
    predict,
    title='Predict Artist and Artistic Style of Drawings 🎨👨🏻‍🎨 (prototype)',
    description='Upload a drawing/image and the model will predict how likely it seems given 10 artists and their trend/style',
    inputs=[
        gr.inputs.Image(label='Upload a drawing/image', type='file')
    ],
    outputs=[
        gr.outputs.Label(label='Artists', num_top_classes=5, type='auto'),
        gr.outputs.Label(label='Styles', num_top_classes=5, type='auto'),
        gr.outputs.Image(label='Prediction with GradCAM')
    ],
    examples=EXAMPLES,
)
iface.launch(debug=True, enable_queue=True)