import numpy as np
import pyvista as pv


ASC_path = "NavisanceDEM.asc"
nodes_path = "node.txt"


def load_asc_XYZ(path):
    tags = dict()
    with open(path) as file:
        for line in range(6):
            key, val = file.readline().split()
            tags[key] = eval(val)

    nx = tags["NCOLS"]
    ny = tags["NROWS"]
    cs = tags["CELLSIZE"]
    xmin = tags["XLLCENTER"]
    ymin = tags["YLLCENTER"]
    xmax = xmin + nx*cs
    ymax = ymin + ny*cs

    x = np.linspace(xmin, xmax, nx, endpoint=True)
    y = np.linspace(ymax, ymin, ny, endpoint=True)
    X, Y = np.meshgrid(x, y)

    Z = np.loadtxt(path, skiprows=len(tags))
    return X, Y, Z

def point_pick(point):
    points.points[:, :] = np.loadtxt(nodes_path)
    print(point, 1+np.argmin(((point - nodes)**2).sum(axis=1)))

nodes = np.loadtxt(nodes_path)
points = pv.PolyData(nodes)
points["id"] = np.arange(nodes.shape[0])

grid = pv.StructuredGrid(*load_asc_XYZ(ASC_path))
grid["Z"] = grid.points[:, 2]  # Z

pl = pv.Plotter()
pl.add_points(points, scalars="id", cmap="jet")
pl.add_mesh(grid, scalars="Z", show_scalar_bar=False)
pl.enable_point_picking(point_pick, left_clicking=True)
pl.show()

