"""NotavelAI: identify Brazilian Real banknotes in photos / identifica cédulas de Real em fotos.

Usage / uso:
    pip install tensorflow pillow numpy
    python predict.py photo.jpg [other.jpg ...]

Runs the model exactly like the Android app: the central square of the photo, resized to 224x224,
and a value is only announced with at least 70% confidence.
Roda o modelo como o app Android: o quadrado central da foto, reduzido a 224x224, e só anuncia
um valor com pelo menos 70% de confiança.
"""

import sys
from pathlib import Path

import numpy as np
import tensorflow as tf
from PIL import Image, ImageOps

HERE = Path(__file__).resolve().parent
MODEL = HERE / "tflite" / "notavelai.tflite"
LABELS = (HERE / "labels.txt").read_text(encoding="utf-8").split()
THRESHOLD = 0.7
SIZE = 224


def load(path: str) -> np.ndarray:
    """Fix EXIF rotation, cap the longer side at 1024 px (like the app), take the central square
    and resize it to 224x224 with bilinear interpolation. Values stay 0-255 (not normalized)."""
    with Image.open(path) as image:
        rgb = ImageOps.exif_transpose(image).convert("RGB")
    rgb.thumbnail((1024, 1024))
    array = np.asarray(rgb, dtype="float32")
    height, width = array.shape[:2]
    side = min(height, width)
    top, left = (height - side) // 2, (width - side) // 2
    square = array[top:top + side, left:left + side]
    return tf.image.resize(square, (SIZE, SIZE), method="bilinear").numpy()[None]


def run(image: np.ndarray) -> np.ndarray:
    interpreter = tf.lite.Interpreter(model_path=str(MODEL))
    interpreter.allocate_tensors()
    interpreter.set_tensor(interpreter.get_input_details()[0]["index"], image)
    interpreter.invoke()
    return interpreter.get_tensor(interpreter.get_output_details()[0]["index"])[0]


def main() -> None:
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    for path in sys.argv[1:]:
        probabilities = run(load(path))
        best = int(np.argmax(probabilities))
        if LABELS[best] == "fundo":
            answer = "no banknote / nenhuma cédula"
        elif probabilities[best] < THRESHOLD:
            answer = "not sure / não consegui identificar"
        else:
            answer = f"R$ {LABELS[best]}"
        print(f"{path}: {answer} ({probabilities[best]:.0%})")


if __name__ == "__main__":
    main()