//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of NVIDIA CORPORATION nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Copyright (c) 2018 NVIDIA Corporation. All rights reserved.


// Copyright (c) 2004-2008 AGEIA Technologies, Inc. All rights reserved.
// Copyright (c) 2001-2004 NovodeX AG. All rights reserved.  

#include "ExtClothConfig.h"
#if APEX_USE_CLOTH_API

#include "ExtClothFabricCooker.h"
#include "ExtClothTetherCooker.h"
#include "PxVec4.h"
#include "PsFoundation.h"
#include "PxErrorCallback.h"
#include "PsArray.h"
#include "PsHashMap.h"
#include "PsSort.h"
#include "PxIO.h"

#include "PxStrideIterator.h"

#pragma warning(disable:4127)

using namespace nvidia;
using namespace physx;

struct nvidia::PxFabricCookerImpl
{
	bool cook(const PxClothMeshDesc& desc, PxVec3 gravity, bool useGeodesicTether);

	PxClothFabricDesc getDescriptor() const;
	void save(physx::PxOutputStream& stream, bool platformMismatch) const;

public:
	uint32_t mNumParticles;

	shdfnd::Array<PxClothFabricPhase> mPhases;
	shdfnd::Array<uint32_t> mSets; // with 0 prefix
	shdfnd::Array<float> mRestvalues;
	shdfnd::Array<uint32_t> mIndices;

	shdfnd::Array<uint32_t> mTetherAnchors;
	shdfnd::Array<float> mTetherLengths;
};

PxClothFabricCooker::PxClothFabricCooker(const PxClothMeshDesc& desc, const PxVec3& gravity, bool useGeodesicTether)
: mImpl(new PxFabricCookerImpl())
{
	mImpl->cook(desc, gravity, useGeodesicTether);
}

PxClothFabricCooker::~PxClothFabricCooker()
{
	delete mImpl;
}

PxClothFabricDesc PxClothFabricCooker::getDescriptor() const
{
	return mImpl->getDescriptor();
}

void PxClothFabricCooker::save(physx::PxOutputStream& stream, bool platformMismatch) const
{
	mImpl->save(stream, platformMismatch);
}


namespace
{
	// calculate the inclusive prefix sum, equivalent of std::partial_sum
	template <typename T>
	void prefixSum(const T* first, const T* last, T* dest)
	{
		if (first != last)
		{	
			*(dest++) = *(first++);
			for (; first != last; ++first, ++dest)
				*dest = *(dest-1) + *first;
		}
	}

	template <typename T>
	void gatherAdjacencies(shdfnd::Array<uint32_t>& valency, shdfnd::Array<uint32_t>& adjacencies, 
		const PxBoundedData& triangles, const PxBoundedData& quads)
	{
		// count number of edges per vertex
		PxStrideIterator<const T> tIt, qIt;
		tIt = physx::PxMakeIterator((const T*)triangles.data, triangles.stride);
		for(uint32_t i=0; i<triangles.count; ++i, ++tIt, ++qIt)
		{
			for(uint32_t j=0; j<3; ++j)
				valency[tIt.ptr()[j]] += 2;
		}
		qIt = physx::PxMakeIterator((const T*)quads.data, quads.stride);
		for(uint32_t i=0; i<quads.count; ++i, ++tIt, ++qIt)
		{
			for(uint32_t j=0; j<4; ++j)
				valency[qIt.ptr()[j]] += 2;
		}

		prefixSum(valency.begin(), valency.end(), valency.begin());
		adjacencies.resize(valency.back());

		// gather adjacent vertices
		tIt = physx::PxMakeIterator((const T*)triangles.data, triangles.stride);
		for(uint32_t i=0; i<triangles.count; ++i, ++tIt)
		{
			for(uint32_t j=0; j<3; ++j)
			{
				adjacencies[--valency[tIt.ptr()[j]]] = tIt.ptr()[(j+1)%3];
				adjacencies[--valency[tIt.ptr()[j]]] = tIt.ptr()[(j+2)%3];
			}
		}
		qIt = physx::PxMakeIterator((const T*)quads.data, quads.stride);
		for(uint32_t i=0; i<quads.count; ++i, ++qIt)
		{
			for(uint32_t j=0; j<4; ++j)
			{
				adjacencies[--valency[qIt.ptr()[j]]] = qIt.ptr()[(j+1)%4];
				adjacencies[--valency[qIt.ptr()[j]]] = qIt.ptr()[(j+3)%4];
			}
		}
	}

	
	struct Edge
	{
		Edge() : mStretching(0.0f), mBending(0.0f), mShearing(0.0f) {}

		void classify()
		{
			mStretching += 0.1f;
		}

		// classify v0-v2 edge based on alternative v0-v1-v2 path 
		void classify(const PxVec4& v0, const PxVec4& v1, const PxVec4& v2)
		{
			const PxVec3& p0 = reinterpret_cast<const PxVec3&>(v0);
			const PxVec3& p1 = reinterpret_cast<const PxVec3&>(v1);
			const PxVec3& p2 = reinterpret_cast<const PxVec3&>(v2);

			float area = (p1-p0).cross(p2-p1).magnitude();
			// triangle height / base length
			// 1.0 = quad edge, 0.2 = quad diagonal + quad edge, 
			float ratio = area / (p2-p0).magnitudeSquared();

			// 0.5 = quad diagonal
			mShearing += PxMax(0.0f, 0.15f - fabsf(0.45f - ratio));
			// 0.0 = collinear points
			mBending += PxMax(0.0f, 0.1f - ratio) * 3;
		}

		float mStretching;
		float mBending;
		float mShearing;
	};

	typedef shdfnd::Pair<uint32_t, uint32_t> Pair;
	typedef shdfnd::Pair<Pair, PxClothFabricPhaseType::Enum> Entry;

	// maintain heap status after elements have been pushed (heapify)
	template<typename T>
	void pushHeap(shdfnd::Array<T> &heap, const T &value)
	{
		heap.pushBack(value);
		T* begin = heap.begin();
		T* end = heap.end();

		if (end <= begin)
			return;
	
		uint32_t current = uint32_t(end - begin) - 1;
		while (current > 0)
		{
			const uint32_t parent = (current - 1) / 2;
			if (!(begin[parent] < begin[current]))
				break;

			shdfnd::swap(begin[parent], begin[current]);
			current = parent;
		}
	}

	// pop one element from the heap
	template<typename T>
	T popHeap(shdfnd::Array<T> &heap)
	{
		T* begin = heap.begin();
		T* end = heap.end();

		shdfnd::swap(begin[0], end[-1]); // exchange elements

		// shift down
		end--;

		uint32_t current = 0;
		while (begin + (current * 2 + 1) < end)
		{
			uint32_t child = current * 2 + 1;
			if (begin + child + 1 < end && begin[child] < begin[child + 1])
				++child;

			if (!(begin[current] < begin[child]))
				break;

			shdfnd::swap(begin[current], begin[child]);
			current = child;
		}

		return heap.popBack();
	}

	// ---------------------------------------------------------------------------------------
	// Heap element to sort constraint based on graph color count
	struct ConstraintGraphColorCount
	{
		ConstraintGraphColorCount(int cid, int count) 
			: constraint((uint32_t)cid), colorCount((uint32_t)count) {}

		uint32_t constraint;
		uint32_t colorCount; 

		bool operator < (const ConstraintGraphColorCount& c) const
		{
			return colorCount < c.colorCount;
		}
	};

	struct ConstraintSorter
	{
	public:

		ConstraintSorter(uint32_t* constraints_) : constraints(constraints_) {}

		bool operator()(uint32_t i, uint32_t j) const
		{
			uint32_t ci = i*2;
			uint32_t cj = j*2;

			if (constraints[ci] == constraints[cj])
				return constraints[ci+1] < constraints[cj+1];
			else
				return constraints[ci] < constraints[cj];
		}

		uint32_t* constraints;
	};

} // anonymous namespace

bool nvidia::PxFabricCookerImpl::cook(const PxClothMeshDesc& desc, PxVec3 gravity, bool useGeodesicTether)
{	
	if(!desc.isValid())
	{
		shdfnd::getFoundation().getErrorCallback().reportError(PxErrorCode::eINVALID_PARAMETER, 
			"PxFabricCookerImpl::cook: desc.isValid() failed!", __FILE__, __LINE__);
		return false;
	}

	gravity = gravity.getNormalized();

	mNumParticles = desc.points.count;

	// assemble points
	shdfnd::Array<PxVec4> particles;
	particles.reserve(mNumParticles);
	PxStrideIterator<const PxVec3> pIt((const PxVec3*)desc.points.data, desc.points.stride);
	PxStrideIterator<const float> wIt((const float*)desc.invMasses.data, desc.invMasses.stride);
	for(uint32_t i=0; i<mNumParticles; ++i)
		particles.pushBack(PxVec4(*pIt++, wIt.ptr() ? *wIt++ : 1.0f));

	// build adjacent vertex list
	shdfnd::Array<uint32_t> valency(mNumParticles+1, 0);
	shdfnd::Array<uint32_t> adjacencies;
	if(desc.flags & PxMeshFlag::e16_BIT_INDICES)
		gatherAdjacencies<uint16_t>(valency, adjacencies, desc.triangles, desc.quads);
	else
		gatherAdjacencies<uint32_t>(valency, adjacencies, desc.triangles, desc.quads);

	// build unique neighbors from adjacencies
	shdfnd::Array<uint32_t> mark(valency.size(), 0);
	shdfnd::Array<uint32_t> neighbors; neighbors.reserve(adjacencies.size());
	for(uint32_t i=1, j=0; i<valency.size(); ++i)
	{
		for(; j<valency[i]; ++j)
		{
			uint32_t k = adjacencies[j];
			if(mark[k] != i)
			{
				mark[k] = i;
				neighbors.pushBack(k);
			}
		}
		valency[i] = neighbors.size();
	}

	// build map of unique edges and classify
	shdfnd::HashMap<Pair, Edge> edges;
	for(uint32_t i=0; i<mNumParticles; ++i)
	{
		float wi = particles[i].w;
		// iterate all neighbors
		uint32_t jlast = valency[i+1];
		for(uint32_t j=valency[i]; j<jlast; ++j)
		{
			// add 1-ring edge
			uint32_t m = neighbors[j];
			if(wi + particles[m].w > 0.0f)
				edges[Pair(PxMin(i, m), PxMax(i, m))].classify();

			// iterate all neighbors of neighbor
			uint32_t klast = valency[m+1];
			for(uint32_t k=valency[m]; k<klast; ++k)
			{
				uint32_t n = neighbors[k];
				if(n != i && wi + particles[n].w > 0.0f)
				{
					// add 2-ring edge
					edges[Pair(PxMin(i, n), PxMax(i, n))].classify(
						particles[i], particles[m], particles[n]);
				}
			}
		}
	}

	// copy classified edges to constraints array
	// build histogram of constraints per vertex
	shdfnd::Array<Entry> constraints; 	
	constraints.reserve(edges.size());
	valency.resize(0); valency.resize(mNumParticles+1, 0);

	const float sqrtHalf = PxSqrt(0.4f);
	for(shdfnd::HashMap<Pair, Edge>::Iterator eIt = edges.getIterator(); !eIt.done(); ++eIt)
	{
		const Edge& edge = eIt->second;
		const Pair& pair = eIt->first;
		if((edge.mStretching + edge.mBending + edge.mShearing) > 0.0f)
		{	
			PxClothFabricPhaseType::Enum type = PxClothFabricPhaseType::eINVALID;
			if(edge.mBending > PxMax(edge.mStretching, edge.mShearing))
				type = PxClothFabricPhaseType::eBENDING;
			else if(edge.mShearing > PxMax(edge.mStretching, edge.mBending))
				type = PxClothFabricPhaseType::eSHEARING;
			else 
			{
				PxVec4 diff = particles[pair.first]-particles[pair.second];
				float dot = gravity.dot(reinterpret_cast<const PxVec3&>(diff).getNormalized());
				type = fabsf(dot) < sqrtHalf ? PxClothFabricPhaseType::eHORIZONTAL : PxClothFabricPhaseType::eVERTICAL;
			}
			++valency[pair.first];
			++valency[pair.second];
			constraints.pushBack(Entry(pair, type));
		}
	} 

	prefixSum(valency.begin(), valency.end(), valency.begin());

	uint32_t numConstraints = constraints.size();

	// build adjacent constraint list
	adjacencies.resize(0); adjacencies.resize(valency.back(), 0);
	for(uint32_t i=0; i<numConstraints; ++i)
	{
		adjacencies[--valency[constraints[i].first.first]] = i;
		adjacencies[--valency[constraints[i].first.second]] = i;
	}
	
	shdfnd::Array<uint32_t>::ConstIterator aFirst = adjacencies.begin();
	shdfnd::Array<uint32_t> colors(numConstraints, numConstraints); // constraint -> color, initialily not colored
	mark.resize(0); mark.resize(numConstraints+1, UINT32_MAX); // color -> constraint index
	shdfnd::Array<uint32_t> adjColorCount(numConstraints, 0); // # of neighbors that are already colored

	shdfnd::Array<ConstraintGraphColorCount> constraintHeap; 
	constraintHeap.reserve(numConstraints); // set of constraints to color (added in edge distance order)

	// Do graph coloring based on edge distance.
	// For each constraint, we add its uncolored neighbors to the heap
	// ,and we pick the constraint with most colored neighbors from the heap.
	while ( true )
	{
		uint32_t constraint = 0;
		while ( (constraint < numConstraints) && (colors[constraint] != numConstraints))
			constraint++; // start with the first uncolored constraint
	
		if (constraint >= numConstraints)
			break;

		constraintHeap.clear();
		pushHeap(constraintHeap, ConstraintGraphColorCount((int)constraint, (int)adjColorCount[constraint]));
		PxClothFabricPhaseType::Enum type = constraints[constraint].second;
		
		while (!constraintHeap.empty())
		{		
			ConstraintGraphColorCount heapItem = popHeap(constraintHeap);
			constraint = heapItem.constraint;
			if (colors[constraint] != numConstraints)
				continue; // skip if already colored 

			const Pair& pair = constraints[constraint].first;			
			for(uint32_t j=0; j<2; ++j)
			{
				uint32_t index = j ? pair.first : pair.second;
				if(particles[index].w == 0.0f)
					continue; // don't mark adjacent particles if attached

				for(shdfnd::Array<uint32_t>::ConstIterator aIt = aFirst + valency[index], aEnd = aFirst + valency[index+1]; aIt != aEnd; ++aIt)
				{				
					uint32_t adjacentConstraint = *aIt;
					if ((constraints[adjacentConstraint].second != type) || (adjacentConstraint == constraint))
						continue;

					mark[colors[adjacentConstraint]] = constraint; 
					++adjColorCount[adjacentConstraint];
					pushHeap(constraintHeap, ConstraintGraphColorCount((int)adjacentConstraint, (int)adjColorCount[adjacentConstraint]));
				}
			}

			// find smallest color with matching type
			uint32_t color = 0;
			while((color < mPhases.size() && mPhases[color].phaseType != type) || mark[color] == constraint)
				++color;

			// create a new color set
			if(color == mPhases.size())
			{
				PxClothFabricPhase phase(type, mPhases.size());
				mPhases.pushBack(phase);
				mSets.pushBack(0);
			}

			colors[constraint] = color;
			++mSets[color];
		} 
	}

#if 0 // PX_DEBUG
	printf("set[%u] = ", mSets.size());
	for(uint32_t i=0; i<mSets.size(); ++i)
		printf("%u ", mSets[i]);
#endif

	prefixSum(mSets.begin(), mSets.end(), mSets.begin());

#if 0 // PX_DEBUG
	printf(" = %u\n", mSets.back());
#endif

	// write indices and rest lengths
	// convert mSets to exclusive sum
	uint32_t back = mSets.back();
	mSets.pushBack(back);
	mIndices.resize(numConstraints*2);
	mRestvalues.resize(numConstraints);
	for(uint32_t i=0; i<numConstraints; ++i)
	{
		uint32_t first = constraints[i].first.first;
		uint32_t second = constraints[i].first.second;

		uint32_t index = --mSets[colors[i]];

		mIndices[2*index  ] = first;
		mIndices[2*index+1] = second;

		PxVec4 diff = particles[second] - particles[first];
		mRestvalues[index] = reinterpret_cast<
			const PxVec3&>(diff).magnitude();
	} 
	
	// reorder constraints and rest values for more efficient cache access (linear)
	shdfnd::Array<uint32_t> newIndices(mIndices.size());
	shdfnd::Array<float> newRestValues(mRestvalues.size());

	// sort each constraint set in vertex order
	for (uint32_t i=0; i < mSets.size()-1; ++i)
	{
		// create a re-ordering list
		shdfnd::Array<uint32_t> reorder(mSets[i+1]-mSets[i]);

		for (uint32_t r=0; r < reorder.size(); ++r)
			reorder[r] = r;

		const uint32_t indicesOffset = mSets[i]*2;
		const uint32_t restOffset = mSets[i];

		ConstraintSorter predicate(&mIndices[indicesOffset]);
		shdfnd::sort(&reorder[0], reorder.size(), predicate);
		
		for (uint32_t r=0; r < reorder.size(); ++r)
		{
			newIndices[indicesOffset + r*2] = mIndices[indicesOffset + reorder[r]*2];
			newIndices[indicesOffset + r*2+1] = mIndices[indicesOffset + reorder[r]*2+1];
			newRestValues[restOffset + r] = mRestvalues[restOffset + reorder[r]];
		}
	}

	mIndices = newIndices;
	mRestvalues = newRestValues;

	PX_ASSERT(mIndices.size() == mRestvalues.size()*2);
	PX_ASSERT(mRestvalues.size() == mSets.back());

#if 0 // PX_DEBUG
	for (uint32_t i = 1; i < mSets.size(); i++)
	{
		PxClothFabricPhase phase = mPhases[i-1];
		printf("%d : type %d, size %d\n", 
			i-1, phase.phaseType, mSets[i] - mSets[i-1]);
	}
#endif

	if (useGeodesicTether)
	{
		PxClothGeodesicTetherCooker tetherCooker(desc);
		if (tetherCooker.getCookerStatus() == 0)
		{
			uint32_t numTethersPerParticle = tetherCooker.getNbTethersPerParticle();
			uint32_t tetherSize = mNumParticles * numTethersPerParticle;
			mTetherAnchors.resize(tetherSize);
			mTetherLengths.resize(tetherSize);
			tetherCooker.getTetherData(mTetherAnchors.begin(), mTetherLengths.begin());
		}
		else
			useGeodesicTether = false;
	}

	if (!useGeodesicTether)
	{
		PxClothSimpleTetherCooker tetherCooker(desc);
		if (tetherCooker.getCookerStatus() == 0)
		{
			mTetherAnchors.resize(mNumParticles);
			mTetherLengths.resize(mNumParticles);
			tetherCooker.getTetherData(mTetherAnchors.begin(), mTetherLengths.begin());
		}
	}

	return true;
}

PxClothFabricDesc nvidia::PxFabricCookerImpl::getDescriptor() const
{
	PxClothFabricDesc result;

	result.nbParticles = mNumParticles;
	result.nbPhases = mPhases.size();
	result.phases = mPhases.begin();
	result.nbSets = mSets.size()-1;
	result.sets = mSets.begin()+1;
	result.restvalues = mRestvalues.begin();
	result.indices = mIndices.begin();
	result.nbTethers = mTetherAnchors.size();
	result.tetherAnchors = mTetherAnchors.begin();
	result.tetherLengths = mTetherLengths.begin();

	return result;
}

void nvidia::PxFabricCookerImpl::save( physx::PxOutputStream& stream, bool /*platformMismatch*/ ) const
{
	// version 1 is equivalent to 0x030300 and 0x030301 (PX_PHYSICS_VERSION of 3.3.0 and 3.3.1).
	// If the stream format changes, the loader code in ScClothFabricCore.cpp 
	// and the version number need to change too. 
	uint32_t version = 1;
	stream.write(&version, sizeof(uint32_t));

	PxClothFabricDesc desc = getDescriptor();

	// write explicit sizes, others are implicit
	stream.write(&mNumParticles, sizeof(uint32_t));
	stream.write(&desc.nbPhases, sizeof(uint32_t));
	stream.write(&desc.nbSets, sizeof(uint32_t));
	stream.write(&desc.nbTethers, sizeof(uint32_t));

	uint32_t nbConstraints = desc.sets[desc.nbSets-1];

	// write actual data
	PX_COMPILE_TIME_ASSERT(sizeof(PxClothFabricPhaseType::Enum) == sizeof(uint32_t));
	stream.write(desc.phases, desc.nbPhases*sizeof(PxClothFabricPhase));
	stream.write(desc.sets, desc.nbSets*sizeof(uint32_t));

	stream.write(desc.restvalues, nbConstraints*sizeof(float));
	stream.write(desc.indices, nbConstraints*2*sizeof(uint32_t));

	stream.write(desc.tetherAnchors, desc.nbTethers*sizeof(uint32_t));
	stream.write(desc.tetherLengths, desc.nbTethers*sizeof(float));
}

#endif //APEX_USE_CLOTH_API