#include "project_surface_points_on_poly.h"
#include "closest_polymesh_point.h"

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

#include <limits>


vtkStandardNewMacro(ProjectSurfacePointsOnPoly);


ProjectSurfacePointsOnPoly::ProjectSurfacePointsOnPoly() {
	SetNumberOfInputPorts(2);

	RadiusScale = 2.;
}


int ProjectSurfacePointsOnPoly::FillInputPortInformation(int port, vtkInformation *info) {
	if (port == 0) {
		vtkUnstructuredGridAlgorithm::FillInputPortInformation(port, info);
	} else if (port == 1) {
		vtkNew<vtkInformation> input2;
		input2->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkPolyData");
		info->Append(input2);
	}
	return 1;
}


void findPointsWithinRadius(double radius, vtkPointSet *pointSet,
                            const double *point, vtkIdList *points) {
	for (vtkIdType i = 0; i < pointSet->GetNumberOfPoints(); i++) {
		double vec[3];
		vtkMath::Subtract(point, pointSet->GetPoint(i), vec);
		double dist = vtkMath::Norm(vec);
		if (dist < radius) {
			points->InsertNextId(i);
		}
	}
}


vtkSmartPointer<vtkIdList> removeSurfacePoints(vtkIdList *list,
                                               vtkIntArray *isSurface) {
	vtkNew<vtkIdList> affectedPoints;
	affectedPoints->Allocate(list->GetNumberOfIds());
	for (vtkIdType i = 0; i < list->GetNumberOfIds(); i++) {
		vtkIdType id = list->GetId(i);
		if (!isSurface->GetValue(id)) {
			affectedPoints->InsertNextId(id);
		}
	}
	return affectedPoints;
}


void ProjectSurfacePointsOnPoly::moveSurfacePoint(
		vtkUnstructuredGrid *tetMesh,
		vtkPolyData *polyMesh,
		vtkIdType pointId,
		std::vector<double> &motionVectors,
		std::vector<unsigned> &numberOfAffectors,
		vtkKdTree *polyMeshKdTree,
		vtkKdTree *tetMeshKdTree,
		vtkStaticCellLinks *links) {
	double point[3];
	tetMesh->GetPoint(pointId, point);
	double direction[3];
	double distance;
	closestPolyMeshPoint(
		polyMesh, point, polyMeshKdTree, links, direction, &distance);
	vtkNew<vtkIdList> pointsInRadius;
	findPointsWithinRadius(distance * RadiusScale, tetMesh, point, pointsInRadius);
	auto affectedPoints = removeSurfacePoints(pointsInRadius, isSurface);
	// tetMeshKdTree->FindPointsWithinRadius(distance * RadiusScale, point, affectedPoints);

	for (vtkIdType j = 0; j < affectedPoints->GetNumberOfIds(); j++) {
		vtkIdType affectedPointId = affectedPoints->GetId(j);
		double affectedPoint[3];
		tetMesh->GetPoint(affectedPointId, affectedPoint);
		double dist2 = vtkMath::Distance2BetweenPoints(affectedPoint, point);
		double expected = (distance * RadiusScale) * (distance * RadiusScale);
		if (dist2 > 2 * expected) {
			std::cerr << dist2 << " " << expected << "\n";
			// throw std::runtime_error("");
			continue;
		}

		double motion[3] = {direction[0], direction[1], direction[2]};
		double factor = 1. - std::sqrt(dist2) / (distance * RadiusScale);
		vtkMath::MultiplyScalar(motion, factor);
		motionVectors[affectedPointId * 3 + 0] += motion[0];
		motionVectors[affectedPointId * 3 + 1] += motion[1];
		motionVectors[affectedPointId * 3 + 2] += motion[2];
		numberOfAffectors[affectedPointId]++;
	}
	vtkMath::Subtract(point, direction, point);
	tetMesh->GetPoints()->SetPoint(pointId, point);
	tetMesh->Modified();
	affectedPoints->Reset();
}


static void applyMotionVectors(vtkUnstructuredGrid *tetMesh,
                               std::vector<double> &motionVectors,
                               const std::vector<unsigned> &numberOfAffectors) {
	for (vtkIdType i = 0; i < tetMesh->GetNumberOfPoints(); i++) {
		if(numberOfAffectors[i] == 0) continue;
		double point[3];
		tetMesh->GetPoint(i, point);
		double *motion = motionVectors.data() + (i * 3);
		vtkMath::MultiplyScalar(motion, 1. / numberOfAffectors[i]);
		vtkMath::Subtract(point, motion, point);
		tetMesh->GetPoints()->SetPoint(i, point);
		tetMesh->Modified();
	}
}


vtkTypeBool ProjectSurfacePointsOnPoly::RequestData(
	vtkInformation *request,
	vtkInformationVector **inputVector,
	vtkInformationVector *outputVector) {
	(void) request;
	vtkUnstructuredGrid* tetMesh =
		vtkUnstructuredGrid::GetData(inputVector[0]);
	vtkPolyData* polyMesh = vtkPolyData::GetData(inputVector[1]);
	vtkUnstructuredGrid* output = vtkUnstructuredGrid::GetData(outputVector);

	output->CopyStructure(tetMesh);
	output->GetPointData()->PassData(tetMesh->GetPointData());
	output->GetCellData()->PassData(tetMesh->GetCellData());

	/* Generate cell links, these tell us which cell a point belongs to. */
	vtkNew<vtkStaticCellLinks> links;
	links->BuildLinks(polyMesh);

	vtkNew<vtkKdTree> polyMeshKdTree;
	polyMeshKdTree->BuildLocatorFromPoints(polyMesh->GetPoints());
	vtkNew<vtkKdTree> tetMeshKdTree;
	tetMeshKdTree->BuildLocatorFromPoints(vtkPointSet::SafeDownCast(tetMesh));

	vtkIdTypeArray *surfacePoints =
		vtkIdTypeArray::SafeDownCast(tetMesh->GetFieldData()->GetArray("surfacePoints"));
	isSurface = vtkIntArray::SafeDownCast(
		output->GetPointData()->GetArray("isSurface"));
	std::vector<double> motionVectors(3 * tetMesh->GetNumberOfPoints(), 0);
	std::vector<unsigned> numberOfAffectors(tetMesh->GetNumberOfPoints(), 0);
	for (vtkIdType i = 0; i < surfacePoints->GetNumberOfTuples(); i++) {
		moveSurfacePoint(tetMesh, polyMesh,
		                 surfacePoints->GetValue(i),
		                 motionVectors,
		                 numberOfAffectors,
		                 polyMeshKdTree, tetMeshKdTree,
		                 links);
	}
	applyMotionVectors(tetMesh, motionVectors, numberOfAffectors);

	return true;
}
