#!/usr/bin/env python3
"""Inference script for the DCSkyCam helicopter binary classifier.

Usage:
    python inference.py <image_path>

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_heli.tflite"
LABELS_PATH = "custom_heli-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):
    """Run inference on an image and return prediction."""
    # Load model
    interpreter = tflite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    # Load labels
    labels = load_labels(labels_path)

    # Prepare image: 224x224, RGB, normalized to [0, 1]
    img = Image.open(image_path).convert("RGB")
    img = img.resize((224, 224))
    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]

    pred_idx = int(np.argmax(output))
    confidence = float(np.max(output))

    return labels[pred_idx], confidence, output.tolist()


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

    image_path = sys.argv[1]
    label, confidence, scores = predict(image_path)

    print(f"Image: {image_path}")
    print(f"Prediction: {label}")
    print(f"Confidence: {confidence:.4f}")
    for i, (lbl, sc) in enumerate(zip(load_labels(LABELS_PATH), scores)):
        print(f"  {lbl}: {sc:.4f}")