import numpy as np
from sklearn.neighbors import NearestNeighbors
import vtk
from scipy.interpolate import griddata
import sys
import os

	
XDIM = 300 # 250 # 512
YDIM = 300 # 250 # 512
ZDIM = 300  # 50 # 512

data_set = 'asteroid_vo2_test_0.02' # 'isabel_pressure_10_percent' # 'nyx_5_percent_'
samp_method = 'kdtree_histgrad_random' # random, hist , hist_grad, kdtree_histgrad_random
cur_samp = 'nearest'
# var_name = "vo2" #"isabel_pressure" # 'nyx'

origin = np.array([-2.3e+6,  -500000,  -1.2e+6 ])
spacing = np.array([4.6e+6, 2.8e+6, 2.4e+6 ])/(300-1)

if len(sys.argv)>1:
	infile = sys.argv[1]
else:
	infile = data_set+'_sampled_'+samp_method+'.vtp'

filename, file_extension = os.path.splitext(infile)

if file_extension=='.vtp':
    poly_reader = vtk.vtkXMLPolyDataReader()
elif file_extension=='.vtk':
    poly_reader = vtk.vtkGenericDataObjectReader()
poly_reader.SetFileName(infile)
poly_reader.Update()

data = poly_reader.GetOutput()

var_name = data.GetPointData().GetArrayName(0)

print("total points:",data.GetNumberOfPoints(),data.GetNumberOfElements(0))

pts = data.GetPoints()

pt_data = data.GetPointData().GetArray(var_name).GetTuple1(100)
print('data:',pt_data)

print(pts.GetPoint(0))

tot_pts = data.GetNumberOfPoints()
feat_arr = np.zeros((tot_pts,3))

print('total points:',tot_pts)

data_vals = np.zeros(tot_pts)


for i in range(tot_pts):
    loc = pts.GetPoint(i)
    feat_arr[i,:] = np.asarray(loc)
    pt_data = data.GetPointData().GetArray(var_name).GetTuple1(i)
    data_vals[i] = pt_data

range_min = np.min(feat_arr,axis=0)
range_max = np.max(feat_arr,axis=0)

print("range:",range_min,range_max)

#n_neighbors = 5
#nbrs = NearestNeighbors(n_neighbors=n_neighbors, algorithm='ball_tree').fit(feat_arr)


#grid_x, grid_y, grid_z = np.mgrid[0:XDIM:5j, 0:YDIM:5j, 0:ZDIM:1j]
# grid_x, grid_y, grid_z = np.mgrid[0:XDIM, 0:YDIM, 0:ZDIM]
#
#
# grid_z0 = griddata(feat_arr, data_vals, (grid_x, grid_y, grid_z), method='nearest')
# #grid_z1 = griddata(feat_arr, data_vals, (grid_x, grid_y, grid_z), method='linear')
#
# grid_z0.tofile('out_full.bin')

cur_loc = np.zeros((XDIM*YDIM*ZDIM,3),dtype='double')

ind = 0
for k in range(ZDIM):
    for j in range(YDIM):
        for i in range(XDIM):
            cur_loc[ind,:] = origin + spacing * np.array([i,j,k])
            ind = ind+1

grid_z0 = griddata(feat_arr, data_vals, cur_loc, method=cur_samp)
#grid_z1 = griddata(feat_arr, data_vals, (grid_x, grid_y, grid_z), method='linear')
print("writing file:")
#grid_z0.tofile('recons_'+samp_method+'.raw')
grid_z0_3d = grid_z0.reshape((ZDIM,YDIM,XDIM))
# write to a vti file
filename = 'recons_'+data_set+'_'+samp_method+'_'+cur_samp+'.vti'
imageData = vtk.vtkImageData()
imageData.SetDimensions(XDIM, YDIM, ZDIM)

imageData.SetOrigin(origin)
imageData.SetSpacing(spacing)


if vtk.VTK_MAJOR_VERSION <= 5:
    imageData.SetNumberOfScalarComponents(1)
    imageData.SetScalarTypeToDouble()
else:
    imageData.AllocateScalars(vtk.VTK_DOUBLE, 1)

dims = imageData.GetDimensions()
print(dims)
# Fill every entry of the image data with "2.0"
for z in range(dims[2]):
    for y in range(dims[1]):
        for x in range(dims[0]):
            imageData.SetScalarComponentFromDouble(x, y, z, 0, grid_z0_3d[z,y,x])

writer = vtk.vtkXMLImageDataWriter()
writer.SetFileName(filename)
if vtk.VTK_MAJOR_VERSION <= 5:
    writer.SetInputConnection(imageData.GetProducerPort())
else:
    writer.SetInputData(imageData)
writer.Write()

# interp_data = np.zeros([XDIM,YDIM,ZDIM])
#
# for k in range(ZDIM):
#     for j in range(YDIM):
#         for i in range(XDIM):
#             cur_loc = np.array([i,j,k])
#             distances, indices = nbrs.kneighbors(cur_loc.reshape(1,-1))
#             vals = data_vals[indices]/
#             interp_data[i,j,k] =
#             print(i,j,k)




