import numpy as np
from time import time
from scipy.spatial import KDTree as kd
from functools import reduce
import matplotlib.pyplot as plt

def euclid(c, cs, r):
    return ((cs[:,0] - c[0]) ** 2 + (cs[:,1] - c[1]) ** 2 + (cs[:,2] - c[2]) ** 2) < r ** 2

def find_nn_naive(cells, radius):
    matrix = np.zeros((cells.shape[0] ** 2,2))
    mi = 0
    for i in range(len(cells)):
        cell = cells[i]
        cands = euclid(cell, cells, radius)
        ids = np.where(cands)[0]
        ids = ids[ids != i]
        for j in range(len(ids)):
            matrix[mi + j, 0] = i
            matrix[mi + j, 1] = ids[j]
        mi += len(ids)
    return matrix[0:mi,:]

def find_nn_kd_seminaive(cells, radius):
    tree = kd(cells)
    matrix = np.zeros((cells.shape[0] ** 2,2))
    mi = 0
    for i in range(len(cells)):
        res = tree.query_ball_point(cells[i], radius)
        if len(res) > 1:
            ids = res[1:]
            for j in range(len(ids)):
                matrix[mi + j, 0] = i
                matrix[mi + j, 1] = ids[j]
            mi += len(ids)
    return matrix[0:mi,:]

def find_nn_kd_by_tree(cells, radius):
    tree = kd(cells)
    res = list(filter(lambda x: len(x) > 1, tree.query_ball_tree(tree, radius)))
    matrix = np.zeros((reduce(lambda count, x: count + len(x) - 1, res, 0), 2))
    mi = 0
    for j in range(len(res)):
        for i in range(1, len(res[j])):
            matrix[mi, 0] = res[j][0]
            matrix[mi, 1] = res[j][i]
            mi += 1

    return matrix[0:mi,:]

min_iter = 5000
max_iter = 10000
step_iter = 1000

rng = range(min_iter, max_iter, step_iter)
elapsed_naive = np.zeros(len(rng))
elapsed_kd_sn = np.zeros(len(rng))
elapsed_kd_tr = np.zeros(len(rng))
shapes = np.zeros((len(rng), 3))
ei = 0
for i in rng:
    random_cells = np.random.rand(i, 3) * 400.
    t = time()
    r1 = find_nn_naive(random_cells, 50.)
    shapes[ei, 0] = r1.shape[0]
    elapsed_naive[ei] = time() - t
    # print('naive shape:', r1.shape)
    # print('naive time:', elapsed_naive[ei])
    t = time()
    r2 = find_nn_kd_seminaive(random_cells, 50.)
    shapes[ei, 1] = r2.shape[0]
    elapsed_kd_sn[ei] = time() - t
    # print('seminaive shape:', r2.shape)
    # print('seminaive time:', elapsed_kd_sn[ei])
    t = time()
    r3 = find_nn_kd_by_tree(random_cells, 50.)
    shapes[ei, 2] = r3.shape[0]
    elapsed_kd_tr[ei] = time() - t
    # print('tree shape:', r3.shape)
    # print('tree time:', elapsed_kd_tr[ei])
    # print(r1,r2,r3)
    # exit()
    ei += 1

# Plot result comparison: Do all 3 implementations yield the same result? -> 3 overlapping lines
plt.plot(rng, shapes[:,0], label='naive')
plt.plot(rng, shapes[:,1], label='semi kd')
plt.plot(rng, shapes[:,2], label='full kd')
plt.legend()
plt.show(block=True)

# What's the runtime for each?
plt.plot(rng, elapsed_naive, label='naive')
plt.plot(rng, elapsed_kd_sn, label='semi kd')
plt.plot(rng, elapsed_kd_tr, label='full kd')
plt.legend()
plt.show(block=True)