# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from modelscope import snapshot_download
from datasets import Dataset
import os
import json
from tqdm import tqdm
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial']
plt.rcParams['axes.unicode_minus'] = False
print("=== Qwen1.5-0.5B-Chat医疗专业能力改进 ===\n")
class MedicalModelImprover:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")
        model_name="qwen/Qwen1.5-1.8B-Chat"
        cache_dir = "D:\\modelscope\\hub"
        local_model_path = snapshot_download(model_name, cache_dir=cache_dir)
        # 加载模型和分词器
        self.tokenizer = AutoTokenizer.from_pretrained(local_model_path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        
        print("加载基础模型...")
        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                local_model_path,
                torch_dtype=torch.float16,
                device_map="auto"
            )
            if not hasattr(self, 'model') or self.model is None:
                raise RuntimeError("模型加载失败，请检查模型路径或配置。")
            print("✅ 模型加载成功")
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            raise
        
        # 训练记录
        self.train_losses = []
        
    def create_simulated_medical_data(self):
        """创建模拟医疗数据"""
        print("创建模拟医疗训练数据...")
        
        medical_data = [
            # 糖尿病相关
            {
                "instruction": "糖尿病患者空腹血糖控制目标是多少？",
                "output": "糖尿病患者空腹血糖控制目标为4.4-7.0mmol/L，餐后血糖应低于10.0mmol/L，糖化血红蛋白控制在7.0%以下。具体目标需根据患者年龄、并发症情况个体化调整。"
            },
            {
                "instruction": "胰岛素治疗的起始剂量如何计算？",
                "output": "基础胰岛素起始剂量为0.1-0.2单位/公斤体重，根据空腹血糖水平调整。餐时胰岛素根据碳水化合物摄入量计算，通常每10-15克碳水化合物需要1单位胰岛素。需在医生指导下个体化调整。"
            },
            {
                "instruction": "二甲双胍的用法用量？",
                "output": "二甲双胍起始剂量为500mg每日一次或二次，随餐服用。最大剂量2550mg/日。主要副作用为胃肠道反应，肾功能不全(eGFR<45)患者禁用。用药期间需定期监测肾功能。"
            },
            
            # 高血压相关
            {
                "instruction": "高血压的诊断标准是什么？",
                "output": "高血压诊断标准：非同日3次诊室血压测量，收缩压≥140mmHg和/或舒张压≥90mmHg。24小时动态血压监测平均血压≥130/80mmHg。家庭自测血压≥135/85mmHg应考虑高血压。"
            },
            {
                "instruction": "常用降压药物有哪些类别？",
                "output": "一线降压药物包括：ACEI类(如培哚普利)、ARB类(如缬沙坦)、钙通道阻滞剂(如氨氯地平)、利尿剂(如氢氯噻嗪)。选择需根据患者合并症个体化决定。"
            },
            {
                "instruction": "老年高血压患者的血压控制目标？",
                "output": "65-79岁老年高血压患者血压目标<140/90mmHg，如能耐受可降至<130/80mmHg。≥80岁老年人血压目标<150/90mmHg。降压过程应平稳，避免过快过低。"
            },
            
            # 心脏病相关
            {
                "instruction": "心肌梗死的典型症状有哪些？",
                "output": "心肌梗死典型症状：胸骨后压榨性疼痛，可放射至左肩、下颌、背部，持续20分钟以上，伴出汗、恶心、呼吸困难。不典型表现可表现为牙痛、上腹痛等，老年人及糖尿病患者症状可能不典型。"
            },
            {
                "instruction": "冠心病患者如何进行二级预防？",
                "output": "冠心病二级预防：阿司匹林100mg每日一次，他汀类药物强化降脂(LDL-C<1.8mmol/L)，β受体阻滞剂，ACEI/ARB。同时控制血压<140/90mmHg，血糖达标，戒烟限酒，规律有氧运动。"
            },
            
            # 用药安全相关
            {
                "instruction": "头孢类药物使用注意事项？",
                "output": "头孢类药物使用前需询问青霉素过敏史，用药期间及停药后7天内禁止饮酒，可能发生双硫仑样反应(面部潮红、头痛、呕吐、呼吸困难，严重可致死)。肾功能不全者需调整剂量。"
            },
            {
                "instruction": "华法林的监测指标和目标？",
                "output": "华法林治疗需监测INR(国际标准化比值)，目标范围通常为2.0-3.0，机械瓣膜患者为2.5-3.5。初始治疗需频繁监测，稳定后每4周监测一次。注意多种药物和食物会影响药效。"
            },
            
            # 诊断鉴别相关
            {
                "instruction": "胸痛的鉴别诊断有哪些？",
                "output": "胸痛需鉴别：心源性(心绞痛、心肌梗死、心包炎)、呼吸系统(肺栓塞、肺炎、气胸)、消化系统(胃食管反流、食管痉挛)、 musculoskeletal(肋软骨炎)、焦虑症等。需结合疼痛性质、持续时间、诱发缓解因素判断。"
            },
            {
                "instruction": "腹痛的定位诊断意义？",
                "output": "右上腹痛：肝胆疾病；上腹痛：胃十二指肠、胰腺；右下腹痛：阑尾炎；左下腹痛：降结肠、妇科疾病；弥漫性腹痛：肠梗阻、腹膜炎。结合体征、实验室检查和影像学综合判断。"
            },
            
            # 检验指标相关
            {
                "instruction": "肝功能检查主要指标和意义？",
                "output": "ALT、AST反映肝细胞损伤，正常值<40U/L；ALP、GGT反映胆道疾病；总胆红素正常值<21μmol/L；白蛋白反映肝脏合成功能，正常值35-55g/L。异常需结合临床表现和影像学检查。"
            },
            {
                "instruction": "肾功能评估指标有哪些？",
                "output": "血肌酐正常值44-133μmol/L，eGFR(估算肾小球滤过率)正常值>90ml/min/1.73m²。尿蛋白/肌酐比值<30mg/g为正常。慢性肾病分期基于eGFR和尿蛋白水平。"
            },
            
            # 急诊处理相关
            {
                "instruction": "高热惊厥的急诊处理？",
                "output": "高热惊厥急诊处理：保持呼吸道通畅，侧卧位防止误吸，地西泮0.3-0.5mg/kg静脉或直肠给药止惊，物理降温，对乙酰氨基酚10-15mg/kg退热。查找感染源，严重或持续惊厥需住院治疗。"
            },
            {
                "instruction": "过敏性休克的急救措施？",
                "output": "过敏性休克急救：立即肌注肾上腺素0.3-0.5mg，保持呼吸道通畅，吸氧，建立静脉通道快速补液，糖皮质激素和抗组胺药物辅助治疗。密切监测生命体征，必要时气管插管。"
            }
        ]
        
        print(f"创建了 {len(medical_data)} 条医疗训练数据")
        return medical_data
    
    def create_precision_test_data(self):
        """创建精度测试数据"""
        test_cases = [
            {
                "category": "数值精度",
                "question": "正常成人的血压范围是多少？",
                "expected_keywords": ["120", "80", "mmHg", "139", "89"],
                "description": "测试具体数值记忆"
            },
            {
                "category": "专业术语", 
                "question": "ACEI类药物的作用机制是什么？",
                "expected_keywords": ["血管紧张素", "转换酶", "抑制", "醛固酮"],
                "description": "测试专业术语理解"
            },
            {
                "category": "安全警示",
                "question": "头孢类药物可以和酒精一起使用吗？",
                "expected_keywords": ["禁止", "双硫仑样反应", "7天", "致死"],
                "description": "测试安全警告"
            },
            {
                "category": "诊断思路",
                "question": "右下腹痛可能是什么疾病？",
                "expected_keywords": ["阑尾炎", "妇科疾病", "泌尿系", "鉴别"],
                "description": "测试诊断思维"
            },
            {
                "category": "用药指导",
                "question": "二甲双胍的禁忌症有哪些？",
                "expected_keywords": ["肾功能", "eGFR", "禁忌", "乳酸酸中毒"],
                "description": "测试用药知识"
            }
        ]
        return test_cases
    
    def evaluate_model(self, test_cases):
        """评估模型表现"""
        print("\n🔍 开始模型评估...")
        self.model.eval()
        
        results = []
        
        for i, test in enumerate(test_cases, 1):
            prompt = f"你是一名专业的医疗AI助手，请根据以下医学问题提供准确、专业的回答。\n医学问题: {test['question']}\n医学回答要求:\n1. 回答必须包含所有关键医学术语和数值。\n2. 回答必须清晰、准确，避免模糊描述。\n3. 如果问题涉及禁忌或安全警示，必须明确提示。\n医学回答:"
            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(self.device)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=600,
                    temperature=0.8,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id,
                    top_k=80,
                    top_p=0.85,
                    num_beams=4,
                    early_stopping=True,
                    repetition_penalty=1.2
                )
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"模型生成结果: {response}")  # 调试输出
            answer = response.split("医学回答:")[-1].strip()
            if not answer:
                answer = "模型未能生成有效回答，请检查输入或模型状态。"
            
            # 计算关键词匹配度
            matched_keywords = []
            missing_keywords = []
            
            for keyword in test['expected_keywords']:
                if keyword in answer:
                    matched_keywords.append(keyword)
                else:
                    missing_keywords.append(keyword)
            
            score = len(matched_keywords) / len(test['expected_keywords']) * 100
            
            print(f"\n{i}. [{test['category']}] {test['question']}")
            print(f"   回答: {answer}")
            print(f"   得分: {score:.0f}%")
            print(f"   匹配: {matched_keywords}")
            print(f"   缺失: {missing_keywords}")
            
            results.append({
                'category': test['category'],
                'score': score,
                'answer': answer,
                'matched': matched_keywords,
                'missing': missing_keywords
            })
        
        return results
    
    def tokenize_function(self, examples):
        """分词函数"""
        prompts = [f"医学问题: {q}\n医学回答:" for q in examples['instruction']]
        answers = examples['output']
        
        # 编码输入和目标
        model_inputs = self.tokenizer(
            prompts,
            text_target=answers,
            max_length=384,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        return model_inputs
    
    def train_model(self, training_data, epochs=3):
        """训练模型"""
        print(f"\n🎯 开始模型训练 ({epochs}个epochs)...")
        
        # 转换为dataset格式
        train_dataset = Dataset.from_list(training_data)
        tokenized_dataset = train_dataset.map(self.tokenize_function, batched=True)
        
        # 训练参数
        training_args = TrainingArguments(
            output_dir="./medical_model_checkpoints",
            per_device_train_batch_size=2,
            num_train_epochs=epochs,
            learning_rate=2e-5,
            warmup_steps=100,
            logging_steps=50,
            save_steps=200,
            save_total_limit=2,
            prediction_loss_only=True,
            remove_unused_columns=False,
            fp16=True,
        )
        
        # 自定义训练器以记录损失
        class CustomTrainer(Trainer):
            def __init__(self, *args, **kwargs):
                self.outer_class = kwargs.pop('outer_class')
                super().__init__(*args, **kwargs)
            
            def log(self, logs):
                super().log(logs)
                if 'loss' in logs:
                    self.outer_class.train_losses.append(logs['loss'])
        
        trainer = CustomTrainer(
            model=self.model,
            args=training_args,
            train_dataset=tokenized_dataset,
            outer_class=self
        )
        
        # 开始训练
        trainer.train()
        
        print("✅ 训练完成")
        return trainer
    
    def plot_training_progress(self):
        """绘制训练进度图"""
        if not self.train_losses:
            print("没有训练损失数据可绘制")
            return
        
        plt.figure(figsize=(12, 4))
        
        # 损失曲线
        plt.subplot(1, 2, 1)
        plt.plot(self.train_losses, 'b-', alpha=0.7, linewidth=1)
        plt.title('训练损失曲线')
        plt.xlabel('训练步数')
        plt.ylabel('损失值')
        plt.grid(True, alpha=0.3)
        
        # 移动平均
        plt.subplot(1, 2, 2)
        if len(self.train_losses) > 10:
            window = 10
            moving_avg = np.convolve(self.train_losses, np.ones(window)/window, mode='valid')
            plt.plot(range(window-1, len(self.train_losses)), moving_avg, 'r-', linewidth=2)
            plt.title('损失移动平均')
            plt.xlabel('训练步数')
            plt.ylabel('损失值')
            plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('./training_progress.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def compare_performance(self, initial_results, final_results):
        """对比改进效果"""
        print("\n📊 性能改进对比报告")
        print("=" * 60)
        
        initial_scores = {r['category']: r['score'] for r in initial_results}
        final_scores = {r['category']: r['score'] for r in final_results}
        
        categories = set(initial_scores.keys()) | set(final_scores.keys())
        
        improvement_data = []
        
        for category in categories:
            initial = initial_scores.get(category, 0)
            final = final_scores.get(category, 0)
            improvement = final - initial
            
            improvement_data.append({
                'category': category,
                'initial': initial,
                'final': final,
                'improvement': improvement
            })
            
            print(f"{category:12} | {initial:5.1f}% → {final:5.1f}% | 提升: {improvement:+.1f}%")
        
        # 绘制对比图
        self.plot_comparison(improvement_data)
        
        return improvement_data
    
    def plot_comparison(self, improvement_data):
        """绘制对比图"""
        categories = [item['category'] for item in improvement_data]
        initial_scores = [item['initial'] for item in improvement_data]
        final_scores = [item['final'] for item in improvement_data]
        
        x = np.arange(len(categories))
        width = 0.35
        
        plt.figure(figsize=(12, 6))
        
        plt.bar(x - width/2, initial_scores, width, label='改进前', alpha=0.7, color='red')
        plt.bar(x + width/2, final_scores, width, label='改进后', alpha=0.7, color='green')
        
        plt.xlabel('能力类别')
        plt.ylabel('得分 (%)')
        plt.title('医疗专业能力改进对比')
        plt.xticks(x, categories, rotation=45)
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        # 添加数值标签
        for i, (init, final) in enumerate(zip(initial_scores, final_scores)):
            plt.text(i - width/2, init + 1, f'{init:.0f}%', ha='center', va='bottom')
            plt.text(i + width/2, final + 1, f'{final:.0f}%', ha='center', va='bottom')
            improvement = final - init
            plt.text(i, max(init, final) + 5, f'+{improvement:.0f}%', 
                    ha='center', va='bottom', fontweight='bold', color='blue')
        
        plt.tight_layout()
        plt.savefig('./improvement_comparison.png', dpi=300, bbox_inches='tight')
        plt.show()
    
    def save_model(self, output_dir="./improved_medical_model"):
        """保存改进后的模型"""
        print(f"\n💾 保存模型到: {output_dir}")
        os.makedirs(output_dir, exist_ok=True)
        
        self.model.save_pretrained(output_dir)
        self.tokenizer.save_pretrained(output_dir)
        
        # 保存训练信息
        training_info = {
            'base_model': 'Qwen/Qwen1.5-0.5B-Chat',
            'training_data_size': len(self.train_losses) * 2,  # 估算
            'final_loss': self.train_losses[-1] if self.train_losses else None,
            'training_epochs': 3,
            'improvement_details': '医疗专业能力增强训练'
        }
        
        with open(os.path.join(output_dir, "training_info.json"), 'w', encoding='utf-8') as f:
            json.dump(training_info, f, ensure_ascii=False, indent=2)
        
        print("✅ 模型保存完成")
    
    def demo_improved_capability(self):
        """展示改进后的能力"""
        demo_questions = [
            "糖尿病患者应该如何制定饮食计划？",
            "高血压急症如何处理？",
            "胸痛患者需要做哪些检查？"
        ]
        
        print("\n🎯 改进后能力演示")
        print("=" * 50)
        
        self.model.eval()
        
        for i, question in enumerate(demo_questions, 1):
            prompt = f"医学问题: {question}\n医学回答:"
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=250,
                    temperature=0.7,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            answer = response.split("医学回答:")[-1].strip()
            
            print(f"\n{i}. 问题: {question}")
            print(f"   回答: {answer}")
            print("-" * 80)
def main():
    """主函数"""
    try:
        print("🚀 开始Qwen1.5-0.5B-Chat医疗专业能力改进流程")
        
        # 初始化改进器
        improver = MedicalModelImprover()
        
        # 1. 创建训练数据
        print("\n" + "="*50)
        print("步骤1: 准备训练数据")
        training_data = improver.create_simulated_medical_data()
        
        # 2. 初始评估
        print("\n" + "="*50)
        print("步骤2: 初始能力评估")
        test_cases = improver.create_precision_test_data()
        initial_results = improver.evaluate_model(test_cases)
        initial_avg_score = np.mean([r['score'] for r in initial_results])
        print(f"\n📈 初始平均得分: {initial_avg_score:.1f}%")
        
        # 3. 模型训练
        print("\n" + "="*50)
        print("步骤3: 专业能力训练")
        improver.train_model(training_data, epochs=3)
        
        # 4. 训练进度可视化
        improver.plot_training_progress()
        
        # 5. 最终评估
        print("\n" + "="*50)
        print("步骤4: 改进后评估")
        final_results = improver.evaluate_model(test_cases)
        final_avg_score = np.mean([r['score'] for r in final_results])
        print(f"\n📈 最终平均得分: {final_avg_score:.1f}%")
        
        # 6. 性能对比
        print("\n" + "="*50)
        print("步骤5: 性能改进分析")
        improvement_data = improver.compare_performance(initial_results, final_results)
        
        # 7. 能力演示
        improver.demo_improved_capability()
        
        # 8. 保存模型
        print("\n" + "="*50)
        print("步骤6: 保存改进模型")
        improver.save_model()
        
        # 总结报告
        print("\n🎉 改进完成总结")
        print("=" * 50)
        print(f"初始平均得分: {initial_avg_score:.1f}%")
        print(f"最终平均得分: {final_avg_score:.1f}%")
        print(f"总体提升: {final_avg_score - initial_avg_score:+.1f}%")
        print(f"提升比例: {(final_avg_score - initial_avg_score) / initial_avg_score * 100:.1f}%")
        
        return improver.model, improver.tokenizer
        
    except Exception as e:
        print(f"❌ 错误: {e}")
        import traceback
        traceback.print_exc()
if __name__ == "__main__":
    result = main()
    if result is not None:
        improved_model, improved_tokenizer = result
    else:
        print("❌ 主函数未返回有效结果，请检查错误日志。")