import numpy as np
import math
import os



import vtk
from vtk.util import numpy_support as VN
import random
from numpy import linalg as LA

import sys


if len(sys.argv) == 5:
	original_filename = sys.argv[1]
	stencil_info_filename = sys.argv[2]
	variable_name = sys.argv[3]
	output_filename = sys.argv[4]

	print("original_filename = " + original_filename)
	print("stencil_info_filename = " + stencil_info_filename)
	print("variable_name = " + variable_name)
	print("output_filename = " + output_filename)

	main_reader = vtk.vtkDataSetReader()
	main_reader.SetFileName(original_filename)
	main_reader.Update()
	main_data = main_reader.GetOutput()

	XDIM, YDIM, ZDIM = main_data.GetDimensions()

	x = np.zeros(main_data.GetNumberOfPoints())
	y = np.zeros(main_data.GetNumberOfPoints())
	z = np.zeros(main_data.GetNumberOfPoints())
	for i in range(main_data.GetNumberOfPoints()):
	    x[i],y[i],z[i] = main_data.GetPoint(i)
	    
	variable_data = main_data.GetPointData().GetArray(variable_name)
	variable_data = VN.vtk_to_numpy(variable_data)

	stencil_data_2d = np.loadtxt(stencil_info_filename, dtype='float32')

	sel_x = x[stencil_data_2d > 0.5]
	sel_y = y[stencil_data_2d > 0.5]
	sel_z = z[stencil_data_2d > 0.5]

	sel_var_data = variable_data[stencil_data_2d > 0.5]

	## create boundary points
	A = [0,XDIM-1]
	B = [0,YDIM-1]
	C = [0,ZDIM-1]
	# j_vals = ((xval,yval,zval) for xval in A for yval in B for zval in C)
	j_vals = (((zval * XDIM * YDIM) + (yval * XDIM) + xval) for xval in A for yval in B for zval in C)
	j_list = np.asanyarray(list(j_vals))

	Points = vtk.vtkPoints()

	for i in range(sel_x.shape[0]):
	    Points.InsertNextPoint(sel_x[i], sel_y[i], sel_z[i])

	for j in j_list:
	    ttx, tty, ttz = main_data.GetPoint(j)
	    Points.InsertNextPoint(ttx, tty, ttz)

	val_arr = vtk.vtkFloatArray()
	val_arr.SetNumberOfComponents(1)
	val_arr.SetName(variable_name)

	for i in range(sel_var_data.shape[0]):
	    val_arr.InsertNextValue(sel_var_data[i])

	for j in j_list:
	    val_arr.InsertNextValue(variable_data[j])


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

	polydata.GetPointData().AddArray(val_arr)
	polydata.Modified()


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


else:
	print("insufficient arg provided")