import gradio as gr
from fastai.vision.all import load_learner, PILImage
import os
from pathlib import Path

# Matomo tracking
MATOMO_SCRIPT = """<!-- Matomo -->
<script>
  var _paq = window._paq = window._paq || [];
  _paq.push(['trackPageView']);
  _paq.push(['enableLinkTracking']);
  (function() {
    var u="https://matomodocker.azurewebsites.net/";
    _paq.push(['setTrackerUrl', u+'matomo.php']);
    _paq.push(['setSiteId', '11']);
    var d=document, g=d.createElement('script'), s=d.getElementsByTagName('script')[0];
    g.async=true; g.src=u+'matomo.js'; s.parentNode.insertBefore(g,s);
  })();
</script>
<!-- End Matomo Code -->
"""

# Load model
try:
    MODEL_PATH = Path("model/art_print_model4.pkl")
    if not MODEL_PATH.exists():
        raise FileNotFoundError(f"Model not found at {MODEL_PATH}")
    learn = load_learner(MODEL_PATH)
except Exception as e:
    print(f"Error loading model: {e}")
    learn = None

def classify_image(img):
    if learn is None:
        return {"Error": 1.0, "message": "Model failed to load"}
    try:
        fastai_img = PILImage.create(img)
        pred, pred_idx, probs = learn.predict(fastai_img)
        return {learn.dls.vocab[i]: float(probs[i]) for i in range(len(probs))}
    except Exception as e:
        return {"Error": 1.0, "message": str(e)}

# Interface
title = "Classification of Historical Prints"
description = """
Automatic classification of pre-digital intaglio/photographic printing techniques.
Upload an image or use the examples below.
"""

example_dir = Path("examples")
example_paths = []
if example_dir.exists():
    example_paths = [[str(f)] for f in example_dir.glob("*.jpg")][:10]

with gr.Blocks(title=title) as demo:
    demo.head = MATOMO_SCRIPT

    with gr.Row(variant="compact"):
        with gr.Column(scale=1):
            gr.Image("ms_logo_wtrasp.png", show_label=False, width=100, show_download_button=False)
        with gr.Column(scale=4):
            gr.Markdown(f"## {title}")
            gr.Markdown(description)

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Upload Image")
            examples = gr.Examples(
                examples=example_paths,
                inputs=img_input,
                label="Example Images"
            )
        with gr.Column():
            label_output = gr.Label(num_top_classes=3, label="Classification Results")

    img_input.change(fn=classify_image, inputs=img_input, outputs=label_output)

demo.launch()