import os
import numpy as np
import time
import faiss
import dashscope
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from tqdm import tqdm
from modelscope import snapshot_download
import torch
import warnings
warnings.filterwarnings("ignore")

# ====================== 1. 配置项 ======================
# 通义千问API配置（用于效果评估）
DASHSCOPE_API_KEY = os.environ.get('DASHSCOPE_API_KEY')
dashscope.api_key = DASHSCOPE_API_KEY

# 大模型配置
MODEL_NAME = "qwen/Qwen1.5-1.8B-Chat"
CACHE_DIR = "D:\\modelscope\\hub"
print("正在下载/校验模型缓存...")
LOCAL_MODEL_PATH = snapshot_download(MODEL_NAME, cache_dir=CACHE_DIR)
print(f"模型路径: {LOCAL_MODEL_PATH}")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# PSO参数配置
POPULATION_SIZE = 20  # 种群规模
MAX_ITER = 50         # 最大迭代次数
W_MAX = 0.9           # 最大惯性权重
W_MIN = 0.4           # 最小惯性权重
C1 = 2.0              # 个体学习因子
C2 = 2.0              # 群体学习因子

# 参数空间配置（优化的参数及范围）
PARAM_CONFIG = {
    "chunk_size": {"type": "int", "min": 256, "max": 1024, "step": 64},
    "temperature": {"type": "float", "min": 0.1, "max": 1.0, "step": 0.1},
    "batch_size": {"type": "int", "min": 1, "max": 16, "step": 1}
}
PARAM_NAMES = list(PARAM_CONFIG.keys())
PARAM_DIM = len(PARAM_NAMES)

# 评估用的知识库问答数据（示例）
QA_DATA = [
    {"question": "什么是粒子群优化算法？", "answer": "粒子群优化（PSO）是一种基于群体智能的随机优化算法，灵感来源于鸟群觅食行为，通过粒子的位置和速度更新搜索最优解。"},
    {"question": "PSO的核心参数有哪些？", "answer": "PSO的核心参数包括惯性权重、个体学习因子、群体学习因子、种群规模、最大迭代次数等。"},
    {"question": "大模型temperature参数的作用是什么？", "answer": "temperature参数控制大模型生成文本的随机性，值越高生成越随机，值越低越确定。"},
    {"question": "向量库chunk_size参数的作用是什么？", "answer": "chunk_size参数控制文本分段的大小，过小会导致上下文碎片化，过大则检索精度降低。"},
    {"question": "batch_size参数对大模型推理有什么影响？", "answer": "batch_size参数控制批量推理的样本数，过小则推理速度慢，过大则显存占用过高。"},
    # 可扩展更多样本
]
EVAL_SAMPLE_NUM = len(QA_DATA)

# ====================== 2. 工具函数 ======================
def init_pso_population():
    """初始化PSO种群：位置和速度"""
    positions = []
    velocities = []
    
    for _ in range(POPULATION_SIZE):
        pos = []
        vel = []
        for name in PARAM_NAMES:
            cfg = PARAM_CONFIG[name]
            # 初始化位置：参数范围内随机采样
            if cfg["type"] == "int":
                pos_val = np.random.randint(cfg["min"], cfg["max"]+1, size=1)[0]
                # 对齐步长
                pos_val = pos_val - (pos_val - cfg["min"]) % cfg["step"]
            else:
                pos_val = np.random.uniform(cfg["min"], cfg["max"], size=1)[0]
                # 对齐步长
                pos_val = round(pos_val / cfg["step"]) * cfg["step"]
            
            pos.append(pos_val)
            
            # 初始化速度：参数范围的±10%
            v_max = 0.1 * (cfg["max"] - cfg["min"])
            vel_val = np.random.uniform(-v_max, v_max, size=1)[0]
            vel.append(vel_val)
        
        positions.append(np.array(pos))
        velocities.append(np.array(vel))
    
    return np.array(positions), np.array(velocities)

def validate_parameters(params):
    """校验参数有效性，返回修正后的参数"""
    validated = []
    for i, name in enumerate(PARAM_NAMES):
        cfg = PARAM_CONFIG[name]
        val = params[i]
        
        # 类型转换
        if cfg["type"] == "int":
            val = int(round(val))
        
        # 边界限制
        val = max(cfg["min"], min(cfg["max"], val))
        
        # 步长对齐
        if cfg["step"] is not None:
            val = val - (val - cfg["min"]) % cfg["step"]
        
        validated.append(val)
    
    return np.array(validated)

def build_vector_database(chunk_size, chunk_overlap=64):
    """构建向量库（示例）"""
    # 提取知识库文本（这里用QA_DATA的answer）
    texts = [item["answer"] for item in QA_DATA]
    tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_PATH, trust_remote_code=True)
    
    # 文本分段
    chunks = []
    for text in texts:
        tokens = tokenizer.encode(text, add_special_tokens=False)
        for i in range(0, len(tokens), chunk_size - chunk_overlap):
            chunk_tokens = tokens[i:i+chunk_size]
            chunk_text = tokenizer.decode(chunk_tokens)
            chunks.append(chunk_text)
    
    # 生成向量
    model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL_PATH, trust_remote_code=True).to(DEVICE)
    embeddings = []
    for chunk in chunks:
        inputs = tokenizer(chunk, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            embedding = model.model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
        embeddings.append(embedding)
    
    embeddings = np.concatenate(embeddings, axis=0).astype(np.float32)
    
    # 构建FAISS索引
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    
    return index, tokenizer, chunks

def retrieve_context(question, index, tokenizer, chunks, chunk_size):
    """检索上下文（示例）"""
    # 生成问题向量
    model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL_PATH, trust_remote_code=True).to(DEVICE)
    inputs = tokenizer(question, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        q_embedding = model.model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
    
    # 检索Top1
    _, indices = index.search(q_embedding, k=1)
    context = chunks[indices[0][0]] if indices[0][0] < len(chunks) else ""
    
    return context

def generate_answer(question, context, temperature, batch_size):
    """生成回答并记录耗时"""
    # 加载生成管道
    tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_PATH, trust_remote_code=True)
    generator = pipeline(
        "text-generation",
        model=LOCAL_MODEL_PATH,
        tokenizer=tokenizer,
        device=DEVICE,
        batch_size=batch_size
    )
    
    # 构建提示词（Qwen1.5-Chat使用对话格式）
    prompt = f"基于以下上下文回答问题：\n上下文：{context}\n问题：{question}\n回答："
    
    # 生成回答并计时
    start_time = time.time()
    outputs = generator(
        prompt,
        max_new_tokens=100,
        temperature=temperature,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        return_full_text=False
    )
    end_time = time.time()
    
    # 提取生成结果和耗时
    answer = outputs[0]["generated_text"].strip()
    generate_time = end_time - start_time
    
    # 计算生成速度（tokens/s）
    token_num = len(tokenizer.encode(answer))
    generate_speed = token_num / generate_time if generate_time > 0 else 0
    
    return answer, generate_speed

def evaluate_fitness(params):
    """计算适应度值：准确率得分×0.7 + 速度得分×0.3"""
    # 1. 校验参数
    params = validate_parameters(params)
    chunk_size, temperature, batch_size = params
    
    # 2. 构建向量库
    index, tokenizer, chunks = build_vector_database(chunk_size)
    
    # 3. 执行问答任务
    total_accuracy_score = 0
    total_speed = 0
    valid_count = 0
    
    for qa in tqdm(QA_DATA, desc=f"Evaluating params: {params}", leave=False):
        try:
            # 检索上下文
            context = retrieve_context(qa["question"], index, tokenizer, chunks, chunk_size)
            
            # 生成回答
            generated_answer, speed = generate_answer(
                qa["question"], context, temperature, batch_size
            )
            
            # 调用GPT-4评估准确率
            accuracy_score = get_gpt4_evaluation(
                qa["question"], qa["answer"], generated_answer
            )
            
            total_accuracy_score += accuracy_score
            total_speed += speed
            valid_count += 1
        except Exception as e:
            print(f"Error evaluating {qa['question']}: {e}")
            continue
    
    if valid_count == 0:
        return 0.0  # 全部失败，适应度值为0
    
    # 4. 计算平均分和平均速度
    avg_accuracy = total_accuracy_score / valid_count
    avg_speed = total_speed / valid_count
    
    # 5. 归一化得分（0-10）
    accuracy_score = min(10.0, max(0.0, avg_accuracy))
    speed_score = min(10.0, max(0.0, avg_speed * 2))  # 速度×2，让满分对应5 tokens/s
    
    # 6. 计算最终适应度值
    fitness = accuracy_score * 0.7 + speed_score * 0.3
    
    print(f"Params: {params} | Accuracy: {accuracy_score:.2f} | Speed: {speed_score:.2f} | Fitness: {fitness:.2f}")
    
    return fitness

def get_gpt4_evaluation(question, answer, generated_answer):
    """调用通义千问大模型评估回答准确率，返回0-10分"""
    prompt = f"""
请作为专业的评估师，对以下知识库问答结果的准确率进行打分（仅输出0-10的整数分数）：

### 打分规则
1. 生成结果完全包含标准答案的所有核心信息，且无错误信息：10分；
2. 生成结果包含标准答案的核心信息，但缺失1个非核心细节：9分；
3. 生成结果包含标准答案的核心信息，但缺失2个非核心细节：8分；
4. 生成结果包含标准答案的核心信息，但有1个非核心错误：7分；
5. 生成结果包含标准答案的核心信息，但有2个非核心错误：6分；
6. 生成结果缺失1个核心信息，无错误：5分；
7. 生成结果缺失1个核心信息，有1个非核心错误：4分；
8. 生成结果缺失2个核心信息：3分；
9. 生成结果核心信息完全错误：2分；
10. 生成结果与问题无关：0-1分。

### 待评估内容
问题：{question}
标准答案：{answer}
生成结果：{generated_answer}

### 输出要求
仅输出整数分数，不要输出其他内容。
    """
    
    try:
        response = dashscope.Generation.call(
            model='qwen-turbo',
            messages=[{"role": "user", "content": prompt}],
            result_format='message',
            temperature=0.0,  # 评估时温度设为0，保证结果稳定
            max_tokens=10
        )
        # 检查响应是否成功
        if response is None:
            print("Qwen evaluation error: API 返回为空，请检查 DASHSCOPE_API_KEY 是否设置正确")
            return 0
        if response.status_code != 200:
            print(f"Qwen evaluation error: API 错误，code={response.status_code}, message={response.message}")
            return 0
        score = int(response.output.choices[0].message.content.strip())
        return score
    except Exception as e:
        print(f"Qwen evaluation error: {e}")
        return 0  # 评估失败，得0分

def update_pso(positions, velocities, pbest, pbest_fitness, gbest, gbest_fitness, iter_num):
    """更新PSO的速度和位置"""
    # 线性递减惯性权重
    w = W_MAX - (W_MAX - W_MIN) * iter_num / MAX_ITER
    
    new_velocities = []
    new_positions = []
    
    for i in range(POPULATION_SIZE):
        # 速度更新
        r1 = np.random.uniform(0, 1, PARAM_DIM)
        r2 = np.random.uniform(0, 1, PARAM_DIM)
        
        vel_cognitive = C1 * r1 * (pbest[i] - positions[i])
        vel_social = C2 * r2 * (gbest - positions[i])
        new_vel = w * velocities[i] + vel_cognitive + vel_social
        
        # 位置更新
        new_pos = positions[i] + new_vel
        
        # 校验新位置
        new_pos = validate_parameters(new_pos)
        
        new_velocities.append(new_vel)
        new_positions.append(new_pos)
    
    return np.array(new_positions), np.array(new_velocities)

# ====================== 3. 主执行函数 ======================
def main():
    print("=== PSO + 大模型调参开始 ===")
    print(f"PSO参数：种群规模={POPULATION_SIZE}, 最大迭代={MAX_ITER}")
    print(f"优化参数：{PARAM_NAMES}")
    
    # 1. 初始化PSO
    positions, velocities = init_pso_population()
    pbest = positions.copy()  # 个体最优位置
    pbest_fitness = np.array([evaluate_fitness(pos) for pos in positions])  # 个体最优适应度
    gbest_idx = np.argmax(pbest_fitness)  # 全局最优索引
    gbest = pbest[gbest_idx].copy()  # 全局最优位置
    gbest_fitness = pbest_fitness[gbest_idx]  # 全局最优适应度
    
    # 记录每轮最优结果
    history = []
    
    # 2. 迭代优化
    for iter_num in range(MAX_ITER):
        print(f"\n=== 迭代 {iter_num+1}/{MAX_ITER} ===")
        print(f"当前全局最优参数：{gbest} | 适应度值：{gbest_fitness:.2f}")
        
        # 更新速度和位置
        positions, velocities = update_pso(
            positions, velocities, pbest, pbest_fitness, gbest, gbest_fitness, iter_num
        )
        
        # 计算新适应度
        current_fitness = np.array([evaluate_fitness(pos) for pos in positions])
        
        # 更新个体最优
        for i in range(POPULATION_SIZE):
            if current_fitness[i] > pbest_fitness[i]:
                pbest[i] = positions[i].copy()
                pbest_fitness[i] = current_fitness[i]
        
        # 更新全局最优
        current_best_idx = np.argmax(current_fitness)
        current_best_fitness = current_fitness[current_best_idx]
        if current_best_fitness > gbest_fitness:
            gbest = positions[current_best_idx].copy()
            gbest_fitness = current_best_fitness
        
        # 记录历史
        history.append({
            "iter": iter_num+1,
            "gbest": gbest.copy(),
            "gbest_fitness": gbest_fitness
        })
    
    # 3. 输出结果
    print("\n=== PSO + 大模型调参完成 ===")
    print(f"最优参数：")
    for i, name in enumerate(PARAM_NAMES):
        print(f"  {name}: {gbest[i]}")
    print(f"最优适应度值：{gbest_fitness:.2f}")
    
    # 保存历史记录
    np.save("pso_history.npy", history)
    print("历史记录已保存到 pso_history.npy")

if __name__ == "__main__":
    main()