import os
import numpy as np
import gradio as gr
from PIL import Image
from pinecone import Pinecone

# 设置 Pinecone API 密钥
api_key = "0faf78dd-2331-43ae-b772-0f2b5986217d"

# 创建 Pinecone 实例
pc = Pinecone(api_key=api_key)

# 设置 Pinecone 索引
index_name = "mnist-index"
index = pc.Index(index_name)

def preprocess_image(image):
    if image is None:
        print("Error: 输入图像为空")
        return None
    
    print(f"Input image type: {type(image)}")
    print(f"Input image shape: {image.shape if isinstance(image, np.ndarray) else 'N/A'}")
    
    try:
        # Ensure the input is a numpy array
        if not isinstance(image, np.ndarray):
            raise ValueError("输入不是NumPy数组")

        # If the image is RGB, convert to grayscale
        if image.ndim == 3:
            img_array = np.mean(image, axis=2).astype(np.uint8)
        elif image.ndim == 2:
            img_array = image
        else:
            raise ValueError(f"意外的图像形状: {image.shape}")

        print(f"处理后的图像形状: {img_array.shape}")
        
        # Convert to PIL Image for resizing
        img = Image.fromarray(img_array)
        
        # Resize to 8x8
        img = img.resize((8, 8), Image.LANCZOS)
        
        # Convert back to numpy array and normalize
        img_array = np.array(img)
        img_array = (img_array / 255.0) * 16
        
        return img_array.flatten()
    except Exception as e:
        print(f"图像预处理中的错误: {e}")
        return None

def predict(image):
    processed_image = preprocess_image(image)
    if processed_image is None:
        return "图像处理失败"
    
    # 使用 Pinecone 查询
    try:
        query_response = index.query(vector=processed_image.tolist(), top_k=11, include_metadata=True)
        neighbors = query_response['matches']
        
        if not neighbors:
            return "未找到邻居"
        
        # 获取最近的k个邻居的标签
        neighbor_labels = [match['metadata']['label'] for match in neighbors]
        
        # 使用多数投票法确定最终标签
        predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
        
        return int(predicted_label)
    except Exception as e:
        return f"预测过程中的错误：{e}"

# 创建Gradio接口
iface = gr.Interface(
    fn=predict,
    inputs=gr.Sketchpad(),
    outputs=gr.Label(label="预测结果"),
    live=False,
    title="手写数字识别",
    description="请在下方的画板上绘制一个手写数字（0-9）"
)

# 启动接口
iface.launch(share=True)