import numpy as np
from swc import read_swc, write_swc
from neuron import h

ptntls = np.loadtxt("neuron_ISO.fld")
neuron = read_swc("input.swc")

coords = neuron['coords'] * 1e-6

ptntl_coords = ptntls[:, :3]

from scipy.spatial import cKDTree
ptntl_tree = cKDTree(ptntl_coords)

# query nearest neighbor for each SWC node
ptntl_idx = ptntl_tree.query(coords, k=1)[1]

ptntls_ordered  = ptntls[ptntl_idx] * 1e3

flags = [not np.isnan(p[3]) for p in ptntls_ordered]

non_nan = {
    'id' : neuron['id'][flags],
    'type' : neuron['type'][flags],
    'coords' : neuron['coords'][flags],
    'radius' : neuron['radius'][flags],
    'parent' : neuron['parent'][flags]
}

write_swc("pruned.swc", non_nan)

ptntls_reduce = ptntls_ordered[flags]


# load preamble
h.load_file("import3d.hoc")
imp = h.Import3d_SWC_read()
imp.input("pruned.swc")

gui = h.Import3d_GUI(imp)
gui.instantiate(None)

h.load_file("parameters.hoc")
h.load_file("interpxyz.hoc")
h.load_file("setup.hoc")

# insert xtra at each section
for sec in h.allsec():
    sec.insert("xtra")

h.load_file("setpointers.hoc")
h('setpointers()')

segs = []
seg_xyz = []

for sec in h.allsec():
    for seg in sec:
        segs.append(seg)
        seg_xyz.append([seg.xtra.x, seg.xtra.y, seg.xtra.z])

seg_xyz = np.array(seg_xyz)

# KD-Tree again to match closest segments
tree = cKDTree(non_nan['coords'])
_, node_idx = tree.query(seg_xyz)

i = 0
for seg, idx in zip(segs, node_idx):
    print(i)
    i += 1
    seg.xtra.es = ptntls_reduce[idx][3]

h.load_file("nrngui.hoc")
h.load_file("stdlib.hoc")
h.load_file("stim.hoc")
h.load_file("exportLocs_seg.hoc")
h.load_file("exportLocs_seg.hoc")
h.load_file("TMS_sim_simple.hoc")
h.load_file("detectFire.hoc")
h.load_file("run.hoc")