import os

import gradio as gr
import tensorflow as tf
from keras_tuner import HyperParameters
from huggingface_hub import hf_hub_download

from src.models import MakeHyperModel
from src.preprocessing import get_data_augmentation
from src.config import IMAGE_SIZE

data_augmentation = get_data_augmentation()
cache_dir = os.path.join('hf_hub')

# Download models
for f in ['checkpoint', 'checkpoint.data-00000-of-00001', 'checkpoint.index']:
    old_name = hf_hub_download(repo_id="eddydecena/cat-vs-dog", filename=f"tuner_model/cat-vs-dog/trial_0484d8d758a5ef7b91ca97d334ba7870/checkpoints/epoch_0/{f}", cache_dir=cache_dir)
    temp_value = old_name.split('/')
    temp_value.pop(-1)
    path = '/'.join(temp_value)
    os.rename(old_name, os.path.join(path, f))

# Download examples images
examples_cache_dir = 'examples'
for image in ['cat1.jpg', 'cat2.jpg', 'dog1.jpeg', 'dog2.jpeg']:
    old_name = hf_hub_download(repo_id="eddydecena/cat-vs-dog", filename=f"examples/{image}", cache_dir=examples_cache_dir)
    temp_value = old_name.split('/')
    temp_value.pop(-1)
    path = '/'.join(temp_value)
    os.rename(old_name, os.path.join(path, image))
    
latest = tf.train.latest_checkpoint(cache_dir)
hypermodel = MakeHyperModel(input_shape=IMAGE_SIZE + (3,), num_classes=2, data_augmentation=data_augmentation)
model = hypermodel.build(hp=HyperParameters())
model.load_weights(latest).expect_partial()

def cat_vs_dog(image):
    img_array = tf.constant(image, dtype=tf.float32)
    img_array = tf.expand_dims(img_array, 0)
    predictions = model.predict(img_array)
    score = predictions[0]
    return {'cat': float((1 - score)), 'dog': float(score)}

iface = gr.Interface(
    cat_vs_dog, 
    gr.inputs.Image(shape=IMAGE_SIZE), 
    gr.outputs.Label(num_top_classes=2),
    capture_session=True,
    interpretation="default",
    examples=[
        [f"{examples_cache_dir}/cat1.jpg"],
        [f"{examples_cache_dir}/cat2.jpg"],
        [f"{examples_cache_dir}/dog1.jpeg"],
        [f"{examples_cache_dir}/dog2.jpeg"]
    ])

if __name__ == "__main__":
    iface.launch()