#include <vector>
#include <list>
#include <set>
#include <limits>
#include <algorithm>
#include <cmath>
#include <cassert>

#include <osg/Geometry>
#include <osg/Array>
#include <osg/Notify>
#include <osgAnimation/MorphGeometry>
#include <osgUtil/MeshOptimizers>

#include "GeometryUniqueVisitor"
#include "TriangleMeshGraph"


// Smoothing steps:
//
// 1. compute the vertex/triangles graph
// 2. compute triangle normals (vertexTriangles::addTriangle)
// 3. determine *piecewise* one-ring for each *unique* vertex (TriangleMeshGraph::vertexOneRing)
//    Each piece of the one-ring contains triangles that are neighbors and do not share a sharp edge
// 4. for each one-ring piece sum the triangle normals (TriangleMeshSmoother::computeVertexNormals)
// 5. if the vertex has been processed already: duplicate and update triangles
//    otherwise set the normal
//
// **triangle normals are normalized but weighted by their area when cumulated over the ring**

class TriangleMeshSmoother {
public:
    enum SmoothingMode {
        recompute      = 1 << 0,
        diagnose       = 1 << 1,
        smooth_flipped = 1 << 2,
        smooth_all     = 1 << 3
    };


    class DuplicateVertex : public osg::ArrayVisitor {
    public:
        unsigned int _i;
        unsigned int _end;

        DuplicateVertex(unsigned int i): _i(i), _end(i)
        {}

        template <class ARRAY>
        void apply_imp(ARRAY& array) {
            _end = array.size();
            array.push_back(array[_i]);
        }

        virtual void apply(osg::ByteArray& array) { apply_imp(array); }
        virtual void apply(osg::ShortArray& array) { apply_imp(array); }
        virtual void apply(osg::IntArray& array) { apply_imp(array); }
        virtual void apply(osg::UByteArray& array) { apply_imp(array); }
        virtual void apply(osg::UShortArray& array) { apply_imp(array); }
        virtual void apply(osg::UIntArray& array) { apply_imp(array); }
        virtual void apply(osg::FloatArray& array) { apply_imp(array); }
        virtual void apply(osg::DoubleArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2Array& array) { apply_imp(array); }
        virtual void apply(osg::Vec3Array& array) { apply_imp(array); }
        virtual void apply(osg::Vec4Array& array) { apply_imp(array); }

        virtual void apply(osg::Vec2bArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3bArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4bArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2sArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3sArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4sArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2iArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3iArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4iArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2dArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3dArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4dArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2ubArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3ubArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4ubArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2usArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3usArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4usArray& array) { apply_imp(array); }

        virtual void apply(osg::Vec2uiArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec3uiArray& array) { apply_imp(array); }
        virtual void apply(osg::Vec4uiArray& array) { apply_imp(array); }

        virtual void apply(osg::MatrixfArray& array) { apply_imp(array); }
        virtual void apply(osg::MatrixdArray& array) { apply_imp(array); }
    };

public:
    TriangleMeshSmoother(osg::Geometry& geometry, float creaseAngle, bool comparePosition=false, int mode=diagnose):
        _geometry(geometry),
        _creaseAngle(creaseAngle),
        _graph(0),
        _mode(mode)
    {
        if(!_geometry.getVertexArray() || !_geometry.getVertexArray()->getNumElements()) {
            return;
        }

        osgUtil::SharedArrayOptimizer deduplicator;
        deduplicator.findDuplicatedUVs(geometry);

        // duplicate shared arrays as it isn't safe to duplicate vertices when arrays are shared.
        if (geometry.containsSharedArrays()) {
            geometry.duplicateSharedArrays();
        }

        if(!_geometry.getNormalArray() || _geometry.getNormalArray()->getNumElements() != _geometry.getVertexArray()->getNumElements()) {
            _geometry.setNormalArray(new osg::Vec3Array(_geometry.getVertexArray()->getNumElements()), osg::Array::BIND_PER_VERTEX);
        }

        // build a unifier to consider deduplicated vertex indices
        _graph = new TriangleMeshGraph(_geometry, comparePosition);

        unsigned int nbTriangles = 0;
        for(unsigned int i = 0 ; i < _geometry.getNumPrimitiveSets() ; ++ i) {
            osg::PrimitiveSet* primitive = _geometry.getPrimitiveSet(i);

            if(!primitive || !primitive->getNumIndices()) {
                continue;
            }
            else if(primitive->getMode() > osg::PrimitiveSet::TRIANGLES) {
                OSG_INFO << "[smoother] Cannot smooth geometry '" << _geometry.getName()
                         << "' due to not tessellated primitives" << std::endl;
                return;
            }
            else if(primitive->getMode() == osg::PrimitiveSet::TRIANGLES) {
                nbTriangles += primitive->getNumIndices() / 3;
            }
        }
        _triangles.reserve(nbTriangles);

        // collect all buffers that are BIND_PER_VERTEX for eventual vertex duplication
        addArray(_geometry.getVertexArray());
        addArray(_geometry.getColorArray());
        addArray(_geometry.getSecondaryColorArray());
        addArray(_geometry.getFogCoordArray());
        for(unsigned int i = 0; i < _geometry.getNumTexCoordArrays(); ++ i) {
           addArray(_geometry.getTexCoordArray(i));
        }
        for(unsigned int i = 0; i < _geometry.getNumVertexAttribArrays(); ++ i) {
           addArray(_geometry.getVertexAttribArray(i));
        }

        switch(_mode) {
            case recompute:
                computeVertexNormals();
                break;
            case smooth_all:
                smoothVertexNormals(true, true);
                break;
            case smooth_flipped:
                smoothVertexNormals(true, false);
                break;
            case diagnose:
                smoothVertexNormals(false, false);
                break;
        };

        // deduplicate UVs array that were only shared within the geometry
        deduplicator.deduplicateUVs(geometry);
    }

    ~TriangleMeshSmoother() {
        if(_graph) {
            delete _graph;
        }
    }

protected:
    unsigned int duplicateVertex(unsigned int index) {
        DuplicateVertex duplicate(index);
        for(ArrayVector::iterator array = _vertexArrays.begin(); array != _vertexArrays.end(); ++ array) {
            (*array)->accept(duplicate);
        }
#if 0
        OSG_INFO << "[normals] [[TriangleMeshSmoother]] vertex " << index
                 << " duplicated => " << duplicate._end << std::endl;
#endif
        _graph->add(duplicate._end, index);
        return duplicate._end;
    }

    void smoothVertexNormals(bool fix=true, bool force=false) {
        _vertexArrays.clear(); // make sure we do not change vertex count
        bool flipped = false;

        osg::Vec3Array* normals = dynamic_cast<osg::Vec3Array*>(_geometry.getNormalArray());
        for(unsigned int index = 0 ; index < _geometry.getVertexArray()->getNumElements() ; ++ index) {
            std::vector<IndexVector> oneRing = _graph->vertexOneRing(_graph->unify(index), _creaseAngle);
            osg::Vec3f smoothedNormal(0.f, 0.f, 0.f);

            // sum normals for each cluster in the one-ring
            for(std::vector<IndexVector>::iterator cluster = oneRing.begin() ; cluster != oneRing.end() ; ++ cluster) {
                smoothedNormal += cumulateTriangleNormals(*cluster);
            }

            float length = smoothedNormal.normalize();
            if(length > 0.) {
                if(force || smoothedNormal * normals->at(index) < 1.e-6) {
                    flipped = true;
                    if(fix) {
                        (*normals)[index] = smoothedNormal;
                    }
                }
            }
        }

        if(flipped) {
            OSG_WARN << std::endl << "Warning: [smoothVertexNormals] [[normals]] Geometry '" << _geometry.getName() << "' ";
            switch(_mode) {
                case diagnose:
                    OSG_WARN << "has some flipped normals; please check that the shading is correct" << std::endl;
                    OSG_WARN << "Monitor: normal.invalid" << std::endl;
                    break;
                case smooth_flipped:
                    OSG_WARN << "has some flipped normals that have been fixed" << std::endl;
                    OSG_WARN << "Monitor: normal.smooth_flipped" << std::endl;
                    break;
                case smooth_all:
                    OSG_WARN << "normals have all been smoothed" << std::endl;
                    OSG_WARN << "Monitor: normal.smooth_all" << std::endl;
                    break;
            }
        }
    }

    void computeVertexNormals() {
        osg::Vec3Array* normals = new osg::Vec3Array(osg::Array::BIND_PER_VERTEX,
                                                     _geometry.getVertexArray()->getNumElements());
        addArray(normals);

        for(unsigned int i = 0 ; i < normals->getNumElements() ; ++ i) {
            (*normals)[i].set(0.f, 0.f, 0.f);
        }

        for(VertexIterator uniqueIndex = _graph->begin() ; uniqueIndex != _graph->end() ; ++ uniqueIndex) {
            unsigned int index = uniqueIndex->_index;
            std::set<unsigned int> processed;

            std::vector<IndexVector> oneRing = _graph->vertexOneRing(index, _creaseAngle);
            for(std::vector<IndexVector>::iterator cluster = oneRing.begin() ; cluster != oneRing.end() ; ++ cluster) {
                osg::Vec3f clusterNormal = cumulateTriangleNormals(*cluster);
                clusterNormal.normalize();

                std::set<unsigned int> duplicates;
                for(IndexVector::const_iterator tri = cluster->begin() ; tri != cluster->end() ; ++ tri) {
                    const Triangle& triangle = _graph->triangle(*tri);

                    if(_graph->unify(triangle.v1()) == index) {
                        duplicates.insert(triangle.v1());
                    }
                    else if(_graph->unify(triangle.v2()) == index) {
                        duplicates.insert(triangle.v2());
                    }
                    else if(_graph->unify(triangle.v3()) == index) {
                        duplicates.insert(triangle.v3());
                    }
                }

                for(std::set<unsigned int>::iterator vertex = duplicates.begin() ; vertex != duplicates.end() ; ++ vertex) {
                    if(processed.find(*vertex) == processed.end()) {
                        // vertex not yet processed
                        (*normals)[*vertex] = clusterNormal;
                        processed.insert(*vertex);
                    }
                    else {
                        // vertex already processed in a previous cluster: need to duplicate
                        unsigned int duplicate = duplicateVertex(*vertex);
                        replaceVertexIndexInTriangles(*cluster, *vertex, duplicate);
                        (*normals)[duplicate] = clusterNormal;

                        processed.insert(duplicate);
                    }
                }
            }
        }

        _geometry.setNormalArray(normals, osg::Array::BIND_PER_VERTEX);
        updateGeometryPrimitives();

        OSG_WARN << std::endl <<"Warning: [computeVertexNormals] [[normals]] Geometry '" << _geometry.getName()
                 << "' normals have been recomputed" << std::endl;
        OSG_WARN << "Monitor: normal.recompute" << std::endl;
    }

    osg::Vec3f cumulateTriangleNormals(const IndexVector& triangles) const {
        osg::Vec3f normal;
        normal.set(0.f, 0.f, 0.f);
        for(IndexVector::const_iterator triangle = triangles.begin() ; triangle != triangles.end() ; ++ triangle) {
            const Triangle& t = _graph->triangle(*triangle);
            normal += (t._normal * t._area);
        }
        return normal;
    }

    void replaceVertexIndexInTriangles(const IndexVector& triangles, unsigned int oldIndex, unsigned int newIndex) {
        for(IndexVector::const_iterator tri = triangles.begin() ; tri != triangles.end() ; ++ tri) {
            Triangle& triangle = _graph->triangle(*tri);
            if(triangle.v1() == oldIndex) {
                triangle.v1() = newIndex;
            }
            else if(triangle.v2() == oldIndex) {
                triangle.v2() = newIndex;
            }
            else if(triangle.v3() == oldIndex) {
                triangle.v3() = newIndex;
            }
        }
    }

    void addArray(osg::Array* array) {
        if (array && array->getBinding() == osg::Array::BIND_PER_VERTEX) {
            _vertexArrays.push_back(array);
        }
    }

    void updateGeometryPrimitives() {
        osg::Geometry::PrimitiveSetList primitives;
        for(unsigned int i = 0 ; i < _geometry.getNumPrimitiveSets() ; ++ i) {
            osg::PrimitiveSet* primitive = _geometry.getPrimitiveSet(i);
            if(primitive && primitive->getMode() < osg::PrimitiveSet::TRIANGLES) {
                primitives.push_back(primitive);
            }
        }

        osg::DrawElementsUInt* triangles = new osg::DrawElementsUInt(osg::PrimitiveSet::TRIANGLES);
        for(unsigned int i = 0 ; i < _graph->getNumTriangles() ; ++ i) {
            const Triangle& triangle = _graph->triangle(i);
            triangles->push_back(triangle.v1());
            triangles->push_back(triangle.v2());
            triangles->push_back(triangle.v3());
        }
        primitives.push_back(triangles);

        _geometry.setPrimitiveSetList(primitives);
    }


    osg::Geometry& _geometry;
    float _creaseAngle;
    TriangleMeshGraph* _graph;
    TriangleVector _triangles;
    ArrayVector _vertexArrays;
    int _mode; // smooth or recompute normals
};



class SmoothNormalVisitor : public GeometryUniqueVisitor {
public:
    SmoothNormalVisitor(float creaseAngle, bool comparePosition=false):
        GeometryUniqueVisitor("SmoothNormalVisitor"),
        _creaseAngle(creaseAngle),
        _comparePosition(comparePosition)
    {}

    void process(osg::Geometry& geometry) {
        if(!geometry.getNormalArray()) {
            TriangleMeshSmoother(geometry, _creaseAngle, _comparePosition, TriangleMeshSmoother::recompute);
        }
        else {
            TriangleMeshSmoother(geometry, _creaseAngle, _comparePosition, TriangleMeshSmoother::diagnose);
        }
    }

    void process(osgAnimation::MorphGeometry& morphGeometry) {
        TriangleMeshSmoother(morphGeometry, 0, true, TriangleMeshSmoother::smooth_all);
        osgAnimation::MorphGeometry::MorphTargetList targets = morphGeometry.getMorphTargetList();
        for(osgAnimation::MorphGeometry::MorphTargetList::iterator target = targets.begin() ; target != targets.end() ; ++ target) {
            // check normal orientation using the same primitives as parent geometry
            osg::Geometry::PrimitiveSetList& primitives = target->getGeometry()->getPrimitiveSetList();
            target->getGeometry()->setPrimitiveSetList(morphGeometry.getPrimitiveSetList());

            TriangleMeshSmoother(*target->getGeometry(), 0, true, TriangleMeshSmoother::smooth_all);

            target->getGeometry()->setPrimitiveSetList(primitives);
        }
    }

protected:
    float _creaseAngle;
    bool _comparePosition;
};