# 导入必要的库
import matplotlib.pyplot as plt  # 用于绘图
import numpy as np  # 用于数值计算
from collections import Counter  # 用于计数
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits  # 用于加载 MNIST 数据集

# 初始化 Pinecone
pinecone = Pinecone(api_key="5f01d48c-9244-4f19-b1a6-1008b8728b56")

# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pinecone.list_indexes()
'''
# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")
'''

# 连接到索引
index = pinecone.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")

# 使用 load_digits 函数加载 MNIST 数据集
digits = load_digits(n_class=10)

# 获取数据集中的特征数据和标签
X = digits.data
y = digits.target

# 初始化一个空列表,用于存储转换后的向量数据
vectors = []

# 遍历所有样本,将数据转换为 Pinecone 可接受的格式
for i in range(len(X)):
    vector_id = str(i)  # 使用样本的索引作为向量的唯一标识符
    vector_values = X[i].tolist()  # 将 NumPy 数组转换为 Python 列表
    metadata = {"label": int(y[i])}  # 创建元数据字典,包含该样本的真实标签
    vectors.append((vector_id, vector_values, metadata))  # 添加到 vectors 列表中

# 定义批处理大小,每批最多包含 1000 个向量
batch_size = 1000

# 分批上传数据到 Pinecone 索引中
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)

# 检查索引中的向量数量
index_stats = index.describe_index_stats()
print(f"Number of vectors in index: {index_stats['total_vector_count']}")

# 创建一个手写数字 3 的图像（8x8）
digit_3 = np.array(
    [[0, 0, 255, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 255, 255, 255, 255, 0, 0]]
)

# 将图像像素值从 [0-255] 缩放到 [0-16]
digit_3_flatten = (digit_3 / 255.0) * 16

# 将二维图像数组展平成一维列表
query_data = digit_3_flatten.ravel().tolist()

# 使用准备好的查询向量在 Pinecone 索引中执行搜索
results = index.query(
    vector=query_data,
    top_k=5,
    include_metadata=True
)

# 检查是否有匹配项并处理结果
if results['matches']:
    labels = [match['metadata']['label'] for match in results['matches']]
    
    # 打印每个匹配结果的详细信息
    for match in results['matches']:
        print(f"id: {match['id']}, distance: {match['score']}, label: {match['metadata']['label']}")

    # 使用投票机制确定最终的分类结果
    final_prediction = Counter(labels).most_common(1)[0][0]
else:
    print("没有找到匹配项。")
    final_prediction = None

# 使用 matplotlib 显示查询图像和预测结果
plt.imshow(digit_3, cmap='gray')  
if final_prediction is not None:
    plt.title(f"Predicted digit: {final_prediction}", size=15)
else:
    plt.title("No prediction available", size=15)
plt.axis('off')  
plt.show()