import tensorflow as tf
import numpy as np
from PIL import Image
import os

MODEL_PATH = os.path.join(
    os.path.dirname(__file__),
    "saved_model",
    "Inception_V3_Animals_Classification.h5"
)

model = tf.keras.models.load_model(MODEL_PATH)

CLASS_NAMES = ["Cat", "Dog", "Snake"]

def preprocess_image(img: Image.Image, target_size=(256, 256)):
    img = img.convert("RGB")
    img = img.resize(target_size)
    img = np.array(img).astype("float32") / 255.0
    img = np.expand_dims(img, axis=0)
    return img

def predict(img: Image.Image):
    input_tensor = preprocess_image(img)
    preds = model.predict(input_tensor)[0]

    class_idx = int(np.argmax(preds))
    confidence = float(np.max(preds))

    prob_dict = {CLASS_NAMES[i]: float(preds[i]) for i in range(len(CLASS_NAMES))}

    return CLASS_NAMES[class_idx], confidence, prob_dict