import numpy as np
import os

# 学号和姓名
id = "202210423118"
name = "刘国庆"

def classify(input_vector, dataSet, labels, k=3):
    """
    kNN分类算法.

    Args:
    - input_vector: 输入待分类的样本.
    - dataSet: 训练数据集.
    - labels: 训练集对应的标签.
    - k: kNN参数.

    Returns:
    - 预测的分类标签.
    """
    # 计算欧氏距离
    distances = np.linalg.norm(dataSet - input_vector, axis=1)
    # 获取距离最近的k个样本
    nearest_neighbors = np.argsort(distances)[:k]
    # 统计每个标签的投票数
    label_count = {}
    for i in nearest_neighbors:
        vote_label = labels[i]
        label_count[vote_label] = label_count.get(vote_label, 0) + 1
    # 返回投票最多的标签
    return max(label_count, key=label_count.get)

def load_data(directory):
    """
    加载手写数字数据集.

    Args:
    - directory: D:\二值手写数字数据集D:\二值手写数字数据集

    Returns:
    - 特征数据集和标签集.
    """
    labels = []
    file_list = os.listdir(directory)
    num_files = len(file_list)
    dataset = np.zeros((num_files, 1024))  # 每个图像为32x32, 转换为1x1024向量
    
    for i, filename in enumerate(file_list):
        # 标签是文件名的第一个字符
        label = int(filename.split('_')[0])
        labels.append(label)
        # 转换图像为向量
        dataset[i, :] = img_to_vector(os.path.join(directory, filename))
    
    return dataset, np.array(labels)

def img_to_vector(filename):
    """
    将32x32的二值图像转换为1x1024向量.

    Args:
    - filename: 图像文件路径.

    Returns:
    - 向量化的图像数据.
    """
    vector = np.zeros((1, 1024))
    with open(filename) as file:
        for i in range(32):
            line = file.readline().strip()
            for j in range(32):
                vector[0, 32 * i + j] = int(line[j])
    return vector


if __name__ == "__main__":
    # 加载训练数据集
    train_data, train_labels = load_data("trainingDigits")  # 使用正确的函数名 load_data
    
    # 加载测试数据集
    test_data, test_labels = load_data("testDigits")  # 加载测试集文件夹中的所有数据

    # 测试集预测结果
    correct_count = 0
    total_count = len(test_labels)

    for i in range(total_count):
        predicted_label = classify(test_data[i], train_data, train_labels)
        print(f"测试样本 {i+1} - 真实标签: {test_labels[i]}, 预测标签: {predicted_label}")
        if predicted_label == test_labels[i]:
            correct_count += 1

    accuracy = correct_count / total_count
    print(f"分类精度: {accuracy * 100:.2f}%")