import tensorflow as tf
import numpy as np
from audio_utils import audio_to_spectrogram

MODEL_PATH = "models/audio_vit_savedmodel"

model = tf.saved_model.load(MODEL_PATH)
infer = model.signatures["serving_default"]


def predict_audio(wav_file):
    spec_img = audio_to_spectrogram(wav_file)

    x = spec_img.astype("float32") / 255.0
    x = np.expand_dims(x, axis=0)

    preds = infer(tf.constant(x))
    prob = list(preds.values())[0].numpy()[0][0]

    label = "Fake" if prob >= 0.5 else "Real"
    confidence = prob * 100

    return label, round(confidence, 2), spec_img