"""
Trashpanzee - Local / Hub Inference Verification Script
"""
import os
import sys
import numpy as np
from PIL import Image

# Multi-fallback interpreter import
try:
    from ai_edge_litert.interpreter import Interpreter
except ImportError:
    try:
        from tflite_runtime.interpreter import Interpreter
    except ImportError:
        try:
            import tensorflow as tf
            Interpreter = tf.lite.Interpreter
        except ImportError:
            sys.exit(
                "Error: No TFLite runtime found. Please install ai-edge-litert or tflite-runtime or tensorflow:\n"
                "  pip install ai-edge-litert pillow numpy"
            )

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_MODEL = os.path.join(SCRIPT_DIR, "model_unquant.tflite")
DEFAULT_LABELS = os.path.join(SCRIPT_DIR, "labels.txt")

def load_labels(labels_file=DEFAULT_LABELS):
    labels = []
    if not os.path.exists(labels_file):
        return ["Trash Can", "Not Trash"]
    with open(labels_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                name = line.split(" ", 1)[-1]
                labels.append("Not Trash" if name.lower() == "not tash" else name)
    return labels

def classify_image(image_path, model_path=DEFAULT_MODEL, labels_path=DEFAULT_LABELS):
    labels = load_labels(labels_path)
    
    interpreter = Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    
    # Load and resize
    image = Image.open(image_path).convert("RGB").resize((224, 224))
    input_data = (np.array(image, dtype=np.float32) / 127.5) - 1.0
    input_data = np.expand_dims(input_data, axis=0)
    
    # Inference
    interpreter.set_tensor(input_details[0]["index"], input_data)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]["index"])[0]
    
    results = {labels[i]: float(output_data[i]) for i in range(len(labels))}
    return results

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python example_inference.py <path_to_image>")
        sys.exit(1)
    
    img_path = sys.argv[1]
    scores = classify_image(img_path)
    print(f"\nClassification results for '{img_path}':")
    for lbl, prob in sorted(scores.items(), key=lambda x: x[1], reverse=True):
        print(f"  {lbl:12s}: {prob * 100:.2f}%")