"""Astro image classifier — Gradio app for a free Hugging Face Space (ZeroGPU).

The Vercel page in web/ calls it through @gradio/client:

    /classify (image) -> {"label": top class, "confidences": [{label, confidence}, ...]}

Weights are pulled from the Hugging Face model repo at startup
(MA29/astro-image-classifier), so this GitHub repo stays small and nothing
large is committed.

Architecture, read directly from the saved ensemble_model.keras:

    vgg_ensemble_input      [None,224,224,3] -> vgg_branch_model      -> softmax(11)
    densenet_ensemble_input [None,224,224,3] -> densenet_branch_model -> softmax(11)
                                             -> Average

Two things the saved graph dictates:

1. `custom_objects={"preprocess_input": ...}` is REQUIRED. The densenet branch
   holds a Lambda serialised as {"config": "preprocess_input"}, so Keras
   resolves that function by name from custom_objects at load time.

2. Both branches preprocess INTERNALLY — densenet via that Lambda, vgg via a
   GetItem/Stack/Add chain (channel swap + mean subtract). Both inputs therefore
   take RAW 0-255 RGB. Preprocessing before feeding does it twice: measured at
   71.8% correct versus 100% on the same images. See verify_model.py.
"""
try:
    import spaces  # ZeroGPU: must be imported before gradio
except ImportError:  # running locally or in Docker
    spaces = None

import json
import os
import re
import tempfile
import zipfile

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "-1")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import gradio as gr
import keras
import tensorflow as tf
from huggingface_hub import hf_hub_download
from tensorflow.keras.applications.densenet import preprocess_input

from utils.labels import CLASS_NAMES
from utils.preprocessing import feed, prepare_image

MODEL_REPO = os.getenv("MODEL_REPO", "MA29/astro-image-classifier")
MODEL_FILE = os.getenv("MODEL_FILE", "ensemble_model.keras")


def _load(path):
    custom = {"preprocess_input": preprocess_input}
    try:
        return tf.keras.models.load_model(path, custom_objects=custom, compile=False)
    except (TypeError, ValueError):
        return tf.keras.models.load_model(
            path, custom_objects=custom, compile=False, safe_mode=False
        )


# Keras error messages that name a config key the installed version doesn't know.
_UNKNOWN_KEY = [
    re.compile(r"unexpected keyword argument '(\w+)'"),
    re.compile(r"Unrecognized keyword arguments? passed to \w+: \{'(\w+)'"),
]


def _drop_key(obj, key):
    if isinstance(obj, dict):
        return {k: _drop_key(v, key) for k, v in obj.items() if k != key}
    if isinstance(obj, list):
        return [_drop_key(v, key) for v in obj]
    return obj


def load_model():
    """The model was saved with Keras 3.15, which needs Python 3.11. On an older
    Keras (e.g. a Python 3.10 Space) loading fails on config options added
    since, such as `input_axes` on initializers or `quantization_config` on
    Dense. Those are left at their defaults, so each one the installed Keras
    rejects is removed from config.json and the load retried."""
    keras.mixed_precision.set_global_policy("float32")  # saved as mixed_float16
    path = original = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    dropped = []
    while True:
        try:
            model = _load(path)
            if path != original:
                os.remove(path)  # weights are in memory now
            if dropped:
                print(f"Keras {keras.__version__}: ignored newer config keys {dropped}")
            return model
        except (TypeError, ValueError) as e:
            key = next((m.group(1) for p in _UNKNOWN_KEY for m in [p.search(str(e))] if m), None)
            if key is None or key in dropped or len(dropped) >= 20:
                raise
            dropped.append(key)
            stripped = _strip_config(path, key)
            if path != original:
                os.remove(path)  # the previous stripped copy
            path = stripped


def _strip_config(src, key):
    fd, dst = tempfile.mkstemp(suffix=".keras")
    os.close(fd)
    with zipfile.ZipFile(src) as zin, zipfile.ZipFile(dst, "w") as zout:
        for item in zin.infolist():
            data = zin.read(item.filename)
            if item.filename == "config.json":
                data = json.dumps(_drop_key(json.loads(data), key)).encode()
            zout.writestr(item, data)
    return dst


# Loaded once at startup, so the first visitor doesn't wait for the 267 MB download.
model = load_model()


# Free Gradio Spaces only come with ZeroGPU hardware, which refuses to start
# unless at least one @spaces.GPU function exists. TensorFlow can't use the
# ZeroGPU device anyway, and the ensemble answers in about a second on CPU, so
# classify() is not decorated: it runs on CPU and spends none of the GPU quota.
# This registered-but-unused function only satisfies the startup check.
if spaces is not None:
    @spaces.GPU(duration=5)
    def _zerogpu_startup_check():
        return None


def classify(image):
    if image is None:
        raise gr.Error("Upload an image.")
    preds = model.predict(feed(model, prepare_image(image)), verbose=0)[0]
    return {name.replace("_", " "): float(p) for name, p in zip(CLASS_NAMES, preds)}


with gr.Blocks(title="Astro Image Classifier") as demo:
    gr.Markdown(
        "# Astro Image Classifier\n"
        "VGG19 + DenseNet201 ensemble over 11 astronomical classes. "
        "[Code](https://github.com/Majd1029/Astro-Image-Classifier) · "
        f"[Weights](https://huggingface.co/{MODEL_REPO})"
    )
    with gr.Row():
        image = gr.Image(label="Astronomical image", type="pil")
        label = gr.Label(label="Prediction")
    run = gr.Button("Classify", variant="primary")
    run.click(classify, inputs=image, outputs=label, api_name="classify")

demo.queue(max_size=20)

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