import streamlit as st
import pandas as pd
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import f1_score


# Метод для расстановки меток
def define_popularity(num_orders):
    if num_orders < 100:
        return 0
    elif 100 <= num_orders:
        return 1


# Добавление метки популярности, где 0(не популярно), 1(популярно)
on_add_pop = st.toggle("Добавить метки популярности")
if on_add_pop:
    data = pd.read_csv('data222.csv', sep=',')
    data['Популярность'] = None
    data['Популярность'] = data['Количество заказов'].apply(define_popularity)
    data.to_csv('data222.csv', index=False)

on_start = st.toggle("Создать и обучить модель классификации по популярности")
if on_start:
    data = pd.read_csv('data222.csv')

    size_mapping = {
        '15х15': 0,
        '20х14': 1,
        '20х20': 2,
        '21х15': 3,
        '30х40': 4,
        '30х60': 5,
        '38х50': 6,
        '40х40': 7,
        '40х50': 8,
        '40х60': 9
    }
    genre_mapping = {
        "Абстракция": 0,
        "Аниме": 1,
        "Графити и стрит-арт": 2,
        "Животные": 3,
        "Историческая живопись": 4,
        "Мультфильмы": 5,
        "Натюрморт": 6,
        "Пейзаж": 7,
        "Поп-арт": 8,
        "Портрет": 9,
        "Религия": 10,
        "Ретро и винтаж": 11,
        "Спорт": 12,
        "Фантастика": 13,
        "Фильмы": 14,
        "Эротика": 15
    }

    # Преобразуем данные в числовые значения
    data['Жанр'] = data['Жанр'].map(genre_mapping)
    data['Размер'] = data['Размер'].map(size_mapping)

    X = data[['Размер', 'Жанр']]
    y = data['Популярность']

    # Разделяем данные на тестовую и обучающую выборку
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Создания дерева решений
    model = DecisionTreeClassifier()

    # Обучение модели на тестовой выборке
    model.fit(X_train, y_train)

    # Получение предсказаний на тестовом наборе данных
    y_pred = model.predict(X_test)

    # Вычисление F1-меры
    f1 = f1_score(y_test, y_pred)

    st.write(f"F1-мера: {f1}")

on_pred = st.toggle('')
if on_pred:
    # Список с вариантами для выпадающего списка
    sizes = ['15х15', '20х14', '20х20', '21х15', '30х40', '30х60', '38х50', '40х40', '40х50', '40х60']
    selected_size = st.selectbox('Выберите размер:', sizes)

    genres = ['Абстракция', 'Аниме', 'Графити и стрит-арт', 'Животные', 'Историческая живопись', 'Мультфильмы',
              'Натюрморт', 'Пейзаж', 'Поп-арт', 'Портрет', 'Религия', 'Ретро и винтаж', 'Спорт', 'Фантастика', 'Фильмы',
              'Эротика']
    selected_genre = st.selectbox('Выберите жанр:', genres)

    # Новая строка данных, для которой нужно сделать предсказание
    new_data = pd.DataFrame({'Размер': [selected_size], 'Жанр': [selected_genre]}, index=[0])

    size_mapping = {
        '15х15': 0,
        '20х14': 1,
        '20х20': 2,
        '21х15': 3,
        '30х40': 4,
        '30х60': 5,
        '38х50': 6,
        '40х40': 7,
        '40х50': 8,
        '40х60': 9
    }
    genre_mapping = {
        "Абстракция": 0,
        "Аниме": 1,
        "Графити и стрит-арт": 2,
        "Животные": 3,
        "Историческая живопись": 4,
        "Мультфильмы": 5,
        "Натюрморт": 6,
        "Пейзаж": 7,
        "Поп-арт": 8,
        "Портрет": 9,
        "Религия": 10,
        "Ретро и винтаж": 11,
        "Спорт": 12,
        "Фантастика": 13,
        "Фильмы": 14,
        "Эротика": 15
    }

    # Преобразование категориальных значений, если нужно
    new_data['Жанр'] = new_data['Жанр'].map(genre_mapping)
    new_data['Размер'] = new_data['Размер'].map(size_mapping)

    # Предсказание для новой строки
    prediction = model.predict(new_data)
    if prediction == 0:
        st.write("Предсказанная популярность: не популярно")
    else:
        st.write("Предсказанная популярность: популярно")