import gradio as gr
from transformers import AutoModel

# Загрузка модели
model = AutoModel.from_pretrained("SnowFlash383935/DigitalEduTransformers", trust_remote_code=True)

def getresult(
    sex: str,
    has_photo: bool,
    has_mobile: bool,
    followers_count: int,
    relation: str,
    life_main: str,
    people_main: str
) -> bool:
    # Конвертация входных данных в числовой формат
    input_data = [[
        1 if sex == "Мужчина" else 0,
        int(has_photo),
        int(has_mobile),
        followers_count,
        RELATION_OPTIONS.index(relation),
        LIFE_MAIN_OPTIONS.index(life_main),
        PEOPLE_MAIN_OPTIONS.index(people_main)
    ]]
    
    # Получение предсказания (модель возвращает [True] или [False])
    return model(input_data)[0]

# Конфигурация выпадающих списков
RELATION_OPTIONS = [
    "Не указано",
    "Не женат/не замужем",
    "Есть друг/подруга",
    "Помолвлен(а)",
    "Женат/замужем",
    "Всё сложно",
    "В активном поиске",
    "Влюблен(а)",
    "В гражданском браке"
]

LIFE_MAIN_OPTIONS = [
    "Не указано",
    "Семья и дети",
    "Карьера и деньги",
    "Развлечения и отдых",
    "Наука и исследования",
    "Совершенствование мира",
    "Саморазвитие",
    "Красота и искусство",
    "Слава и влияние"
]

PEOPLE_MAIN_OPTIONS = [
    "Не указано",
    "Ум и креативность",
    "Доброта и честность",
    "Красота и здоровье",
    "Власть и богатство",
    "Смелость и упорство",
    "Юмор и жизнелюбие"
]

# Примеры с предварительно вычисленными результатами
EXAMPLES = [
    ["Мужчина", True, True, 150, "Помолвлен(а)", "Не указано", "Не указано"],
    ["Женщина", True, True, 294, "Не указано", "Не указано", "Не указано"],
    ["Мужчина", False, False, 0, "В активном поиске", "Карьера и деньги", "Власть и богатство"],
    ["Женщина", True, False, 1000, "Женат/замужем", "Семья и дети", "Доброта и честность"]
]

# Создание интерфейса
with gr.Blocks(title="Прогнозирование покупки курса") as demo:
    gr.Markdown("## Прогнозирование покупки обучающего курса")
    
    with gr.Row():
        with gr.Column():
            inputs = [
                gr.Radio(["Женщина", "Мужчина"], label="Пол", value="Женщина"),
                gr.Checkbox(label="Есть фото профиля"),
                gr.Checkbox(label="Привязан телефон"),
                gr.Number(label="Количество подписчиков", precision=0),
                gr.Dropdown(RELATION_OPTIONS, label="Семейное положение", value="Не указано"),
                gr.Dropdown(LIFE_MAIN_OPTIONS, label="Главное в жизни", value="Не указано"),
                gr.Dropdown(PEOPLE_MAIN_OPTIONS, label="Главное в людях", value="Не указано")
            ]
            
        with gr.Column():
            output = gr.Checkbox(label="Купит ли курс", interactive=False)
            submit_btn = gr.Button("Предсказать", variant="primary")

    # Примеры с автоматическим вычислением
    gr.Examples(
        examples=EXAMPLES,
        inputs=inputs,
        outputs=output,
        fn=getresult,
        cache_examples=False,
        label="Примеры:",
        run_on_click=True  # Вычислять при клике на пример
    )

    # Обработчик кнопки
    submit_btn.click(
        fn=getresult,
        inputs=inputs,
        outputs=output
    )
    
    demo.launch()