import os
import pandas as pd
import numpy as np
import json
import torch
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, roc_curve, confusion_matrix
from transformers import AutoTokenizer, AutoModelForCausalLM
import warnings
warnings.filterwarnings("ignore")
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
# ---- 全局配置区（医疗场景可自行修改） ----
RANDOM_SEED = 42              # 随机种子，保证实验可复现
RF_ESTIMATORS = 200           # 随机森林树数量
RF_MAX_DEPTH = 8              # 随机森林最大深度
LLM_MODEL_NAME = "/home/model/Qwen/Qwen-7B-Chat"  # 开源医疗微调大模型基座路径
DISEASE_NAME = "2型糖尿病"     # 当前筛查的疾病名称
PROMPT_TEMPLATE = """
你是专业内分泌临床医师，严格依据《中国2型糖尿病防治指南》完成解读，禁止无依据主观推断。
【患者全部临床检验指标】：{patient_indicators}
【随机森林AI模型预测结果】：{predict_result}，预测置信概率：{predict_prob}
【各检验指标对本次风险判定的贡献权重（0~1，数值越高影响越大）】：{feature_weights}
输出要求：分两段，第一段阐述核心高权重致病指标归因，标注指标临床正常参考范围；
第二段说明次要辅助指标影响，语言通俗专业，总字数控制500字以内。
"""

# ---- 步骤1：模拟生成医疗结构化表格数据 ----
def generate_medical_data(sample_num=2000):
    """模拟生成糖尿病筛查患者表格数据集"""
    np.random.seed(RANDOM_SEED)
    data = {}
    # 连续数值临床特征
    data["age"] = np.random.randint(30, 75, sample_num)  # 年龄30-75
    data["bmi"] = np.round(np.random.normal(24, 3.5, sample_num), 1)  # BMI指数
    data["fast_glucose"] = np.round(np.random.normal(6.2, 1.8, sample_num), 1)  # 空腹血糖
    data["hba1c"] = np.round(np.random.normal(5.8, 1.2, sample_num), 1)  # 糖化血红蛋白
    data["triglyceride"] = np.round(np.random.normal(1.5, 0.6, sample_num), 1)  # 甘油三酯
    # 无序分类特征
    data["gender"] = np.random.randint(0, 2, sample_num)  # 0女 1男
    data["smoke"] = np.random.randint(0, 3, sample_num)  # 0无 1轻度 2重度吸烟
    # 目标标签：1=确诊糖尿病高风险，0=健康低风险
    label = np.zeros(sample_num)
    # 模拟患病逻辑：血糖、糖化、BMI偏高则标记为1
    risk_mask = (data["fast_glucose"] > 6.8) & (data["hba1c"] > 6.2)
    label[risk_mask] = 1
    data["diabetes_risk"] = label
    df = pd.DataFrame(data)
    return df
# 生成2000条模拟患者数据
medical_df = generate_medical_data(2000)
print("原始医疗表格数据集前5行：")
print(medical_df.head())
# ---- 步骤2：医疗数据标准化预处理 ----
def medical_data_preprocess(df):
    """医疗表格清洗预处理，适配随机森林训练"""
    # 1.分离特征与分类标签
    X = df.drop("diabetes_risk", axis=1)
    y = df["diabetes_risk"]
    # 2.医疗数据缺失值填充（模拟少量缺失）
    X = X.mask(np.random.random(X.shape) < 0.03)  # 随机制造3%缺失数据
    for col in X.columns:
        if X[col].dtype in [np.float64, np.int64]:
            X[col] = X[col].fillna(X[col].mean())
    # 3.7:3划分训练集、测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=RANDOM_SEED, stratify=y
    )
    return X_train, X_test, y_train, y_test, X
X_train, X_test, y_train, y_test, full_X = medical_data_preprocess(medical_df)

# ---- 步骤3：训练医疗定制随机森林模型 ----
def train_medical_rf(X_train, y_train):
    """初始化并训练适配医疗表格的随机森林分类器"""
    rf_model = RandomForestClassifier(
        n_estimators=RF_ESTIMATORS,
        max_depth=RF_MAX_DEPTH,
        max_features="sqrt",
        random_state=RANDOM_SEED,
        n_jobs=-1  # 启用全部CPU核心加速训练
    )
    rf_model.fit(X_train, y_train)
    # 提取全局归一化特征重要性权重
    feature_names = X_train.columns.tolist()
    feature_importance = rf_model.feature_importances_
    weight_dict = dict(zip(feature_names, np.round(feature_importance, 4)))
    return rf_model, weight_dict
rf_model, global_feature_weight = train_medical_rf(X_train, y_train)
print("\n各临床指标全局贡献权重：")
for k, v in global_feature_weight.items():
    print(f"{k}: {v}")
# ---- 步骤4：模型精度评估（医疗专用指标） ----
def evaluate_rf_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    y_pred_proba = model.predict_proba(X_test)[:, 1]
    acc = accuracy_score(y_test, y_pred)
    prec = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)
    auc = roc_auc_score(y_test, y_pred_proba)
    print("\n随机森林糖尿病分类模型测试集评估结果：")
    print(f"准确率Accuracy：{np.round(acc,4)}")
    print(f"精确率Precision：{np.round(prec,4)}")
    print(f"召回率Recall：{np.round(recall,4)}")
    print(f"F1分数：{np.round(f1,4)}")
    print(f"AUC值：{np.round(auc,4)}")
    return y_pred, y_pred_proba
test_pred, test_proba = evaluate_rf_model(rf_model, X_test, y_test)

# ---- 步骤5：单患者样本导出结构化中转JSON数据 ----
def get_patient_struct_data(model, weight_dict, sample_x, sample_index=0):
    """提取单患者指标、预测结果、权重，生成标准化JSON"""
    patient_indicators = sample_x.iloc[sample_index].to_dict()
    pred_label_num = model.predict(sample_x.iloc[[sample_index]])[0]
    pred_prob = np.round(model.predict_proba(sample_x.iloc[[sample_index]])[0][int(pred_label_num)], 3)
    pred_text = f"{DISEASE_NAME}高风险" if pred_label_num == 1 else "健康低风险"
    struct_data = {
        "patient_id": sample_index,
        "indicators": patient_indicators,
        "feature_weight": weight_dict,
        "predict_text": pred_text,
        "predict_prob": float(pred_prob)
    }
    return struct_data
# 选取测试集第1位患者作为示例样本
patient_json_data = get_patient_struct_data(rf_model, global_feature_weight, X_test, sample_index=0)
print("\n单患者结构化中转JSON数据：")
print(json.dumps(patient_json_data, ensure_ascii=False, indent=2))

# ---- 步骤6：绘制特征权重可视化柱状图 ----
def plot_feature_weight_chart(weight_dict):
    """绘制临床指标贡献权重可视化图，配套归因报告展示"""
    plt.rcParams["font.sans-serif"] = ["WenQuanYi Micro Hei"]
    plt.rcParams["axes.unicode_minus"] = False
    features = list(weight_dict.keys())
    weights = list(weight_dict.values())
    plt.figure(figsize=(10, 6))
    bars = plt.bar(features, weights, color="#2E86AB")
    plt.title("糖尿病筛查-临床指标特征贡献权重分布", fontsize=14)
    plt.xlabel("检验指标名称", fontsize=12)
    plt.ylabel("归一化贡献权重（0~1）", fontsize=12)
    plt.xticks(rotation=30)
    # 柱状图标注权重数值
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height,
                f"{np.round(height,3)}", ha="center", va="bottom")
    plt.tight_layout()
    output_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "195.糖尿病筛查-临床指标特征贡献权重分布.png")
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.show()
plot_feature_weight_chart(global_feature_weight)
# ---- 步骤7：辅助诊断可视化（4张独立图 + 数据标注） ----
BASE_IMG_DIR = os.path.dirname(os.path.abspath(__file__))
plt.rcParams["font.sans-serif"] = ["WenQuanYi Micro Hei"]
plt.rcParams["axes.unicode_minus"] = False
# ---- 图1：ROC曲线 ----
def plot_roc_chart(y_test, y_pred_proba):
    """ROC曲线：评估模型对高低风险患者的区分能力"""
    fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
    auc_val = roc_auc_score(y_test, y_pred_proba)
    fig, ax = plt.subplots(figsize=(11, 8))
    ax.plot(fpr, tpr, color="#2E86AB", lw=2.5, label=f"随机森林 AUC = {auc_val:.4f}")
    ax.plot([0, 1], [0, 1], color="#AAA", lw=1.5, linestyle="--", label="随机猜测基线")
    ax.fill_between(fpr, tpr, alpha=0.12, color="#2E86AB")
    # 标注最佳约登指数点
    youden = tpr - fpr
    best_idx = np.argmax(youden)
    best_thresh = thresholds[best_idx]
    ax.scatter(fpr[best_idx], tpr[best_idx], s=120, color="#E74C3C", zorder=5, edgecolors="white", lw=2)
    ax.annotate(f"最佳阈值={best_thresh:.2f}\nTPR={tpr[best_idx]:.2f} FPR={fpr[best_idx]:.2f}",
                xy=(fpr[best_idx], tpr[best_idx]),
                xytext=(fpr[best_idx]+0.15, tpr[best_idx]-0.15),
                fontsize=11, color="#E74C3C", fontweight="bold",
                arrowprops=dict(arrowstyle="->", color="#E74C3C", lw=1.5))
    # 结论文本
    if auc_val >= 0.9:
        quality = "优秀"
    elif auc_val >= 0.8:
        quality = "良好"
    elif auc_val >= 0.7:
        quality = "尚可"
    else:
        quality = "需优化"
    n_test = len(y_test)
    ax.text(0.98, 0.15, f"AUC={auc_val:.4f}，模型区分能力：{quality}\n测试集{n_test}例，高风险{np.sum(y_test.values)}例，低风险{n_test-np.sum(y_test.values)}例",
            transform=ax.transAxes, fontsize=11, ha="right", va="bottom",
            bbox=dict(boxstyle="round,pad=0.4", facecolor="#FFF9C4", edgecolor="#F9A825", alpha=0.9))
    ax.set_title("ROC曲线 — 糖尿病风险分类模型区分能力评估", fontsize=14, fontweight="bold")
    ax.set_xlabel("假阳性率 (FPR) — 实际低风险被误判为高风险的比例", fontsize=11)
    ax.set_ylabel("真阳性率 (TPR) — 实际高风险被正确识别的比例", fontsize=11)
    ax.legend(fontsize=11, loc="lower right")
    ax.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(BASE_IMG_DIR, "195.1-ROC曲线.png"), dpi=150, bbox_inches="tight")
    plt.show()
# ---- 图2：混淆矩阵热力图 ----
def plot_confusion_matrix_chart(y_test, y_pred):
    """混淆矩阵：展示真阳性/假阳性/假阴性/真阴性的具体分布"""
    cm = confusion_matrix(y_test, y_pred)
    tn, fp, fn, tp = cm.ravel()
    fig, ax = plt.subplots(figsize=(10, 8))
    im = ax.imshow(cm, cmap="Blues", interpolation="nearest")
    labels_text = ["低风险(0)", "高风险(1)"]
    ax.set_xticks([0, 1])
    ax.set_xticklabels(labels_text, fontsize=12)
    ax.set_yticks([0, 1])
    ax.set_yticklabels(labels_text, fontsize=12)
    ax.set_xlabel("预测标签", fontsize=12)
    ax.set_ylabel("真实标签", fontsize=12)
    # 格子内标注
    cell_info = [[f"真阴性 TN={tn}", f"假阳性 FP={fp}"],
                 [f"假阴性 FN={fn}", f"真阳性 TP={tp}"]]
    for i in range(2):
        for j in range(2):
            color = "white" if cm[i, j] > cm.max() / 2 else "black"
            ax.text(j, i - 0.08, f"{cm[i, j]}", ha="center", va="center",
                    fontsize=22, fontweight="bold", color=color)
            ax.text(j, i + 0.18, cell_info[i][j], ha="center", va="center",
                    fontsize=11, color=color, style="italic")
    plt.colorbar(im, ax=ax, fraction=0.046)
    # 底部结论
    total = tn + fp + fn + tp
    sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
    specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
    precision_val = tp / (tp + fp) if (tp + fp) > 0 else 0
    conclusion = (f"测试集共{total}例 | 灵敏度(召回率)={sensitivity:.2%} | 特异度={specificity:.2%} | 阳性预测值(精确率)={precision_val:.2%}\n"
                  f"漏诊{fn}例（假阴性） — 误诊{fp}例（假阳性）")
    ax.text(0.5, -0.15, conclusion, transform=ax.transAxes, fontsize=11, ha="center",
            bbox=dict(boxstyle="round,pad=0.4", facecolor="#FFF9C4", edgecolor="#F9A825", alpha=0.9))
    ax.set_title("混淆矩阵 — 糖尿病风险分类结果分布", fontsize=14, fontweight="bold")
    plt.tight_layout()
    plt.savefig(os.path.join(BASE_IMG_DIR, "195.2-混淆矩阵.png"), dpi=150, bbox_inches="tight")
    plt.show()
# ---- 图3：风险概率分布直方图 ----
def plot_risk_distribution(y_test, y_pred_proba):
    """展示高低风险人群的预测概率分布，评估模型分离度"""
    prob_neg = y_pred_proba[y_test.values == 0]
    prob_pos = y_pred_proba[y_test.values == 1]
    fig, ax = plt.subplots(figsize=(12, 7))
    ax.hist(prob_neg, bins=30, alpha=0.65, color="#27AE60",
            label=f"实际低风险(n={len(prob_neg)}，均值={np.mean(prob_neg):.3f})", edgecolor="white")
    ax.hist(prob_pos, bins=30, alpha=0.65, color="#E74C3C",
            label=f"实际高风险(n={len(prob_pos)}，均值={np.mean(prob_pos):.3f})", edgecolor="white")
    ax.axvline(x=0.5, color="black", linestyle="--", lw=2, label="决策阈值0.5")
    # 标注重叠区域
    overlap_low = np.sum((prob_neg > 0.3) & (prob_neg < 0.7))
    overlap_high = np.sum((prob_pos > 0.3) & (prob_pos < 0.7))
    ax.axvspan(0.3, 0.7, alpha=0.08, color="#F39C12", label=f"模糊区间(0.3~0.7)：低风险{overlap_low}例/高风险{overlap_high}例")
    # 标注均值
    ax.axvline(x=np.mean(prob_neg), color="#27AE60", linestyle=":", lw=1.8)
    ax.axvline(x=np.mean(prob_pos), color="#E74C3C", linestyle=":", lw=1.8)
    ax.text(np.mean(prob_neg), ax.get_ylim()[1]*0.9, f"低风险均值\n{np.mean(prob_neg):.3f}", fontsize=10, ha="center", color="#27AE60")
    ax.text(np.mean(prob_pos), ax.get_ylim()[1]*0.85, f"高风险均值\n{np.mean(prob_pos):.3f}", fontsize=10, ha="center", color="#E74C3C")
    # 结论
    separation = np.mean(prob_pos) - np.mean(prob_neg)
    quality = "优秀" if separation > 0.6 else "良好" if separation > 0.4 else "一般"
    ax.text(0.02, 0.98, f"高低风险均值差={separation:.3f}，分离度：{quality}",
            transform=ax.transAxes, fontsize=11, ha="left", va="top",
            bbox=dict(boxstyle="round,pad=0.4", facecolor="#FFF9C4", edgecolor="#F9A825", alpha=0.9))
    ax.set_title("预测风险概率分布 — 高低风险人群分离度", fontsize=14, fontweight="bold")
    ax.set_xlabel("预测为高风险的概率", fontsize=12)
    ax.set_ylabel("样本数", fontsize=12)
    ax.legend(fontsize=10, loc="center right")
    plt.tight_layout()
    plt.savefig(os.path.join(BASE_IMG_DIR, "195.3-风险概率分布.png"), dpi=150, bbox_inches="tight")
    plt.show()
# ---- 图4：单患者雷达图 ----
def plot_patient_radar(patient_data):
    """单患者雷达图：各指标与临床正常范围对比，超限区域红色突出"""
    indicator_names = {"age": "年龄", "bmi": "BMI", "fast_glucose": "空腹血糖",
                       "hba1c": "糖化血红蛋白", "triglyceride": "甘油三酯",
                       "gender": "性别", "smoke": "吸烟"}
    # 临床正常参考范围
    normal_ref = {"age": ("—", 75), "bmi": ("18.5~24", 28), "fast_glucose": ("3.9~6.1", 7.0),
                  "hba1c": ("4.0~6.0", 6.5), "triglyceride": ("0.56~1.70", 2.0),
                  "gender": ("—", 1), "smoke": ("0", 2)}
    patient_vals = patient_data["indicators"]
    radar_features = ["fast_glucose", "hba1c", "bmi", "triglyceride", "age", "smoke"]
    radar_labels = [indicator_names[f] for f in radar_features]
    # 归一化到0~1（相对正常上限）
    normal_max = {k: v[1] for k, v in normal_ref.items()}
    patient_norm = [min(float(patient_vals[f]) / normal_max[f], 1.5) / 1.5 for f in radar_features]
    normal_norm = [1.0 / 1.5 for _ in radar_features]
    # 标记超限
    exceeded = [float(patient_vals[f]) > normal_max[f] for f in radar_features]
    val_labels = [f"{patient_vals[f]}{'↑' if e else ''}" for f, e in zip(radar_features, exceeded)]
    angles = np.linspace(0, 2 * np.pi, len(radar_features), endpoint=False).tolist()
    patient_norm += patient_norm[:1]
    normal_norm += normal_norm[:1]
    angles += angles[:1]
    fig, ax = plt.subplots(figsize=(11, 9), subplot_kw=dict(polar=True))
    ax.plot(angles, patient_norm, "o-", color="#E74C3C", lw=2, markersize=7, label="患者指标")
    ax.fill(angles, patient_norm, alpha=0.2, color="#E74C3C")
    ax.plot(angles, normal_norm, "s--", color="#2E86AB", lw=2, label="临床正常上限")
    ax.fill(angles, normal_norm, alpha=0.08, color="#2E86AB")
    # 超限点标红放大
    for i, is_exc in enumerate(exceeded):
        if is_exc:
            ax.scatter(angles[i], patient_norm[i], s=180, color="#E74C3C", edgecolors="black", lw=1.5, zorder=5)
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels([f"{l}\n{v}" for l, v in zip(radar_labels, val_labels)], fontsize=11)
    ax.set_ylim(0, 1.15)
    # 结论文本
    exc_count = sum(exceeded)
    exc_names = "、".join([radar_labels[i] for i, e in enumerate(exceeded) if e])
    conclusion = f"患者预测：{patient_data['predict_text']}（置信{patient_data['predict_prob']:.1%}）\n{exc_count}项指标超正常上限：{exc_names if exc_names else '无'}"
    ax.set_title(f"单患者临床指标雷达图\n{conclusion}", fontsize=13, fontweight="bold", pad=25)
    ax.legend(loc="upper right", bbox_to_anchor=(1.28, 1.12), fontsize=11)
    plt.tight_layout()
    plt.savefig(os.path.join(BASE_IMG_DIR, "195.4-单患者雷达图.png"), dpi=150, bbox_inches="tight")
    plt.show()
# 依次绘制4张独立图
plot_roc_chart(y_test, test_proba)
plot_confusion_matrix_chart(y_test, test_pred)
plot_risk_distribution(y_test, test_proba)
plot_patient_radar(patient_json_data)

# ---- 步骤8：拼接大模型推理Prompt模板 ----
def build_llm_prompt(struct_data):
    indicators_text = json.dumps(struct_data["indicators"], ensure_ascii=False, indent=1)
    weight_text = json.dumps(struct_data["feature_weight"], ensure_ascii=False, indent=1)
    prompt = PROMPT_TEMPLATE.format(
        patient_indicators=indicators_text,
        predict_result=struct_data["predict_text"],
        predict_prob=struct_data["predict_prob"],
        feature_weights=weight_text
    )
    return prompt
llm_input_prompt = build_llm_prompt(patient_json_data)
print("\n输入大模型完整Prompt文本：")
print(llm_input_prompt)
# ---- 步骤9：加载医疗大模型并生成归因解读 ----
def load_medical_llm(model_name):
    """加载轻量化医疗大模型与分词器"""
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    llm_model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True
    )
    return tokenizer, llm_model
def generate_medical_explain(tokenizer, llm, prompt):
    """输入Prompt，生成医学归因解读文本"""
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = llm.generate(
        **inputs, max_new_tokens=600, temperature=0.3, top_p=0.7
    )
    explain_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return explain_text
# 加载模型并生成归因（无GPU环境可注释本段，直接查看Prompt逻辑）
tokenizer, llm = load_medical_llm(LLM_MODEL_NAME)
medical_explain_report = generate_medical_explain(tokenizer, llm, llm_input_prompt)
print("\n=====大模型输出完整医学归因解读报告=====")
print(medical_explain_report)