from sklearn.cluster import MiniBatchKMeans
from sklearn.cluster import KMeans
import numpy as np

# 1. 数据集准备
X = np.array([[1, 2], [1, 4], [1, 0],
              [4, 2], [4, 0], [4, 4],
              [4, 5], [0, 1], [2, 2],
              [3, 2], [5, 5], [1, -1]])

print("请输入0或1 (0: MiniBatchKMeans, 1: KMeans):")
n = int(input())

# 根据输入选择模型
if n == 0:
    # MiniBatchKMeans 模块
    # ********** Begin **********#
    
    # 创建 MiniBatchKMeans 模型实例
    # n_clusters=2 表示我们想把数据聚成2类。
    # random_state=0 保证每次运行的随机初始化质心都是一样的，结果可复现。
    # n_init='auto' 会自动选择最佳的初始化次数。
    # batch_size 是每次迭代使用的数据子集大小，这是它与标准 KMeans 的主要区别。
    model = MiniBatchKMeans(n_clusters=2, random_state=0, n_init='auto', batch_size=6)
    
    # ********** End **********#
    print("使用 MiniBatchKMeans")

else:
    # KMeans 模块
    # ********** Begin **********#
    
    # 创建 KMeans 模型实例
    # 参数与 MiniBatchKMeans 类似，但没有 batch_size。
    model = KMeans(n_clusters=2, random_state=0, n_init='auto')
    
    # ********** End **********#
    print("使用 KMeans")

# 对数据进行训练（拟合）
model.fit(X)

# 输出所有点的类别、两类的中心点并预测[0,0],[4,4]的类别
# ********** Begin **********#

# 1. 输出每个数据点所属的类别 (簇)
# model.labels_ 存储了训练后每个点被分配到的簇索引
print("\n数据点的类别标签:")
print(model.labels_)

# 2. 输出两个簇的中心点坐标
# model.cluster_centers_ 存储了每个簇的中心点坐标
print("\n簇的中心点:")
print(model.cluster_centers_)

# 3. 预测新的数据点 [0,0] 和 [4,4] 的类别
new_points = np.array([[0, 0], [4, 4]])
predictions = model.predict(new_points)
print("\n新数据点 [0,0] 和 [4,4] 的预测类别:")
print(predictions)

# ********** End **********#