import gradio as gr
import shap
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import io
import joblib
import os

# Установка бэкенда matplotlib
plt.switch_backend('Agg')

# Загрузка моделей
model_dir = './models'
xgb_reg_model = joblib.load(os.path.join(model_dir, 'xgb_reg_model.joblib'))
xgb_clf_model = joblib.load(os.path.join(model_dir, 'xgb_clf_model.joblib'))
umap_reducer = joblib.load(os.path.join(model_dir, 'umap_reducer.joblib'))
hdbscan_clusterer = joblib.load(os.path.join(model_dir, 'hdbscan_clusterer.joblib'))

# Определение колонок (ВАЖНО: должны совпадать с обучением)
feature_columns = [
    'RIDAGEYR', 'RIAGENDR', 'BMXBMI', 'BMXWAIST', 'BPXSY1', 'BPXDI1',
    'LBXGH', 'LBXHSCRP', 'LBDTCSI', 'LBDHDD', 'LBXIN',
    'SLQ050', 'HSQ510', 'PHQ9_score',
    'LBXBPB', 'LBXBCD', 'LBXTHG', 'URXMHH'
]

# Списки признаков для конкретных моделей
X_reg_features_cols = feature_columns
X_clf_features_cols_without_AL_score = feature_columns

# Медианы для заполнения пропусков
x_train_medians_values = {
    'RIDAGEYR': 31.0, 'RIAGENDR': 2.0, 'BMXBMI': 25.8, 'BMXWAIST': 91.2, 'BPXSY1': 118.0,
    'BPXDI1': 70.0, 'LBXGH': 5.5, 'LBXHSCRP': 1.35, 'LBDTCSI': 4.55, 'LBDHDD': 51.0,
    'SLQ050': 2.0, 'HSQ510': 2.0, 'LBXIN': 10.04, 'PHQ9_score': 0.0,
    'LBXBPB': 0.76, 'LBXBCD': 0.22, 'LBXTHG': 0.51, 'URXMHH': 5.5
}
X_train_medians = pd.Series(x_train_medians_values)

thresholds_recommendation = {
    'LBXHSCRP': 2.71, 'SLQ050': 2.0, 'PHQ9_score': 2.0, 'HSQ510': 2.0,
    'LBXBPB': 1.07, 'LBXBCD': 0.35, 'LBXTHG': 0.91, 'URXMHH': 5.5
}

def generate_recommendations(participant_data):
    recommendations = []
    stage = participant_data.get('stage_label', 0)
    
    if stage == 3.0:
        recommendations.append("🔴 Уровень адаптационной нагрузки крайне высокий. Необходима немедленная консультация врача.")
    elif stage == 2.0:
        recommendations.append("🟠 Уровень адаптационной нагрузки высокий. Рекомендуется изменение образа жизни.")
    elif stage == 1.0:
        recommendations.append("🟡 Наблюдается повышенная адаптационная нагрузка. Следите за здоровьем.")

    if participant_data.get('LBXHSCRP', 0) > thresholds_recommendation['LBXHSCRP']:
        recommendations.append(f"⚠️ Повышен уровень C-реактивного белка ({participant_data['LBXHSCRP']:.2f}).")

    # Сбор симптомов
    symptoms = []
    if participant_data.get('SLQ050', 0) > thresholds_recommendation['SLQ050']: symptoms.append("сонливость")
    if participant_data.get('PHQ9_score', 0) > thresholds_recommendation['PHQ9_score']: symptoms.append("высокий стресс (PHQ-9)")
    
    if symptoms:
        recommendations.append(f"💡 Обратите внимание на: {', '.join(symptoms)}.")

    if not recommendations:
        recommendations.append("✅ Ваши показатели в норме.")
    
    return recommendations

def predict_and_recommend(*args):
    # Создание DataFrame
    input_df = pd.DataFrame([args], columns=feature_columns)
    
    # Заполнение пропусков
    for col in feature_columns:
        if pd.isna(input_df[col]).any():
            input_df[col] = input_df[col].fillna(X_train_medians.get(col, 0))

    # Предсказание AL Score
    predicted_al_score = xgb_reg_model.predict(input_df[X_reg_features_cols])[0]

    # Подготовка для классификатора и UMAP
    input_df_clf = input_df[X_clf_features_cols_without_AL_score].copy()
    input_df_clf.insert(0, 'AL_score', predicted_al_score)

    predicted_stage_label = xgb_clf_model.predict(input_df_clf)[0]
    
    # UMAP и Кластеризация
    umap_emb = umap_reducer.transform(input_df_clf)
    cluster = hdbscan_clusterer.predict(umap_emb)[0]
    
    # SHAP
    explainer = shap.TreeExplainer(xgb_reg_model)
    shap_vals = explainer.shap_values(input_df[X_reg_features_cols])
    
    plt.figure(figsize=(10, 3))
    shap.force_plot(explainer.expected_value, shap_vals[0], input_df.iloc[0], matplotlib=True, show=False)
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    plt.close()
    
    # Рекомендации
    diag_data = input_df.iloc[0].to_dict()
    diag_data['stage_label'] = predicted_stage_label
    recs = generate_recommendations(diag_data)

    return (
        round(float(predicted_al_score), 2),
        int(predicted_stage_label),
        f"{umap_emb[0,0]:.2f}, {umap_emb[0,1]:.2f}",
        str(cluster),
        buf.getvalue(),
        "\n".join(recs)
    )

# Описание интерфейса Gradio
feature_labels = {
    'RIDAGEYR': 'Возраст', 'RIAGENDR': 'Пол (1=М, 2=Ж)', 'BMXBMI': 'ИМТ',
    'BMXWAIST': 'Талия', 'BPXSY1': 'Сист. АД', 'BPXDI1': 'Диаст. АД',
    'LBXGH': 'HbA1c', 'LBXHSCRP': 'hsCRP', 'LBDTCSI': 'Холестерин',
    'LBDHDD': 'ЛПВП', 'LBXIN': 'Инсулин', 'SLQ050': 'Сонливость',
    'HSQ510': 'Боль', 'PHQ9_score': 'PHQ-9', 'LBXBPB': 'Свинец',
    'LBXBCD': 'Кадмий', 'LBXTHG': 'Ртуть', 'URXMHH': 'Фталаты'
}

with gr.Blocks() as demo:
    gr.Markdown("# Система оценки адаптационной нагрузки")
    
    inputs = []
    with gr.Row():
        with gr.Column():
            for col in feature_columns[:9]:
                inputs.append(gr.Number(label=feature_labels[col], value=X_train_medians[col]))
        with gr.Column():
            for col in feature_columns[9:]:
                inputs.append(gr.Number(label=feature_labels[col], value=X_train_medians[col]))

    btn = gr.Button("Рассчитать")
    
    with gr.Row():
        al_out = gr.Textbox(label="AL Score")
        stage_out = gr.Textbox(label="Стадия")
    
    coords_out = gr.Textbox(label="UMAP Координаты")
    cluster_out = gr.Textbox(label="Кластер")
    shap_out = gr.Image(label="Влияние признаков (SHAP)")
    recs_out = gr.Textbox(label="Рекомендации", lines=5)

    btn.click(predict_and_recommend, inputs=inputs, outputs=[al_out, stage_out, coords_out, cluster_out, shap_out, recs_out])

demo.launch()