import matplotlib.pyplot as plt

import numpy as np

def ls(p):

    # 返回左子节点

    return p << 1

def rs(p):

    # 返回右子节点

    return p << 1 | 1

def build_kd_tree(data_x, tree, p, dim):

    """

    建立二叉树的过程

    :param data: 建立二叉树所需要的数据

    :param tree: 存二叉树的数组

    :param p: 所在的节点

    :param dim: 现在所在的维度

    :return: None

    """

    # 根据dim对数据进行排序并取其中位数

    # data[data[:,i].argsort()],根据第i维对数据进行排序

    if len(data_x) == 1:

        tree[p] = data_x

        return

    if len(data_x) == 0:

        return

    length = len(data_x[0]) - 1

    data_x = data_x[data_x[:, dim].argsort()]

    mid = len(data_x) >> 1

    tree[p] = data_x[mid]

    build_kd_tree(data_x[:mid, :], tree, ls(p), ((dim + 1) % length))

    build_kd_tree(data_x[mid + 1:, :], tree, rs(p), ((dim + 1) % length))

    return

def find_leaf_node(data, tree):

    # 从根节点出发，循环向下访问kd树，返回其叶子节点

    p, dim, length = 1, 0, len(tree[1]) - 1

    if tree[p, dim].sum() <= 0.0:

        return 1

    while True:

        if data[dim] > tree[p, dim]:

            if tree[rs(p), dim].sum() <= 0.0:

                return p

            p = rs(p)

        else:

            if tree[ls(p), dim].sum() <= 0.0:

                return p

            p = ls(p)

        dim = (dim + 1) % length

    return 1

def distance(a, b, p=1):

    # 我只对数据进行p方运算，不进行开方运算

    sum = 0

    if p == 1:

        c = a - b

    else:

        c = a - b

        c=c.__pow__(p)

    c=c.__abs__()

    sum=c.sum()

    return sum

def find_label(data,tree,nowp,mer):

    if tree[nowp].sum()<=0.0:

        return

    len_of_mer=len(mer)

    len_of_data=len(data)

    mer=mer[mer[:,0].argsort()]

    dis=distance(data,tree[nowp,:-1])

    if dis<=mer[0,0]:

        for i in range(1,len_of_mer-2):

            mer[i]=mer[i+1]

        mer[0]=[dis,tree[nowp,-1]]

        find_label(data,tree,ls(nowp),mer)

        find_label(data,tree,rs(nowp),mer)

    for i in range(1,len_of_mer):

        if dis>mer[len_of_mer-i,0]:

            if i!=1:

                mer[len_of_mer-1]=[dis,tree[nowp,-1]]

            find_label(data, tree, nowp >> 1, mer)

            return

def k_NN(data, tree, p, k):

    label = np.zeros((len(data), 1))

    for i in range(len(data)):

        # 先找到其对应的叶子节点

        #直接默认p=1吧，p对这个影响不是很大吧，主要是k的影响

        pointer = find_leaf_node(data[i], tree)

        mer=np.zeros((k,2))

        for j in range(k):

            mer[j,0]=9999999

        find_label(data[i],tree,pointer,mer)

        # d=np.argmax(np.bincount(int(mer[:,1])))

        # label[i]=d

    return  label

if __name__ == '__main__':

    # train_x, train_y = np.load("data//train_x.npy"), np.load("data//train_y.npy")

    # test_x, test_y = np.load("data//test_x.npy"), np.load("data//test_y.npy")

    # kd树的建立过程，即为二叉树的建立过程，只需要将二叉树对位置的划分变成对某个维度的排序再取其中位数位置的作为划分中点即可

    # 下面为验证数据集，可以看到完美符合

    train_x = np.array(((2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)))

    train_y = np.array((4, 2, 3, 5, 6, 1))

    # 为了防止操作过于繁琐，将标签直接加到数据的最后一列

    test_x=np.array(((1,2),(3,4)))

    data = np.insert(train_x, len(train_x[0]), train_y, axis=1)

    tree = np.zeros(((len(train_x) << 2) + 10, len(data[0])))

    build_kd_tree(data, tree, 1, 0)

    p, k = 1, 4

    test_label = k_NN(test_x, tree, p, k)

    print(test_label)