#!/usr/bin/env python
# coding: utf-8

# In[1]:


import h5py
import numpy as np
import vtk

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


# In[6]:


blockname = '510'
infilename = 'density_sampled.cycle_000001/domain_000'+blockname+'.hdf5'
outfilename = 'out'+blockname+'.vtp'


# In[9]:


f = h5py.File(infilename, 'r')

np_data = np.asarray(f['fields']['Density']['values'])
print(np_data)

np_coordsX = np.asarray(f['coordsets']['coords']['values']['x'])
print(np_coordsX)
np_coordsY = np.asarray(f['coordsets']['coords']['values']['y'])
print(np_coordsY)
np_coordsZ = np.asarray(f['coordsets']['coords']['values']['z'])
print(np_coordsZ)

f.close()

tot_points = np.size(np_coordsX)

vtp_out = False

if vtp_out:

    name = 'Density'
    Points = vtk.vtkPoints()
    val_arr = vtk.vtkFloatArray()
    val_arr.SetNumberOfComponents(1)
    val_arr.SetName(name)

    #tot_points = np.size(np_coordsX)

    for i in range(tot_points):
        Points.InsertNextPoint(np_coordsX[i],np_coordsY[i],np_coordsZ[i])
        val_arr.InsertNextValue(np_data[i])

    polydata = vtk.vtkPolyData()
    polydata.SetPoints(Points)

    polydata.GetPointData().AddArray(val_arr)

    polydata.Modified()

    writer = vtk.vtkXMLPolyDataWriter();
    writer.SetFileName(outfilename);
    if vtk.VTK_MAJOR_VERSION <= 5:
        writer.SetInput(polydata)
    else:
        writer.SetInputData(polydata)
    writer.Write()


# In[ ]:


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

data_set = 'nyx' # 'isabel_pressure_10_percent' # 'nyx_5_percent_'
samp_method = 'hist' # random, hist , hist_grad, kdtree_histgrad_random
cur_samp = 'linear'
# 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)


feat_arr = np.zeros((tot_points,3))

print('total points:',tot_points)

data_vals = np.zeros(tot_points)


# 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
    
data_vals = np_data
    
feat_arr[:,0] = np_coordsX
feat_arr[:,1] = np_coordsY
feat_arr[:,2] = np_coordsZ

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

print("range:",range_min,range_max)

origin = range_min
spacing = (range_max- range_min) / (XDIM-1)


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_z0 = griddata(feat_arr, data_vals, cur_loc, method='nearest')
grid_z1 = griddata(feat_arr, data_vals, cur_loc, method='linear')
grid_z1[np.isnan(grid_z1)]=grid_z0[np.isnan(grid_z1)]
grid_z0 = grid_z1
#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_hdf5_'+blockname+'_'+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()


# In[ ]:




