import pymysql
import numpy as np
import matplotlib.pyplot as plt
from sentence_transformers import SentenceTransformer
from datetime import date
from typing import List, Tuple
# ===================== 1. 数据库配置与连接 =====================
DB_CONFIG = {
    "host": "localhost", 
    "user": "root",
    "password": "**********",
    "database": "数据库名称",
    "charset": "utf8mb4"
}
# 初始化Embedding模型（生成384维向量，降低计算成本）
model = SentenceTransformer('D:/modelscope/hub/models/sentence-transformers/all-MiniLM-L6-v2')
VECTOR_DIM = 384  # 向量维度
# ===================== 2. 工具函数 =====================
def get_mysql_connection():
    """创建并返回MySQL连接"""
    conn = pymysql.connect(**DB_CONFIG)
    conn.autocommit(False)
    return conn
def init_database():
    """初始化数据库表结构（兼容低版本MySQL）"""
    conn = get_mysql_connection()
    cursor = conn.cursor()
    
    # 1. 创建表：向量存储为字符串（逗号分隔的float值）
    cursor.execute("""
        DROP TABLE IF EXISTS tech_docs;
    """)
    cursor.execute("""                   
        CREATE TABLE tech_docs (
            id INT AUTO_INCREMENT PRIMARY KEY,
            title VARCHAR(255) NOT NULL,
            content TEXT NOT NULL,
            publish_date DATE NOT NULL,
            tag VARCHAR(50) NOT NULL,
            embedding TEXT NOT NULL  -- 向量存储为字符串："x1,x2,...x384"
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
    """)
    
    # 2. 创建结构化字段索引（向量无原生索引，用暴力检索）
    cursor.execute("""
        CREATE INDEX idx_tech_docs_publish_date ON tech_docs(publish_date);
    """)
    cursor.execute("""
        CREATE INDEX idx_tech_docs_tag ON tech_docs(tag);
    """)
    
    # 3. 创建自定义余弦相似度计算函数（核心兼容逻辑）
    cursor.execute("""
        DROP FUNCTION IF EXISTS calc_cosine_similarity;
    """)
    cursor.execute("""
        CREATE FUNCTION calc_cosine_similarity(vec1 TEXT, vec2 TEXT) 
        RETURNS FLOAT
        DETERMINISTIC
        BEGIN
            DECLARE dot_product FLOAT DEFAULT 0.0;
            DECLARE norm1 FLOAT DEFAULT 0.0;
            DECLARE norm2 FLOAT DEFAULT 0.0;
            DECLARE i INT DEFAULT 1;
            DECLARE v1 FLOAT;
            DECLARE v2 FLOAT;
            
            -- 拆分向量为数组并计算点积、模长
            WHILE i <= %d DO
                SET v1 = CAST(SUBSTRING_INDEX(SUBSTRING_INDEX(vec1, ',', i), ',', -1) AS FLOAT);
                SET v2 = CAST(SUBSTRING_INDEX(SUBSTRING_INDEX(vec2, ',', i), ',', -1) AS FLOAT);
                
                SET dot_product = dot_product + (v1 * v2);
                SET norm1 = norm1 + (v1 * v1);
                SET norm2 = norm2 + (v2 * v2);
                
                SET i = i + 1;
            END WHILE;
            
            -- 计算余弦相似度（避免除以0）
            IF norm1 = 0 OR norm2 = 0 THEN
                RETURN 0.0;
            ELSE
                RETURN dot_product / (SQRT(norm1) * SQRT(norm2));
            END IF;
        END;
    """ % VECTOR_DIM)
    
    conn.commit()
    cursor.close()
    conn.close()
    print("数据库表、索引、自定义函数初始化完成！")

def embedding_to_str(embedding: List[float]) -> str:
    """将向量数组转为逗号分隔的字符串"""
    return ','.join(map(str, embedding))
def str_to_embedding(embedding_str: str) -> np.ndarray:
    """将字符串转回向量数组"""
    return np.array([float(x) for x in embedding_str.split(',')])
def insert_docs(docs: List[dict]):
    """插入文档数据（生成向量并入库）"""
    conn = get_mysql_connection()
    cursor = conn.cursor()
    
    for doc in docs:
        # 生成384维向量并转为字符串
        embedding = model.encode(doc["content"], normalize_embeddings=True).tolist()  # 归一化提升相似度计算精度
        embedding_str = embedding_to_str(embedding)
        
        # 插入数据
        cursor.execute("""
            INSERT INTO tech_docs (title, content, publish_date, tag, embedding)
            VALUES (%s, %s, %s, %s, %s)
        """, (doc["title"], doc["content"], doc["publish_date"], doc["tag"], embedding_str))
    
    conn.commit()
    cursor.close()
    conn.close()
    print(f"成功插入{len(docs)}条文档数据！")

def hybrid_query_filter_first(query_text: str, publish_date_after: date, top_k: int = 2) -> List[Tuple]:
    """过滤优先：先按发布时间过滤，再向量检索"""
    # 生成查询向量并转为字符串
    query_embedding = model.encode(query_text, normalize_embeddings=True).tolist()
    query_embedding_str = embedding_to_str(query_embedding)
    
    conn = get_mysql_connection()
    cursor = conn.cursor()
    
    # 使用自定义函数计算余弦相似度
    cursor.execute("""
        SELECT 
            title, content, publish_date, tag,
            calc_cosine_similarity(embedding, %s) AS similarity,
            1 - calc_cosine_similarity(embedding, %s) AS distance
        FROM tech_docs
        WHERE publish_date >= %s
        ORDER BY similarity DESC  -- 相似度越高越靠前
        LIMIT %s
    """, (query_embedding_str, query_embedding_str, publish_date_after, top_k))
    
    results = cursor.fetchall()
    cursor.close()
    conn.close()
    return results
def hybrid_query_search_first(query_text: str, target_tag: str, top_k: int = 3) -> List[Tuple]:
    """检索优先：先向量检索Top-K，再过滤标签"""
    # 生成查询向量并转为字符串
    query_embedding = model.encode(query_text, normalize_embeddings=True).tolist()
    query_embedding_str = embedding_to_str(query_embedding)
    
    conn = get_mysql_connection()
    cursor = conn.cursor()
    
    # 先暴力检索Top-K（按相似度排序）
    cursor.execute("""
        SELECT 
            title, content, publish_date, tag,
            calc_cosine_similarity(embedding, %s) AS similarity,
            1 - calc_cosine_similarity(embedding, %s) AS distance
        FROM tech_docs
        ORDER BY similarity DESC
        LIMIT %s
    """, (query_embedding_str, query_embedding_str, top_k))
    
    raw_results = cursor.fetchall()
    cursor.close()
    conn.close()
    
    # 过滤标签
    filtered_results = [r for r in raw_results if r[3] == target_tag]
    return filtered_results

def plot_similarity(results: List[Tuple], title: str):
    """绘制相似度分布柱状图（图片输出）"""
    if not results:
        print("无数据可绘制图表！")
        return
    
    # 提取数据
    doc_titles = [r[0][:10] + "..." for r in results]  # 标题截断
    similarities = [round(r[4], 4) for r in results]   # 相似度
    distances = [round(r[5], 4) for r in results]      # 距离
    
    # 设置中文字体
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    
    # 创建子图
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # 子图1：余弦相似度
    ax1.bar(doc_titles, similarities, color="#34A853", alpha=0.8)
    ax1.set_title(f"{title} - 余弦相似度")
    ax1.set_xlabel("文档标题")
    ax1.set_ylabel("相似度（越大越相似）")
    ax1.grid(axis="y", linestyle="--", alpha=0.7)
    for i, v in enumerate(similarities):
        ax1.text(i, v + 0.01, f"{v}", ha="center")
    
    # 子图2：余弦距离
    ax2.bar(doc_titles, distances, color="#4285F4", alpha=0.8)
    ax2.set_title(f"{title} - 余弦距离")
    ax2.set_xlabel("文档标题")
    ax2.set_ylabel("距离值（越小越相似）")
    ax2.grid(axis="y", linestyle="--", alpha=0.7)
    for i, v in enumerate(distances):
        ax2.text(i, v + 0.01, f"{v}", ha="center")
    
    plt.tight_layout()
    plt.savefig(f"{title.replace(' ', '_')}.png", dpi=300, bbox_inches="tight")
    plt.show()
    print(f"相似度图表已保存为：{title.replace(' ', '_')}.png")
def print_formatted_results(results: List[Tuple], scene_name: str):
    """格式化输出检索结果"""
    print(f"\n===== {scene_name} 查询结果 =====")
    if not results:
        print("无符合条件的结果！")
        return
    
    for i, res in enumerate(results):
        title, content, publish_date, tag, similarity, distance = res
        print(f"\n【结果{i+1}】")
        print(f"标题：{title}")
        print(f"标签：{tag}")
        print(f"发布时间：{publish_date}")
        print(f"余弦相似度：{similarity:.4f}（越大越相似）")
        print(f"余弦距离：{distance:.4f}（越小越相似）")
        print(f"内容：{content[:100]}...")

# ===================== 3. 主流程执行 =====================
if __name__ == "__main__":
    # 步骤1：初始化数据库
    init_database()
    
    # 步骤2：准备测试数据
    test_docs = [
        {
            "title": "RAG技术核心原理与实践",
            "content": "RAG技术通过检索外部知识增强大模型生成能力，解决幻觉问题，是当前大模型落地的核心技术之一",
            "publish_date": date(2024, 5, 10),
            "tag": "RAG"
        },
        {
            "title": "MySQL向量检索兼容方案",
            "content": "低版本MySQL可通过字符串存储向量，自定义函数计算余弦相似度，实现基础的向量检索能力",
            "publish_date": date(2024, 6, 15),
            "tag": "向量数据库"
        },
        {
            "title": "PostgreSQL pgvector插件实战",
            "content": "pgvector为PostgreSQL提供向量存储和检索能力，支持余弦、欧氏等多种距离计算方式，适配RAG场景",
            "publish_date": date(2024, 7, 20),
            "tag": "向量数据库"
        },
        {
            "title": "大模型微调技术LoRA实战",
            "content": "LoRA通过低秩适配实现高效微调，无需训练整个模型参数，大幅降低大模型微调的硬件成本",
            "publish_date": date(2023, 12, 5),
            "tag": "大模型"
        }
    ]
    
    # 步骤3：插入测试数据
    insert_docs(test_docs)
    
    # 步骤4：场景A - 过滤优先查询
    query_a = "MySQL向量检索"
    publish_date_a = date(2024, 1, 1)
    results_a = hybrid_query_filter_first(query_a, publish_date_a, top_k=2)
    print_formatted_results(results_a, "场景A（过滤优先）")
    plot_similarity(results_a, "场景A - MySQL向量检索文档相似度分布")
    
    # 步骤5：场景B - 检索优先查询
    query_b = "大模型应用"
    target_tag_b = "RAG"
    results_b = hybrid_query_search_first(query_b, target_tag_b, top_k=3)
    print_formatted_results(results_b, "场景B（检索优先）")
    plot_similarity(results_b, "场景B - 大模型应用-RAG文档相似度分布")