import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from scipy.spatial import KDTree
from networkx import Graph, shortest_path
import networkx as nx

# Generate data
N = 2000
K=7
x = np.random.rand(N,2) 
def F(X):
    if X.ndim == 1:
      X = X[None,:]
    X1 = X[:,0]
    X2 = X[:,1]
    Y1 = (X1 + 0.5)*np.cos(10*X1)
    Y2 = X2
    Y3 = (X1 + 0.5)*np.sin(10*X1)
    return np.column_stack((Y1, Y2, Y3))
y = F(x)

# Find points
x_left_id = np.argmin(np.sum((x - [0.15,0.15])**2, axis=1))
x_right_id = np.argmin(np.sum((x - [0.85,0.85])**2, axis=1))
x_left = x[x_left_id] 
x_right = x[x_right_id]

# Build KNN graph on y
knn = KDTree(y)
G_y = nx.Graph()
for i in range(N):
    idx = knn.query([y[i]], k=K)[1][0]
    for j in idx:
        G_y.add_edge(i,j)

G_x = nx.Graph()
# Convert edges from G_y to G_x
for e in G_y.edges:
    G_x.add_edge(e[0], e[1])

# Get shortest path nodes
path_x = shortest_path(G_x, x_left_id, x_right_id, 'weight')
path_y = [p for p in shortest_path(G_y, x_left_id, x_right_id, 'weight')]

# Plot
fig = plt.figure(figsize=(20, 20))

# A
ax = fig.add_subplot(221, projection='3d')
ax.scatter(y[:,0], y[:,1], y[:,2], c='k')
ax.scatter(F(x_left)[0,0], F(x_left)[0,1], F(x_left)[0,2], c='b', marker='o', s=100) 
ax.scatter(F(x_right)[0,0], F(x_right)[0,1], F(x_right)[0,2], c='b', marker='o', s=100)
t = np.linspace(0,1,1000)
ax.plot(F(x_left + t[:,None]*(x_right-x_left))[:,0], 
        F(x_left + t[:,None]*(x_right-x_left))[:,1],
        F(x_left + t[:,None]*(x_right-x_left))[:,2], c='b')
ax.set_title('A')

# B            
ax = fig.add_subplot(222, projection='3d')
ax.scatter(y[:,0], y[:,1], y[:,2], c='k')
for e in G_y.edges:
    ax.plot(y[e,0], y[e,1], y[e,2], c='grey')
ax.plot(np.array(y)[path_y,0], np.array(y)[path_y,1], np.array(y)[path_y,2], c='r')   
ax.set_title('B')

# C
ax = fig.add_subplot(212)
ax.scatter(x[:,0], x[:,1], c='k')
for e in G_x.edges:
    ax.plot(x[e,0], x[e,1], c='grey')  
ax.plot([x_left[0], x_right[0]], [x_left[1], x_right[1]], c='b')
ax.plot(np.array(x)[path_x,0], np.array(x)[path_x,1], c='r')
ax.set_title('C')

plt.tight_layout()
plt.show()