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('Классификатор пациентов')

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

# Боковая панель для ввода данных
with st.sidebar:
    st.header('Ввод данных')
    
    # Создаем вкладки для разных типов параметров
    tab1, tab2 = st.tabs(["📊 Числовые", "📋 Категориальные"])
    
    input_data = {}
    
    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]
                    default_value = float(numeric_defaults[feature])
                    input_data[feature] = st.number_input(
                        f'{feature}',
                        value=default_value,
                        format="%.2f",
                        help=f"Среднее: {default_value:.2f}"
                    )
            with col2:
                if i + 1 < len(numeric_features):
                    feature = numeric_features[i + 1]
                    default_value = float(numeric_defaults[feature])
                    input_data[feature] = st.number_input(
                        f'{feature}',
                        value=default_value,
                        format="%.2f",
                        help=f"Среднее: {default_value:.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_idx = options.index(categorical_defaults[feature]) if categorical_defaults[feature] 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_idx = options.index(categorical_defaults[feature]) if categorical_defaults[feature] 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% рекомендуется дополнительное обследование
                - Результаты модели носят рекомендательный характер
            """)

# Добавляем footer
st.markdown("""
    ---
    <div class="small-font">
        © 2023 Классификатор пациентов | Версия 1.0
    </div>
""", unsafe_allow_html=True)