#include "crackFreePrimitiveGenerator.h"
#include <maya/MStatus.h>
#include <maya/MFnMesh.h>
#include <maya/MIntArray.h>
#include <maya/MFloatArray.h>
#include <maya/MPxVertexBufferGenerator.h>
#include <maya/MHWGeometry.h>
#include <maya/MDrawRegistry.h>
#include <maya/MComponentDataIndexing.h>
#include <map>
namespace
{
    struct Edge
    {
        Edge(unsigned int v0_ = 0, unsigned int v1_ = 0) : v0(v0_), v1(v1_) {}
        Edge reversed() const
        {
            return Edge(v1, v0);
        }
        bool isEqual(const Edge &rhs) const
        {
            return (v0 == rhs.v0 && v1 == rhs.v1);
        }
        bool isReverse(const Edge &rhs) const
        {
            return (v0 == rhs.v1 && v1 == rhs.v0);
        }
        unsigned int v0;
        unsigned int v1;
    };
    bool operator< (const Edge& lhs, const Edge& rhs)
    {
        return (lhs.v0 < rhs.v0) || (lhs.v0 == rhs.v0 && lhs.v1 < rhs.v1);
    }
    class EdgeMapping
    {
    public:
        EdgeMapping();
        void addTriangle(unsigned int faceVertexId0, unsigned int faceVertexId1, unsigned int faceVertexId2, unsigned int polyVertexId0, unsigned int polyVertexId1, unsigned int polyVertexId2);
        void addEdge(const Edge& faceEdge, const Edge& polyEdge);
        void addPositionUV(unsigned int faceVertexId, unsigned int polyVertexId, float u, float v);
        bool adjacentEdge(const Edge& faceEdge, Edge& adjacentEdge) const;
        bool dominantEdge(const Edge& faceEdge, Edge& dominantEdge) const;
        bool dominantPosition(unsigned int faceVertexId, unsigned int& dominantVertexId) const;
    private:
        
        typedef std::map< Edge, Edge > FaceEdge2PolyEdgeMap;
        FaceEdge2PolyEdgeMap fFaceEdge2PolyEdgeMap;
        
        typedef std::map< Edge, std::pair< Edge, Edge > > PolyEdge2FaceEdgeMap;
        PolyEdge2FaceEdgeMap fPolyEdge2FaceEdgesMap;
        
        typedef std::map< unsigned int, unsigned int > FaceVertex2PolyVertexMap;
        FaceVertex2PolyVertexMap fFaceVertex2PolyVertexMap;
        
        typedef std::pair< unsigned int, std::pair< float, float > > VertexUVPair;
        typedef std::map< unsigned int, VertexUVPair > PolyVertex2FaceVertexUVMap;
        PolyVertex2FaceVertexUVMap fPolyVertex2FaceVertexUVMap;
    };
    EdgeMapping::EdgeMapping()
    {
    }
    void EdgeMapping::addTriangle(unsigned int faceVertexId0, unsigned int faceVertexId1, unsigned int faceVertexId2, unsigned int polyVertexId0, unsigned int polyVertexId1, unsigned int polyVertexId2)
    {
        addEdge( Edge(faceVertexId0, faceVertexId1), Edge(polyVertexId0, polyVertexId1) );
        addEdge( Edge(faceVertexId1, faceVertexId2), Edge(polyVertexId1, polyVertexId2) );
        addEdge( Edge(faceVertexId2, faceVertexId0), Edge(polyVertexId2, polyVertexId0) );
    }
    
    
    void EdgeMapping::addEdge(const Edge& faceEdge, const Edge& polyEdge)
    {
        if(polyEdge.v1 < polyEdge.v0)
        {
            
            Edge faceEdgeR = faceEdge.reversed();
            Edge polyEdgeR = polyEdge.reversed();
            addEdge(faceEdgeR, polyEdgeR);
            return;
        }
        fFaceEdge2PolyEdgeMap[faceEdge] = polyEdge;
        
        PolyEdge2FaceEdgeMap::iterator itP2F = fPolyEdge2FaceEdgesMap.find( polyEdge );
        if(itP2F == fPolyEdge2FaceEdgesMap.end())
        {
            fPolyEdge2FaceEdgesMap[polyEdge] = std::make_pair( faceEdge, faceEdge );
        }
        else
        {
            itP2F->second.second = faceEdge;
        }
    }
    void EdgeMapping::addPositionUV(unsigned int faceVertexId, unsigned int polyVertexId, float u, float v)
    {
        fFaceVertex2PolyVertexMap[faceVertexId] = polyVertexId;
        PolyVertex2FaceVertexUVMap::iterator it = fPolyVertex2FaceVertexUVMap.find(polyVertexId);
        if(it == fPolyVertex2FaceVertexUVMap.end())
        {
            fPolyVertex2FaceVertexUVMap[polyVertexId] = std::make_pair( faceVertexId, std::make_pair(u,v) );
        }
        else
        {
            VertexUVPair& vertexUVPair = it->second;
            float lastU = vertexUVPair.second.first;
            float lastV = vertexUVPair.second.second;
            if( u < lastU || ( u == lastU && v < lastV ) )
            {
                vertexUVPair.first = faceVertexId;
                vertexUVPair.second = std::make_pair(u,v);
            }
        }
    }
    
    
    
    bool EdgeMapping::adjacentEdge(const Edge& faceEdge, Edge& adjacentEdge) const
    {
        FaceEdge2PolyEdgeMap::const_iterator itF2P = fFaceEdge2PolyEdgeMap.find( faceEdge );
        if(itF2P == fFaceEdge2PolyEdgeMap.end())
        {
            Edge faceEdgeR = faceEdge.reversed();
            itF2P = fFaceEdge2PolyEdgeMap.find( faceEdgeR );
            if(itF2P == fFaceEdge2PolyEdgeMap.end())
                return false;
        }
        const Edge& polyEdge = itF2P->second;
        PolyEdge2FaceEdgeMap::const_iterator itP2F = fPolyEdge2FaceEdgesMap.find( polyEdge );
        if(itP2F == fPolyEdge2FaceEdgesMap.end())
            return false;
        const Edge& faceEdge0 = itP2F->second.first;
        const Edge& faceEdge1 = itP2F->second.second;
        bool foundMatch = false;
        if(faceEdge.isEqual(faceEdge0))
        {
            adjacentEdge = faceEdge1;
            foundMatch = true;
        }
        else if(faceEdge.isReverse(faceEdge0))
        {
            adjacentEdge = faceEdge1.reversed();
            foundMatch = true;
        }
        else if(faceEdge.isEqual(faceEdge1))
        {
            adjacentEdge = faceEdge0;
            foundMatch = true;
        }
        else if(faceEdge.isReverse(faceEdge1))
        {
            adjacentEdge = faceEdge0.reversed();
            foundMatch = true;
        }
        return foundMatch;
    }
    bool EdgeMapping::dominantEdge(const Edge& faceEdge, Edge& dominantEdge) const
    {
        bool returnReversed = true;
        FaceEdge2PolyEdgeMap::const_iterator itF2P = fFaceEdge2PolyEdgeMap.find( faceEdge );
        if(itF2P == fFaceEdge2PolyEdgeMap.end())
        {
            Edge faceEdgeR = faceEdge.reversed();
            itF2P = fFaceEdge2PolyEdgeMap.find( faceEdgeR );
            if(itF2P == fFaceEdge2PolyEdgeMap.end())
                return false;
            returnReversed = false;
        }
        const Edge& polyEdge = itF2P->second;
        PolyEdge2FaceEdgeMap::const_iterator itP2F = fPolyEdge2FaceEdgesMap.find( polyEdge );
        if(itP2F == fPolyEdge2FaceEdgesMap.end())
            return false;
        const Edge& faceEdge0 = itP2F->second.first;
        const Edge& faceEdge1 = itP2F->second.second;
        if(faceEdge0 < faceEdge1)
        {
            dominantEdge = (returnReversed ? faceEdge0.reversed() : faceEdge0);
        }
        else
        {
            dominantEdge = (returnReversed ? faceEdge1.reversed() : faceEdge1);
        }
        return true;
    }
    bool EdgeMapping::dominantPosition(unsigned int faceVertexId, unsigned int& dominantVertexId) const
    {
        FaceVertex2PolyVertexMap::const_iterator itF2P = fFaceVertex2PolyVertexMap.find(faceVertexId);
        if(itF2P == fFaceVertex2PolyVertexMap.end())
            return false;
        unsigned polyVertexId = itF2P->second;
        PolyVertex2FaceVertexUVMap::const_iterator itP2FUV = fPolyVertex2FaceVertexUVMap.find(polyVertexId);
        if(itP2FUV == fPolyVertex2FaceVertexUVMap.end())
            return false;
        dominantVertexId = itP2FUV->second.first;
        return true;
    }
    struct VertexF
    {
        static const float kTolerance;
        VertexF(const float* buffer, unsigned int index)
        {
            unsigned int bufferPos = index * 3;
            x = buffer[bufferPos++];
            y = buffer[bufferPos++];
            z = buffer[bufferPos++];
        }
        bool isEqual(const VertexF &rhs) const
        {
            return (fabs(x - rhs.x) < kTolerance && fabs(y - rhs.y) < kTolerance && fabs(z - rhs.z) < kTolerance);
        }
        float x, y, z;
    };
    const float VertexF::kTolerance = 1e-5f;
    bool operator< (const VertexF& lhs, const VertexF& rhs)
    {
        return ((lhs.x - rhs.x) < -VertexF::kTolerance) || 
               (fabs(lhs.x - rhs.x) < VertexF::kTolerance && (lhs.y - rhs.y) < -VertexF::kTolerance) || 
               (fabs(lhs.x - rhs.x) < VertexF::kTolerance && fabs(lhs.y - rhs.y) < VertexF::kTolerance && (lhs.z - rhs.z) < -VertexF::kTolerance);
    }
    struct VertexFMap
    {
        unsigned int getVertexId( const VertexF& v );
        typedef std::map<VertexF, unsigned int> TVtxMap;
        TVtxMap vertexMap;
    };
    unsigned int VertexFMap::getVertexId( const VertexF& v )
    {
        VertexFMap::TVtxMap::const_iterator itVtx = vertexMap.find(v);
        if (itVtx != vertexMap.end())
            return itVtx->second;
        unsigned int nextId = (unsigned int)vertexMap.size();
        vertexMap.insert(TVtxMap::value_type(v,nextId));
        return nextId;
    }
}
CrackFreePrimitiveGenerator::CrackFreePrimitiveGenerator(bool addAdjacentEdges, bool addDominantEdges, bool addDominantPosition)
: fAddAdjacentEdges(addAdjacentEdges)
, fAddDominantEdges(addDominantEdges)
, fAddDominantPosition(addDominantPosition)
{
}
CrackFreePrimitiveGenerator::~CrackFreePrimitiveGenerator() {}
unsigned int CrackFreePrimitiveGenerator::computeTriangleSize(bool bAddAdjacentEdges,
                        bool bAddDominantEdges,
                        bool bAddDominantPosition)
{
    return 3                                
        + (bAddAdjacentEdges ? 3 * 2 : 0)   
        + (bAddDominantEdges ? 3 * 2 : 0)   
        + (bAddDominantPosition ? 3 : 0);   
}
void CrackFreePrimitiveGenerator::mutateIndexBuffer( 
const MUintArray& originalBufferIndices, 
 
                        const float* positionBufferFloat, 
                        const float* uvBufferFloat, 
                        bool bAddAdjacentEdges,
                        bool bAddDominantEdges,
                        bool bAddDominantPosition,
                        void* indexData )
{
    unsigned int numTriVerts = originalBufferIndices.
length();
 
    EdgeMapping edges;
    {
        VertexFMap vertexMap;
        
        unsigned int vertexIndex = 0;
        while (vertexIndex < numTriVerts)
        {
            unsigned int faceVertexId0 = originalBufferIndices[vertexIndex++];
            unsigned int polyVertexId0 = vertexMap.getVertexId(VertexF(positionBufferFloat, faceVertexId0));
            unsigned int faceVertexId1 = originalBufferIndices[vertexIndex++];
            unsigned int polyVertexId1 = vertexMap.getVertexId(VertexF(positionBufferFloat, faceVertexId1));
            unsigned int faceVertexId2 = originalBufferIndices[vertexIndex++];
            unsigned int polyVertexId2 = vertexMap.getVertexId(VertexF(positionBufferFloat, faceVertexId2));
            edges.addTriangle(faceVertexId0, faceVertexId1, faceVertexId2, polyVertexId0, polyVertexId1, polyVertexId2);
            if(bAddDominantPosition && uvBufferFloat)
            {
                unsigned int uvIndex;
                uvIndex = faceVertexId0 * 2;
                edges.addPositionUV(faceVertexId0, polyVertexId0, uvBufferFloat[uvIndex], uvBufferFloat[uvIndex+1]);
                uvIndex = faceVertexId1 * 2;
                edges.addPositionUV(faceVertexId1, polyVertexId1, uvBufferFloat[uvIndex], uvBufferFloat[uvIndex+1]);
                uvIndex = faceVertexId2 * 2;
                edges.addPositionUV(faceVertexId2, polyVertexId2, uvBufferFloat[uvIndex], uvBufferFloat[uvIndex+1]);
            }
        }
    }
    unsigned int newTriId = 0;
    for(unsigned int triId = 0; triId < numTriVerts; )
    {
        unsigned int vertexId0 = originalBufferIndices[triId++];
        unsigned int vertexId1 = originalBufferIndices[triId++];
        unsigned int vertexId2 = originalBufferIndices[triId++];
            
            ((unsigned int*)indexData)[newTriId++] = vertexId0;
            ((unsigned int*)indexData)[newTriId++] = vertexId1;
            ((unsigned int*)indexData)[newTriId++] = vertexId2;
            
            if(bAddAdjacentEdges)
            {
                Edge adjacentEdge;
                
                edges.adjacentEdge(Edge(vertexId0, vertexId1), adjacentEdge);
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v1;
                
                edges.adjacentEdge(Edge(vertexId1, vertexId2), adjacentEdge);
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v1;
                
                edges.adjacentEdge(Edge(vertexId2, vertexId0), adjacentEdge);
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = adjacentEdge.v1;
            }
            
            if(bAddDominantEdges)
            {
                Edge dominantEdge;
                
                edges.dominantEdge(Edge(vertexId0, vertexId1), dominantEdge);
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v1;
                
                edges.dominantEdge(Edge(vertexId1, vertexId2), dominantEdge);
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v1;
                
                edges.dominantEdge(Edge(vertexId2, vertexId0), dominantEdge);
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v0;
                ((unsigned int*)indexData)[newTriId++] = dominantEdge.v1;
            }
            
            if(bAddDominantPosition)
            {
                unsigned int dominantVertexId;
                edges.dominantPosition(vertexId0, dominantVertexId);
                ((unsigned int*)indexData)[newTriId++] = dominantVertexId;
                edges.dominantPosition(vertexId1, dominantVertexId);
                ((unsigned int*)indexData)[newTriId++] = dominantVertexId;
                edges.dominantPosition(vertexId2, dominantVertexId);
                ((unsigned int*)indexData)[newTriId++] = dominantVertexId;
            }
        }
            ((unsigned short*)indexData)[newTriId++] = (unsigned short)vertexId0;
            ((unsigned short*)indexData)[newTriId++] = (unsigned short)vertexId1;
            ((unsigned short*)indexData)[newTriId++] = (unsigned short)vertexId2;
            
            if(bAddAdjacentEdges)
            {
                Edge adjacentEdge;
                
                edges.adjacentEdge(Edge(vertexId0, vertexId1), adjacentEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v1;
                
                edges.adjacentEdge(Edge(vertexId1, vertexId2), adjacentEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v1;
                
                edges.adjacentEdge(Edge(vertexId2, vertexId0), adjacentEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)adjacentEdge.v1;
            }
            
            if(bAddDominantEdges)
            {
                Edge dominantEdge;
                
                edges.dominantEdge(Edge(vertexId0, vertexId1), dominantEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v1;
                
                edges.dominantEdge(Edge(vertexId1, vertexId2), dominantEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v1;
                
                edges.dominantEdge(Edge(vertexId2, vertexId0), dominantEdge);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v0;
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantEdge.v1;
            }
            
            if(bAddDominantPosition)
            {
                unsigned int dominantVertexId;
                edges.dominantPosition(vertexId0, dominantVertexId);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantVertexId;
                edges.dominantPosition(vertexId1, dominantVertexId);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantVertexId;
                edges.dominantPosition(vertexId2, dominantVertexId);
                ((unsigned short*)indexData)[newTriId++] = (unsigned short)dominantVertexId;
            }
        }
    }
}
        int& primitiveStride) const
{
    
    for (
unsigned int ivb = 0; ivb < vertexBuffers.
count() && (positionBuffer == NULL || uvBuffer == NULL); ++ivb)
 
    {
            positionBuffer = currBuffer;
            uvBuffer = currBuffer;
    }
    if (positionBuffer == NULL)
        
    float* positionBufferFloat = (
float*)positionBuffer->
map();
 
    float* uvBufferFloat = NULL;
    if (uvBuffer)
        uvBufferFloat = (
float*)uvBuffer->
map();
    for (
int x = 0; x < sourceIndexBuffers.
length(); ++x)
 
    {
        if (sourceIndexBuffers[x]->componentType() != MHWRender::MComponentDataIndexing::kFaceVertex)
            continue;
        const MUintArray& originalBufferIndices = sourceIndexBuffers[x]->indices();
 
        unsigned int numTriVerts = originalBufferIndices.
length();
 
        unsigned int numTri = numTriVerts / 3;
        unsigned int triSize = computeTriangleSize(fAddAdjacentEdges, fAddDominantEdges, fAddDominantPosition);
        unsigned int bufferSize = numTri * triSize;
        void* indexData = indexBuffer.
acquire(bufferSize, 
true );
 
        if (indexData != NULL)
        {
            mutateIndexBuffer( originalBufferIndices, positionBufferFloat, uvBufferFloat, 
                               fAddAdjacentEdges, fAddDominantEdges, fAddDominantPosition,
        }
        if (positionBuffer) positionBuffer->
unmap();
 
        if (uvBuffer) uvBuffer->
unmap();
 
        indexBuffer.
commit(indexData);
        primitiveStride = triSize;
    }
    if (positionBuffer) positionBuffer->
unmap();
 
    if (uvBuffer) uvBuffer->
unmap();
 
}
{
    return new CrackFreePrimitiveGenerator(true , true , true );
}
{
    return new CrackFreePrimitiveGenerator(true , false , false );
}