import streamlit as st
import pandas as pd
import numpy as np
import pickle
import json
import plotly.graph_objects as go
import plotly.express as px

# Настройка страницы
st.set_page_config(
    page_title="Классификатор пациентов",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'About': "Классификатор пациентов для определения группы лечения"
    }
)

# Применяем кастомный CSS
st.markdown("""
    <style>
    .stNumberInput input {
        width: 150px;
    }
    .stSelectbox select {
        width: 150px;
    }
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    div[data-testid="stSidebarNav"] {
        padding-top: 0rem;
    }
    .sidebar .sidebar-content {
        width: 300px;
    }
    .metric-card {
        background-color: #f0f2f6;
        border-radius: 0.5rem;
        padding: 1rem;
        margin: 0.5rem 0;
    }
    .small-font {
        font-size: 0.8rem;
    }
    </style>
    """, unsafe_allow_html=True)

# Загрузка модели и данных
@st.cache_resource
def load_model():
    with open('model.pkl', 'rb') as file:
        model = pickle.load(file)
    with open('scaler.pkl', 'rb') as file:
        scaler = pickle.load(file)
    with open('data.json', 'r', encoding='utf-8') as f:
        data = json.load(f)
    return model, scaler, data

def create_gauge_chart(value, title):
    fig = go.Figure(go.Indicator(
        mode = "gauge+number",
        value = value * 100,
        domain = {'x': [0, 1], 'y': [0, 1]},
        title = {'text': title},
        gauge = {
            'axis': {'range': [None, 100]},
            'bar': {'color': "darkblue"},
            'steps': [
                {'range': [0, 33], 'color': "lightgray"},
                {'range': [33, 66], 'color': "gray"},
                {'range': [66, 100], 'color': "darkgray"}
            ],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': value * 100
            }
        }
    ))
    fig.update_layout(height=200)
    return fig

model, scaler, data = load_model()
feature_list = data['features']
categorical_features = data['categorical_features']
numeric_features = data['numeric_features']
categorical_options = data['categorical_options']
group_names = data['group_names']
numeric_defaults = data['numeric_defaults']
categorical_defaults = data['categorical_defaults']

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

## Примеры:
# Примеры типичных случаев
st.header('Примеры типичных случаев')

# Создаем примеры для каждой группы (замените значения на реальные типичные случаи)
example_cases = {
    'контр': {
        'title': 'Пример контрольной группы',
        'description': 'Типичный случай для контрольной группы',
        'values': {
            'Шкала_ПВА 0': 3,
            'HADS_Тревога 0': 5,
            'HADS_Тревога 3': 5,
            'HADS_Депрессия 0': 7,
            'HADS_Депрессия 2': 3,
            'Пенсильванская_шкала_ПВА 4': 0,
            'HADS_Депрессия 3': 7,
            'HADS_Депрессия 4': 4,
            'CIWAAr 0': 7,
            'CIWAAr 1': 4,
            'ЭМОТИВ_ЛЕОНГАРД': 18,
            'Отяг_наслед_алк': 0,
            'Форма_употреб': 0,
            'Патология_родов': 0,
            'СП_в_ран_дестве': 0,
            'ЧМТ': 0,
        }
    },
    'топирамат': {
        'title': 'Пример группы топирамата',
        'description': 'Типичный случай для группы топирамата',
        'values': {
            'Шкала_ПВА 0': 12,
            'HADS_Тревога 0': 3,
            'HADS_Тревога 3': 2,
            'HADS_Депрессия 0': 7,
            'HADS_Депрессия 2': 5,
            'Пенсильванская_шкала_ПВА 4': 8,
            'HADS_Депрессия 3': 3,
            'HADS_Депрессия 4': 3,
            'CIWAAr 0': 12,
            'CIWAAr 1': 17,
            'ЭМОТИВ_ЛЕОНГАРД': 21,
            'Отяг_наслед_алк': 1,
            'Форма_употреб': 1,
            'Патология_родов': 0,
            'СП_в_ран_дестве': 0,
            'ЧМТ': 1,
        }
    },
    'леветирацетам': {
        'title': 'Пример группы леветирацетама',
        'description': 'Типичный случай для группы леветирацетама',
        'values': {
            'Шкала_ПВА 0': 6,
            'HADS_Тревога 0': 2,
            'HADS_Тревога 3': 2,
            'HADS_Депрессия 0': 6,
            'HADS_Депрессия 2': 6,
            'Пенсильванская_шкала_ПВА 4': 5,
            'HADS_Депрессия 3': 4,
            'HADS_Депрессия 4': 3,
            'CIWAAr 0': 11,
            'CIWAAr 1': 2,
            'ЭМОТИВ_ЛЕОНГАРД': 18,
            'Отяг_наслед_алк': 1,
            'Форма_употреб': 2,
            'Патология_родов': 0,
            'СП_в_ран_дестве': 0,
            'ЧМТ': 0,
        }
    }
}

# Отображение примеров в виде карточек
col1, col2, col3 = st.columns(3)

def create_example_card(title, description, values, key):
    with st.container():
        st.markdown(f"""
            <div style="
                padding: 1rem;
                border-radius: 0.5rem;
                border: 1px solid #e0e0e0;
                margin: 0.5rem 0;
                background-color: white;
                box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                <h3>{title}</h3>
                <p>{description}</p>
            </div>
        """, unsafe_allow_html=True)
        if st.button('Применить этот пример', key=key):
            st.session_state['current_example'] = values

with col1:
    create_example_card(
        example_cases['контр']['title'],
        example_cases['контр']['description'],
        example_cases['контр']['values'],
        'example_control'
    )

with col2:
    create_example_card(
        example_cases['топирамат']['title'],
        example_cases['топирамат']['description'],
        example_cases['топирамат']['values'],
        'example_topirama'
    )

with col3:
    create_example_card(
        example_cases['леветирацетам']['title'],
        example_cases['леветирацетам']['description'],
        example_cases['леветирацетам']['values'],
        'example_levetiracetam'
    )

# # Описание в основной части
# with st.container():
#     st.markdown("""
#         ### О системе
#         Данная система помогает классифицировать пациентов по группам на основе введенных параметров.
        
#         #### Как использовать:
#         1. Заполните все поля в форме слева
#         2. Нажмите кнопку "Выполнить классификацию"
#         3. Получите результат с вероятностями принадлежности к каждой группе
        
#         #### Группы классификации:
#         - Контрольная группа
#         - Группа топирамата
#         - Группа леветирацетама
#     """)

# Боковая панель для ввода данных
with st.sidebar:
    st.header('Ввод данных')
    
    # Создаем вкладки для разных типов параметров
    tab1, tab2 = st.tabs(["📊 Непрерывные", "📋 Категориальные"])
    
    input_data = {}
    
    # Получаем значения из примера, если он выбран
    current_example = st.session_state.get('current_example', {})
    
    with tab1:
        # Числовые признаки
        for i in range(0, len(numeric_features), 2):
            col1, col2 = st.columns(2)
            with col1:
                if i < len(numeric_features):
                    feature = numeric_features[i]
                    # Явно преобразуем значение в float
                    default_value = float(current_example.get(feature, numeric_defaults[feature]))
                    input_data[feature] = st.number_input(
                        f'{feature}',
                        value=default_value,
                        step=0.01,  # Добавляем шаг
                        format="%.2f",
                        help=f"Среднее: {float(numeric_defaults[feature]):.2f}"
                    )
            with col2:
                if i + 1 < len(numeric_features):
                    feature = numeric_features[i + 1]
                    # Явно преобразуем значение в float
                    default_value = float(current_example.get(feature, numeric_defaults[feature]))
                    input_data[feature] = st.number_input(
                        f'{feature}',
                        value=default_value,
                        step=0.01,  # Добавляем шаг
                        format="%.2f",
                        help=f"Среднее: {float(numeric_defaults[feature]):.2f}"
                    )
    
    with tab2:
        # Категориальные признаки
        for i in range(0, len(categorical_features), 2):
            col1, col2 = st.columns(2)
            with col1:
                if i < len(categorical_features):
                    feature = categorical_features[i]
                    options = categorical_options[feature]
                    default_value = current_example.get(feature, categorical_defaults[feature])
                    default_idx = options.index(default_value) if default_value in options else 0
                    input_data[feature] = st.selectbox(
                        f'{feature}',
                        options,
                        index=default_idx,
                        help=f"Типичное: {categorical_defaults[feature]}"
                    )
            with col2:
                if i + 1 < len(categorical_features):
                    feature = categorical_features[i + 1]
                    options = categorical_options[feature]
                    default_value = current_example.get(feature, categorical_defaults[feature])
                    default_idx = options.index(default_value) if default_value in options else 0
                    input_data[feature] = st.selectbox(
                        f'{feature}',
                        options,
                        index=default_idx,
                        help=f"Типичное: {categorical_defaults[feature]}"
                    )
    
    # Кнопка классификации
    if st.button('Выполнить классификацию', use_container_width=True):
        with st.spinner('Выполняется классификация...'):
            # Преобразование входных данных
            input_df = pd.DataFrame([input_data])
            
            # One-hot encoding для категориальных признаков
            input_df_encoded = pd.get_dummies(input_df, columns=categorical_features)
            
            # Масштабирование числовых признаков
            if numeric_features:
                input_df_encoded[numeric_features] = scaler.transform(input_df_encoded[numeric_features])
            
            # Убедитесь, что все необходимые столбцы присутствуют
            for col in feature_list:
                if col not in input_df_encoded.columns:
                    input_df_encoded[col] = 0
            
            X_pred = input_df_encoded[feature_list]
            prediction = model.predict(X_pred)
            probabilities = model.predict_proba(X_pred)[0]
            
            # Сохраняем результаты в session state
            st.session_state['prediction'] = prediction[0]
            st.session_state['probabilities'] = probabilities
            st.session_state['input_data'] = input_df

# Отображение результатов в основной части
if 'prediction' in st.session_state:
    st.header('Результаты классификации')
    
    # Создаем контейнер для основных метрик
    with st.container():
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.markdown("""
                <div class="metric-card">
                    <h3>Предсказанная группа</h3>
                    <h2 style="color: darkblue;">{}</h2>
                </div>
            """.format(group_names[str(st.session_state['prediction'])]), unsafe_allow_html=True)
        
        with col2:
            max_prob = max(st.session_state['probabilities'])
            st.markdown("""
                <div class="metric-card">
                    <h3>Уверенность модели</h3>
                    <h2 style="color: darkblue;">{:.1%}</h2>
                </div>
            """.format(max_prob), unsafe_allow_html=True)
        
        with col3:
            second_best_prob = sorted(st.session_state['probabilities'])[-2]
            st.markdown("""
                <div class="metric-card">
                    <h3>Вторая вероятная группа</h3>
                    <h2 style="color: darkblue;">{:.1%}</h2>
                </div>
            """.format(second_best_prob), unsafe_allow_html=True)

    # Визуализация вероятностей
    st.subheader('Распределение вероятностей по группам')
    
    # График вероятностей
    prob_df = pd.DataFrame({
        'Группа': [group_names[str(i)] for i in range(len(group_names))],
        'Вероятность': st.session_state['probabilities']
    })
    
    fig = px.bar(prob_df, 
                 x='Группа', 
                 y='Вероятность',
                 color='Вероятность',
                 color_continuous_scale='Blues',
                 text=prob_df['Вероятность'].apply(lambda x: f'{x:.1%}'))
    
    fig.update_layout(
        height=400,
        yaxis_range=[0, 1],
        yaxis_tickformat='.0%',
        showlegend=False
    )
    
    st.plotly_chart(fig, use_container_width=True)
    
    # Детальная информация
    with st.expander("Детальная информация"):
        tab1, tab2 = st.tabs(["📊 Введенные данные", "ℹ️ Пояснение"])
        
        with tab1:
            # Отображаем введенные данные в виде двух таблиц
            col1, col2 = st.columns(2)
            
            with col1:
                st.subheader("Числовые параметры")
                numeric_data = st.session_state['input_data'][numeric_features]
                st.dataframe(numeric_data.T.style.format("{:.2f}"))
            
            with col2:
                st.subheader("Категориальные параметры")
                categorical_data = st.session_state['input_data'][categorical_features]
                st.dataframe(categorical_data.T)
        
        with tab2:
            st.markdown("""
                ### Как интерпретировать результаты:
                
                - **Предсказанная группа** - основная рекомендованная группа для пациента
                - **Уверенность модели** - вероятность принадлежности к предсказанной группе
                - **Вторая вероятная группа** - вероятность принадлежности ко второй наиболее вероятной группе
                
                ### Примечания:
                - Чем выше уверенность модели, тем надежнее предсказание
                - При уверенности ниже 50% рекомендуется дополнительное обследование
                - Результаты модели носят рекомендательный характер
            """)
# Добавьте кнопку сброса примера
    if st.session_state.get('current_example'):
        if st.button('Сбросить пример'):
            del st.session_state['current_example']
            st.experimental_rerun()
# Добавляем footer
st.markdown("""
    ---
    <div class="small-font">
        © 2023 Классификатор пациентов | Версия 1.0
    </div>
""", unsafe_allow_html=True)