import logging
from pinecone import Pinecone, ServerlessSpec
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import numpy as np
from collections import Counter

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化 Pinecone 客户端
pinecone = Pinecone(api_key="51dd9722-cc97-4429-861b-4583d49d6264")
index_name = "mnist-index"

# 检查并删除现有索引
existing_indexes = pinecone.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    pinecone.delete_index(index_name)

# 创建新索引
pinecone.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
logging.info(f"成功创建索引 '{index_name}'。")

# 连接到索引
index = pinecone.Index(index_name)

# 加载 MNIST 数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 拆分数据集，80%用于训练，20%用于测试
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 转换数据为 Pinecone 可接受的格式
vectors = [(str(i), X_train[i].tolist(), {"label": int(y_train[i])}) for i in range(len(X_train))]

# 上传数据到 Pinecone
for i in tqdm(range(0, len(vectors), 1000), desc="Uploading data"):
    batch = vectors[i:i + 1000]
    index.upsert(batch)

logging.info(f"成功上传了 {len(vectors)} 条数据。")

# 测试模型的准确率
correct_predictions = 0

for i in tqdm(range(len(X_test)), desc="Testing"):
    query_data = X_test[i].tolist()
    results = index.query(vector=query_data, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]

    if labels:
        final_prediction = Counter(labels).most_common(1)[0][0]
        if final_prediction == y_test[i]:
            correct_predictions += 1

accuracy = correct_predictions / len(X_test) * 100
logging.info(f"当 k=11 时，使用 Pinecone 的准确率为 {accuracy:.2f}%。")