import streamlit as st
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
import pickle
import joblib

# Загрузка сохраненной модели и других необходимых объектов
@st.cache_resource
def load_model():
    with open('model.pkl', 'rb') as file:
        model = pickle.load(file)
    scaler = joblib.load('scaler.joblib')
    feature_list = joblib.load('features.joblib')
    categorical_features = joblib.load('categorical_features.joblib')
    categorical_options = joblib.load('categorical_options.joblib')
    group_names = joblib.load('group_names.joblib')
    return model, scaler, feature_list, categorical_features, categorical_options, group_names

model, scaler, feature_list, categorical_features, categorical_options, group_names = load_model()

st.title('Классификатор пациентов')

# Создание формы ввода для всех необходимых признаков
def get_user_input():
    input_data = {}
    
    # Числовые признаки
    st.subheader('Числовые параметры')
    for feature in feature_list:
        if feature not in categorical_features:  # список категориальных признаков нужно определить
            input_data[feature] = st.number_input(f'{feature}', value=0.0)
    
    # Категориальные признаки
    st.subheader('Категориальные параметры')
    for feature in feature_list:
        if feature in categorical_features:  # список категориальных признаков
            options = categorical_options[feature]  # словарь с возможными значениями для каждого признака
            input_data[feature] = st.selectbox(f'{feature}', options)
    
    return input_data

# Получение данных от пользователя
user_input = get_user_input()

# Кнопка для выполнения предсказания
if st.button('Выполнить классификацию'):
    # Преобразование входных данных
    input_df = pd.DataFrame([user_input])
    
    # One-hot encoding для категориальных признаков
    input_df_encoded = pd.get_dummies(input_df, columns=categorical_features)
    
    # Убедитесь, что все необходимые столбцы присутствуют
    for col in feature_list:
        if col not in input_df_encoded.columns:
            input_df_encoded[col] = 0
    
    # Масштабирование числовых признаков
    numeric_cols = [col for col in input_df_encoded.columns if col not in categorical_features]
    input_df_encoded[numeric_cols] = scaler.transform(input_df_encoded[numeric_cols])
    
    # Получение предсказания
    prediction = model.predict(input_df_encoded[feature_list])
    
    # Отображение результата
    group_names = {0: 'контр', 1: 'топирамат', 2: 'леветирацетам'}
    st.success(f'Предсказанная группа: {group_names[prediction[0]]}')