from huggingface_hub import hf_hub_download
import tensorflow as tf
import numpy as np
import gradio as gr

# Загрузка файла модели из Hugging Face Hub
model_path = hf_hub_download(repo_id="neuronetties/Wine", filename="wine_model.keras")

# Загрузка модели
model = tf.keras.models.load_model(model_path)

# Функция предсказания
def predict(data_row):
    try:
        # Преобразуем данные в массив float
        input_data = np.array(data_row, dtype=float).reshape(1, -1)
        
        # Предсказания модели
        predictions = model.predict(input_data)

        # Находим индекс класса с наибольшей вероятностью
        primary_class_index = predictions[0].argmax()  # Основной класс
        predicted_class = f"№{primary_class_index + 1}"  # Первый предсказанный класс

        # Находим второй предсказанный сорт (второй по вероятности)
        sorted_indices = predictions[0].argsort()
        second_class_index = sorted_indices[-2]
        second_class_probability = predictions[0][second_class_index]

        # Если вероятность второго класса больше 0.1, добавляем его
        if second_class_probability > 0.1:
            predicted_class += f" (or №{second_class_index + 1})"

        return predicted_class
    except ValueError:
        return "Пожалуйста, убедитесь, что все значения числовые."
    
# Gradio интерфейс
interface = gr.Interface(
    fn=predict,
    inputs=gr.Dataframe(
        headers=["alcohol", "malic_acid", "ash", "alcalinity_of_ash", "magnesium", 
                 "total_phenols", "flavanoids", "nonflavanoid_phenols", "proanthocyanins", 
                 "color_intensity", "hue", "od280/od315_of_diluted_wines", "proline"], 
        row_count=1, col_count=13,
        type="array"  # Гарантирует получение данных в виде массива
    ),
    outputs=gr.Textbox(label="Predicted Class"),  # Вывод через Textbox
    description="Введите данные (13 признаков) для предсказания сорта вина"
)

# Запуск приложения
interface.launch()