#include "relaxation_filter.h"

#include <vtkUnstructuredGrid.h>
#include <vtkPointData.h>
#include <vtkCellData.h>
#include <vtkDoubleArray.h>
#include <vtkCellIterator.h>

#include <set>

vtkStandardNewMacro(RelaxationFilter);

RelaxationFilter::RelaxationFilter()
	:IterCount(1) {
}

vtkTypeBool RelaxationFilter::RequestData(
	vtkInformation *request,
	vtkInformationVector **inputVector,
	vtkInformationVector *outputVector) {
	(void) request;
	vtkUnstructuredGrid* input =
		vtkUnstructuredGrid::GetData(inputVector[0]);
	vtkUnstructuredGrid* output = vtkUnstructuredGrid::GetData(outputVector);
	output->CopyStructure(input);
	output->GetPointData()->PassData(input->GetPointData());
	output->GetCellData()->PassData(input->GetCellData());
	vtkNew<vtkPoints> newPoints;
	newPoints->DeepCopy(output->GetPoints());
	vtkNew<vtkIdList> neighborCells;
	vtkNew<vtkIdList> cellPoints;
	vtkIdTypeArray *surfacePoints = vtkIdTypeArray::SafeDownCast(
		output->GetFieldData()->GetArray("surfacePoints"));
	vtkIntArray *isSurface = vtkIntArray::SafeDownCast(
		output->GetPointData()->GetArray("isSurface"));
	std::set<vtkIdType> neighbors;
	output->BuildLinks();
	double avg[3];

	for (int iter = 0; iter < IterCount; iter++) {
		for (vtkIdType i = 0; i < surfacePoints->GetNumberOfValues(); i++) {
			vtkIdType id = surfacePoints->GetValue(i);

			output->GetPointCells(id, neighborCells);

			if (neighborCells->GetNumberOfIds() != 0) {

				for (vtkIdType j = 0; j < neighborCells->GetNumberOfIds(); j++) {

					vtkIdType cellId = neighborCells->GetId(j);
					output->GetCellPoints(cellId, cellPoints);

					for (vtkIdType k = 0; k < cellPoints->GetNumberOfIds(); k++) {

						vtkIdType neighborId = cellPoints->GetId(k);
						if (isSurface->GetValue(neighborId)) {
							neighbors.insert(neighborId);
						}
					}
					cellPoints->Reset();
				}
				neighbors.erase(neighbors.find(id));
				if (neighbors.size() != 0) {
					avg[0] = 0; avg[1] = 0; avg[2] = 0;
					for (const vtkIdType &j : neighbors) {
						double point[3];
						output->GetPoints()->GetPoint(j, point);
						vtkMath::Add(point, avg, avg);
					}
					vtkMath::MultiplyScalar(avg, 1. / neighbors.size());
					newPoints->SetPoint(id, avg);
				}
			}
			neighborCells->Reset();
			neighbors.clear();
		}

		output->SetPoints(newPoints);
	}
	return true;
}
