#!/usr/bin/env python3
"""Inference script for the DCSkyCam helicopter type multi-class classifier.

Usage:
    python inference.py <image_path> [--top-k 3]

Requires: tflite-runtime (on Raspberry Pi) or tensorflow (on desktop).
"""

import sys
import numpy as np
from PIL import Image

try:
    import tflite_runtime.interpreter as tflite
except ImportError:
    import tensorflow.lite as tflite

MODEL_PATH = "custom_multi_e2m.tflite"
LABELS_PATH = "custom_multi_e2m-labels.txt"


def load_labels(path):
    with open(path, "r") as f:
        return [line.strip() for line in f.readlines()]


def predict(image_path, model_path=MODEL_PATH, labels_path=LABELS_PATH, top_k=3):
    """Run inference on an image and return predictions."""
    # Load model
    interpreter = tflite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    # Load labels
    labels = load_labels(labels_path)

    # Prepare image: 480x480, RGB, normalized to [0, 1]
    img = Image.open(image_path).convert("RGB")
    img = img.resize((480, 480))
    input_data = np.expand_dims(np.array(img, dtype=np.float32) / 255.0, axis=0)

    # Run inference
    input_details = interpreter.get_input_details()[0]
    interpreter.set_tensor(input_details["index"], input_data)
    interpreter.invoke()

    # Get results
    output = interpreter.get_tensor(interpreter.get_output_details()[0]["index"])[0]

    # Get top-k predictions
    top_k_idx = np.argsort(output)[::-1][:top_k]

    return labels, top_k_idx, output.tolist()


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <image_path> [--top-k N]")
        sys.exit(1)

    image_path = sys.argv[1]
    top_k = 3
    if "--top-k" in sys.argv:
        idx = sys.argv.index("--top-k")
        if idx + 1 < len(sys.argv):
            top_k = int(sys.argv[idx + 1])

    labels, top_k_idx, scores = predict(image_path, top_k=top_k)

    print(f"Image: {image_path}")
    for i, idx in enumerate(top_k_idx):
        print(f"  #{i+1} {labels[idx]}: {scores[idx]:.4f}")