#!/usr/bin/env python3
"""
MobileNetV2 Image Classification using PyWebNN

- model_id: "tarekziade/mobilenet-webnn"
- backend: "cpu"
- force_download: False

Usage:

    $ pip install pywebnn pillow
    $ python demo.py /path/to/image.jpg

"""

import sys
import time
from pathlib import Path

import numpy as np
from PIL import Image
import webnn


MODEL_ID = "tarekziade/mobilenet-webnn"
BACKEND = "cpu"  # "cpu" | "gpu" | "coreml"
FORCE_DOWNLOAD = False


def load_imagenet_labels():
    labels_file = Path(__file__).with_name("imagenet_classes.txt")
    if not labels_file.exists():
        import urllib.request

        url = (
            "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
        )
        print("   [DOWNLOAD] ImageNet labels...")
        urllib.request.urlretrieve(url, labels_file)
        print("   [OK] Labels downloaded")

    return [line.strip() for line in labels_file.read_text().splitlines()]


IMAGENET_CLASSES = load_imagenet_labels()


def preprocess_image(image_path: Path) -> np.ndarray:
    img = Image.open(image_path).convert("RGB")
    img = img.resize((224, 224), Image.Resampling.BILINEAR)

    x = np.asarray(img, dtype=np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    x = (x - mean) / std

    x = np.transpose(x, (2, 0, 1))  # HWC -> CHW
    x = np.expand_dims(x, 0)  # add batch
    return x


def backend_settings(backend: str):
    if backend == "cpu":
        return False, "default", "ONNX CPU"
    if backend == "gpu":
        return True, "high-performance", "ONNX GPU"
    if backend == "coreml":
        return True, "high-performance", "CoreML (Neural Engine)"
    raise ValueError(f"Unknown BACKEND: {backend!r}")


def main():
    if len(sys.argv) != 2:
        print(f"Usage: {Path(sys.argv[0]).name} /path/to/image.jpg")
        sys.exit(2)

    image_path = Path(sys.argv[1])
    if not image_path.exists():
        print(f"Error: Image not found: {image_path}")
        sys.exit(1)

    accelerated, power, backend_name = backend_settings(BACKEND)

    print("=" * 70)
    print("MobileNetV2 Image Classification (Hugging Face Hub)")
    print("=" * 70)
    print(f"Image:   {image_path}")
    print(f"Model:   {MODEL_ID}")
    print(f"Backend: {backend_name}")
    print()

    hub = webnn.Hub()

    print("Downloading model from Hugging Face Hub...")
    t0 = time.time()
    model_files = hub.download_model(MODEL_ID, force=FORCE_DOWNLOAD)
    download_ms = (time.time() - t0) * 1000
    print(f"   [OK] Downloaded ({download_ms:.2f}ms)")
    print(f"   - Graph: {Path(model_files['graph']).name}")
    print()

    print("Loading graph...")
    t0 = time.time()
    graph = webnn.MLGraph.load(
        model_files["graph"],
        manifest_path=model_files["manifest"],
        weights_path=model_files["weights"],
    )
    load_ms = (time.time() - t0) * 1000
    print(f"   [OK] Loaded ({load_ms:.2f}ms)")
    print()

    print("Preprocessing image...")
    t0 = time.time()
    input_data = preprocess_image(image_path)
    prep_ms = (time.time() - t0) * 1000
    print(f"   [OK] {input_data.shape} ({prep_ms:.2f}ms)")
    print()

    print("Creating WebNN context...")
    ml = webnn.ML()
    context = ml.create_context(power_preference=power, accelerated=accelerated)
    print(f"   [OK] Context created (accelerated={context.accelerated})")
    print()

    print("Running inference...")
    t0 = time.time()
    results = context.compute(graph, {"input": input_data})
    inf_ms = (time.time() - t0) * 1000
    print(f"   [OK] Done ({inf_ms:.2f}ms)")
    print()

    probs = results["output"][0]
    top5 = np.argsort(probs)[-5:][::-1]

    print("Top 5 Predictions:")
    print("-" * 70)
    for rank, idx in enumerate(top5, 1):
        name = IMAGENET_CLASSES[int(idx)]
        conf = float(probs[int(idx)]) * 100.0
        print(f"   {rank}. {name:50s} {conf:6.2f}%")
    print()

    total_ms = download_ms + load_ms + prep_ms + inf_ms
    print("=" * 70)
    print("Performance Summary:")
    print(f"  - Model Download: {download_ms:.2f}ms")
    print(f"  - Graph Load:     {load_ms:.2f}ms")
    print(f"  - Preprocessing:  {prep_ms:.2f}ms")
    print(f"  - Inference:      {inf_ms:.2f}ms")
    print(f"  - Total Time:     {total_ms:.2f}ms")
    print("=" * 70)
    print()
    print(f"[OK] Done. Cache dir: {hub.cache_dir}")


if __name__ == "__main__":
    main()