import json
import numpy as np
import matplotlib.pyplot as plt
import os
import dashscope
from dashscope.api_entities.dashscope_response import Role

# --------------------------
# 第一步：配置基础参数
# --------------------------
# 配置通义千问客户端
api_key = os.environ.get('DASHSCOPE_API_KEY')
dashscope.api_key = api_key
# 领域知识库（医疗问诊）
MEDICAL_KNOWLEDGE_BASE = {
    "咳嗽病因": {
        "支气管炎": {
            "prior_prob": 0.25,
            "likelihood": {
                "咳嗽时长": 0.9,
                "发热": 0.8,
                "痰颜色": 0.85,
                "咳嗽时段": 0.75
            },
            "non_likelihood": {
                "咳嗽时长": 0.1,
                "发热": 0.6,
                "痰颜色": 0.2,
                "咳嗽时段": 0.3
            }
        },
        "上呼吸道感染": {
            "prior_prob": 0.2,
            "likelihood": {
                "咳嗽时长": 0.3,
                "发热": 0.7,
                "痰颜色": 0.7,
                "咳嗽时段": 0.4
            },
            "non_likelihood": {
                "咳嗽时长": 0.2,
                "发热": 0.5,
                "痰颜色": 0.3,
                "咳嗽时段": 0.5
            }
        },
        "过敏性咳嗽": {
            "prior_prob": 0.15,
            "likelihood": {
                "咳嗽时长": 0.8,
                "发热": 0.95,
                "痰颜色": 0.9,
                "咳嗽时段": 0.8
            },
            "non_likelihood": {
                "咳嗽时长": 0.15,
                "发热": 0.1,
                "痰颜色": 0.1,
                "咳嗽时段": 0.2
            }
        },
        "肺炎": {
            "prior_prob": 0.1,
            "likelihood": {
                "咳嗽时长": 0.7,
                "发热": 0.1,
                "痰颜色": 0.1,
                "咳嗽时段": 0.6
            },
            "non_likelihood": {
                "咳嗽时长": 0.2,
                "发热": 0.8,
                "痰颜色": 0.8,
                "咳嗽时段": 0.3
            }
        },
        "胃食管反流": {
            "prior_prob": 0.08,
            "likelihood": {
                "咳嗽时长": 0.6,
                "发热": 0.9,
                "痰颜色": 0.8,
                "咳嗽时段": 0.85
            },
            "non_likelihood": {
                "咳嗽时长": 0.2,
                "发热": 0.2,
                "痰颜色": 0.3,
                "咳嗽时段": 0.2
            }
        }
    }
}

# --------------------------
# 第二步：大模型模块（特征提取+结论生成+结果优化）
# --------------------------
def extract_medical_features(user_input):
    """
    大模型提取医疗特征
    :param user_input: 用户自然语言输入
    :return: 结构化特征字典
    """
    prompt = f"""
    你是一名资深呼吸科医生，请从以下用户输入中提取核心医疗特征，输出JSON格式，特征包括：
    1. 咳嗽时长：转换为天数（如两周=14）
    2. 发热：是/否（无发热=否，有发热=是）
    3. 痰颜色：白/黄/无/未知
    4. 咳嗽时段：白天/夜间/全天/未知

    用户输入：{user_input}
    输出格式：{{"咳嗽时长": 数字, "发热": "是/否", "痰颜色": "白/黄/无/未知", "咳嗽时段": "白天/夜间/全天/未知"}}
    注意：只输出JSON，不要输出其他内容。
    """

    response = dashscope.Generation.call(
        model='qwen-max',
        messages=[{"role": "user", "content": prompt}],
        result_format='message',
        temperature=0.1  # 低温度保证输出稳定
    )

    # 解析JSON输出
    content = response.output.choices[0].message.content.strip()

    # 尝试直接解析
    try:
        features = json.loads(content)
    except json.JSONDecodeError:
        # 如果失败，尝试提取JSON部分
        import re
        json_match = re.search(r'\{.*\}', content, re.DOTALL)
        if json_match:
            features = json.loads(json_match.group())
        else:
            raise ValueError(f"无法从大模型响应中提取JSON: {content}")

    return features

def generate_candidate_conclusions(features):
    """
    大模型生成候选结论集
    :param features: 结构化特征字典
    :return: 候选结论列表
    """
    prompt = f"""
    你是一名资深呼吸科医生，基于以下用户特征，生成可能的咳嗽病因候选集（最多5个），符合医学常识，按常见程度排序。
    用户特征：{json.dumps(features, ensure_ascii=False)}
    输出格式：["病因1", "病因2", ...]
    注意：只输出JSON列表，不要输出其他内容。
    """

    response = dashscope.Generation.call(
        model='qwen-max',
        messages=[{"role": "user", "content": prompt}],
        result_format='message',
        temperature=0.1
    )

    content = response.output.choices[0].message.content.strip()

    # 尝试直接解析
    try:
        conclusions = json.loads(content)
    except json.JSONDecodeError:
        # 如果失败，尝试提取JSON部分
        import re
        json_match = re.search(r'\[.*\]', content, re.DOTALL)
        if json_match:
            conclusions = json.loads(json_match.group())
        else:
            raise ValueError(f"无法从大模型响应中提取JSON: {content}")

    return conclusions

def optimize_output(results):
    """
    大模型优化最终输出（结论+置信概率+依据）
    :param results: 贝叶斯计算结果（结论: 置信概率）
    :return: 优化后的自然语言输出
    """
    prompt = f"""
    你是一名资深呼吸科医生，基于以下贝叶斯计算结果，生成专业、易懂的问诊结论，格式要求：
    1. 标题：问诊结论
    2. 列出置信概率Top3的结论，标注置信概率（百分比，保留1位小数）
    3. 核心依据：分点说明每个结论的置信度分析，包括先验依据、症状匹配依据
    4. 建议：基于最高置信度结论，给出就医、检查建议

    贝叶斯计算结果：{json.dumps(results, ensure_ascii=False)}
    注意：语言通俗易懂，避免专业术语过多，结构清晰。
    """

    response = dashscope.Generation.call(
        model='qwen-max',
        messages=[{"role": "user", "content": prompt}],
        result_format='message',
        temperature=0.3
    )

    output = response.output.choices[0].message.content
    return output

# --------------------------
# 第三步：贝叶斯推理模块
# --------------------------
def calculate_bayesian_prob(conclusion, features):
    """
    计算单个结论的贝叶斯后验概率
    :param conclusion: 候选结论（如支气管炎）
    :param features: 结构化特征字典
    :return: 后验概率（0-1）
    """
    # 1. 获取先验概率
    prior_prob = MEDICAL_KNOWLEDGE_BASE["咳嗽病因"][conclusion]["prior_prob"]
    
    # 2. 特征映射（将用户特征转换为似然概率的键）
    feature_mapping = {
        "咳嗽时长": features["咳嗽时长"] == 14,  # 两周=14天
        "发热": features["发热"] == "否",          # 无发热=否
        "痰颜色": features["痰颜色"] == "白",      # 白痰=白
        "咳嗽时段": features["咳嗽时段"] == "夜间" # 夜间咳嗽=夜间
    }
    
    # 3. 计算似然概率乘积（P(B|A)）
    likelihood_probs = []
    for key, is_match in feature_mapping.items():
        if is_match:
            likelihood = MEDICAL_KNOWLEDGE_BASE["咳嗽病因"][conclusion]["likelihood"][key]
        else:
            likelihood = MEDICAL_KNOWLEDGE_BASE["咳嗽病因"][conclusion]["non_likelihood"][key]
        likelihood_probs.append(likelihood)
    likelihood_product = np.prod(likelihood_probs)
    
    # 4. 计算非似然概率乘积（P(B|¬A)）
    non_likelihood_probs = []
    for key, is_match in feature_mapping.items():
        if is_match:
            non_likelihood = MEDICAL_KNOWLEDGE_BASE["咳嗽病因"][conclusion]["non_likelihood"][key]
        else:
            non_likelihood = MEDICAL_KNOWLEDGE_BASE["咳嗽病因"][conclusion]["likelihood"][key]
        non_likelihood_probs.append(non_likelihood)
    non_likelihood_product = np.prod(non_likelihood_probs)
    
    # 5. 计算边缘概率（P(B)）
    non_prior_prob = 1 - prior_prob
    marginal_prob = (likelihood_product * prior_prob) + (non_likelihood_product * non_prior_prob)
    
    # 6. 计算后验概率（P(A|B)）
    if marginal_prob == 0:
        posterior_prob = 0
    else:
        posterior_prob = (likelihood_product * prior_prob) / marginal_prob
    
    return posterior_prob

def bayesian_inference(conclusions, features):
    """
    对所有候选结论进行贝叶斯计算
    :param conclusions: 候选结论列表
    :param features: 结构化特征字典
    :return: 结论: 置信概率字典（已排序）
    """
    results = {}
    for conclusion in conclusions:
        if conclusion in MEDICAL_KNOWLEDGE_BASE["咳嗽病因"]:
            prob = calculate_bayesian_prob(conclusion, features)
            results[conclusion] = round(prob, 4)  # 保留4位小数
    
    # 按概率降序排序
    sorted_results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
    return sorted_results

# --------------------------
# 第四步：可视化模块
# --------------------------
def plot_confidence_prob(results):
    """
    绘制置信概率柱状图
    :param results: 结论: 置信概率字典
    """
    # 准备数据
    conclusions = list(results.keys())[:5]  # 取前5个
    probs = [results[con] * 100 for con in conclusions]  # 转换为百分比
    
    # 设置中文字体（避免乱码）
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    
    # 绘制柱状图
    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.bar(conclusions, probs, color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd'])
    
    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2., height + 1,
                f'{height:.1f}%', ha='center', va='bottom')
    
    # 设置标题和标签
    ax.set_title('咳嗽病因置信概率分布', fontsize=14)
    ax.set_xlabel('病因', fontsize=12)
    ax.set_ylabel('置信概率（%）', fontsize=12)
    ax.set_ylim(0, 100)
    
    # 保存图片
    plt.tight_layout()
    plt.savefig('confidence_prob.png', dpi=300)
    plt.show()

# --------------------------
# 第五步：主流程
# --------------------------
def main(user_input):
    """
    大模型+贝叶斯推理的主流程
    :param user_input: 用户自然语言输入
    """
    print("===== 1. 提取用户特征 =====")
    features = extract_medical_features(user_input)
    print(f"结构化特征：{json.dumps(features, ensure_ascii=False, indent=2)}")
    
    print("\n===== 2. 生成候选结论 =====")
    conclusions = generate_candidate_conclusions(features)
    print(f"候选结论：{conclusions}")
    
    print("\n===== 3. 贝叶斯计算置信概率 =====")
    results = bayesian_inference(conclusions, features)
    print(f"置信概率结果：{json.dumps(results, ensure_ascii=False, indent=2)}")
    
    print("\n===== 4. 可视化置信概率 =====")
    plot_confidence_prob(results)
    
    print("\n===== 5. 优化输出最终结论 =====")
    final_output = optimize_output(results)
    print(final_output)
    
    return final_output

# --------------------------
# 运行示例
# --------------------------
if __name__ == "__main__":
    # 用户输入示例
    user_input = "我持续咳嗽两周了，没有发热，有白痰，晚上咳嗽特别厉害"
    # 执行主流程
    main(user_input)