#include "instancedmapgamestate.hpp"
#include "findfilecallbackdsmap.hpp"

#include <osgDB/ReadFile>
#include <osg/MatrixTransform>
#include <iostream>

namespace siege
{
    const std::string instanceVert =

        "uniform mat4 instanceModelMatrix[MAX_INSTANCES];\n"

        "smooth out vec2 texCoord;\n"

        "void main()\n"
        "{\n"
        "mat4 _instanceModelMatrix = instanceModelMatrix[gl_InstanceID];\n"
        "gl_Position = gl_ModelViewProjectionMatrix * _instanceModelMatrix * gl_Vertex;\n"
        "texCoord = gl_MultiTexCoord0.xy;\n"
        "}\n";

    const std::string instanceFrag =

        "#version 150 compatibility\n"

        "uniform sampler2D colorTexture;\n"
        "smooth in vec2 texCoord;\n"

        "void main()\n"
        "{\n"
        "vec4 textureColor = texture2D(colorTexture, texCoord);\n"
        "gl_FragColor = vec4(textureColor.rgb, textureColor.a);\n"
        "}\n";

    struct ModelInstance
    {
        ModelInstance() : mObjectId(-1) {}
        osg::Matrix matrix;
        uint32_t mObjectId;
    };

    typedef std::map< osg::ref_ptr<osg::Node>, std::vector<ModelInstance> > ModelInstanceMap;

    struct ComputeInstancedBoundingBoxCallback : public osg::Drawable::ComputeBoundingBoxCallback
    {
        public:

            ComputeInstancedBoundingBoxCallback(osg::ref_ptr<osg::Uniform> instanceMatrices): mInstanceMatrices(instanceMatrices)
            {
            }

            virtual osg::BoundingBox computeBound(const osg::Drawable& drawable) const
            {
                osg::BoundingBox bounds;
                const osg::Geometry* geometry = dynamic_cast<const osg::Geometry*>(&drawable);

                if (!geometry)
                    return bounds;

                const osg::Vec3Array* vertices = dynamic_cast<const osg::Vec3Array*>(geometry->getVertexArray());

                for (unsigned int i = 0; i < mInstanceMatrices->getNumElements(); ++i)
                {
                    osg::Matrixd matrix;
                    mInstanceMatrices->getElement(i, matrix);

                    for (auto it = vertices->begin(); it != vertices->end(); ++it)
                    {
                        bounds.expandBy(*it * matrix);
                    }
                }

                return bounds;
            }

        private:
            osg::ref_ptr<osg::Uniform> mInstanceMatrices;

    };

    class ConvertToDrawInstanced : public osg::NodeVisitor
    {
    public:

        ConvertToDrawInstanced(unsigned numInstances, const osg::BoundingBox &bbox, bool optimize) : mNumInstances(numInstances), mOptimize(optimize)
        {
            setTraversalMode(TRAVERSE_ALL_CHILDREN);
            setNodeMaskOverride(~0);
        }

        void apply(osg::Geode &geode)
        {
            for (unsigned d = 0; d < geode.getNumDrawables(); ++d)
            {
                osg::Geometry *geom = geode.getDrawable(d)->asGeometry();
                if (geom != nullptr)
                {
                    if (mOptimize)
                    {
                        geom->setUseDisplayList(false);
                        geom->setUseVertexBufferObjects(true);
                    }

                    for (unsigned int p = 0; p < geom->getNumPrimitiveSets(); ++p)
                    {
                        osg::PrimitiveSet *ps = geom->getPrimitiveSet(0);
                        ps->setNumInstances(mNumInstances);
                        mPrimitiveSets.push_back(ps);
                    }
                }
            }

            traverse(geode);
        }

    protected:

        unsigned mNumInstances;
        bool mOptimize;
        std::list<osg::PrimitiveSet *> mPrimitiveSets;
    };

    osg::ref_ptr<osg::Node> InstancedMapGameState::createHardwareInstancedNode(osg::Node *node, const std::vector <ModelInstance> &instances, int start, int end)
    {
        osg::Geode *geode = new osg::Geode(*node->asGeode(), osg::CopyOp::DEEP_COPY_ALL);

        // Enable the geometry to be instanced
        ConvertToDrawInstanced cdi(end - start, osg::BoundingBox(), true);
        geode->accept(cdi);

        // Create a uniform to pass to the shader
        osg::ref_ptr<osg::Uniform> instanceMatrixUniform = new osg::Uniform(osg::Uniform::FLOAT_MAT4, "instanceModelMatrix");

        // Iterate over each instance to get matrix information
        for (unsigned m = start, j = 0; m < end; ++m, ++j)
        {
            const osg::Matrix &mat = instances[m].matrix;

            // Assign matrix to the uniform
            instanceMatrixUniform->setElement(j, mat);
        }

        geode->getOrCreateStateSet()->addUniform(instanceMatrixUniform);

        geode->getDrawable(0)->setComputeBoundingBoxCallback(new ComputeInstancedBoundingBoxCallback(instanceMatrixUniform));

        return geode;
    }

    bool InstancedMapGameState::convertGraphToUseDrawInstanced(osg::Group *parent)
    {
        // Store our hardware limit and actual matrix limit
        int32_t hardwareLimit = app().getMaxNumUniforms();
        int32_t maxInstanceMatrices = (hardwareLimit - 16) / 16;

        // Define an instance map to hold all instances of models in a map
        ModelInstanceMap models;

        // Iterate over the group
        for (unsigned i = 0; i < parent->getNumChildren(); ++i)
        {
            osg::MatrixTransform *mt = dynamic_cast<osg::MatrixTransform *> (parent->getChild(i));

            if (mt != nullptr)
            {
                osg::Node *n = mt->getChild(0);

                ModelInstance instance;
                instance.matrix = mt->getMatrix();

                // Store the instance
                models[n].push_back(instance);
            }
        }

        std::cout << "There are " << models.size() << " unique models in the map" << std::endl;

        // Clear out current scene graph
        parent->removeChildren(0, parent->getNumChildren());

        // Create our shader program
        osg::ref_ptr<osg::Program> program = new osg::Program();

        // Add preprocessor definition to the instance vertex shader
        std::stringstream preprocessor;
        preprocessor << "#version 150 compatibility\n #define MAX_INSTANCES " << maxInstanceMatrices << "\n";
        std::cout << preprocessor.str().c_str() << std::endl;
        preprocessor << instanceVert;

        // Add our vertex and fragment shaders
        program->addShader(new osg::Shader(osg::Shader::VERTEX, preprocessor.str()));
        program->addShader(new osg::Shader(osg::Shader::FRAGMENT, instanceFrag));

        // Assign our shader to the parent of the scene graph
        parent->getOrCreateStateSet()->setAttribute(program, osg::StateAttribute::ON);

        // Iterate over each model and attempt to create a group out of it for instancing
        for (ModelInstanceMap::iterator i = models.begin(); i != models.end(); ++i)
        {
            osg::Node *node = i->first.get();
            std::vector<ModelInstance> &instances = i->second;

            // Check to make sure our instance size can fit in one group
            if (instances.size() <= maxInstanceMatrices)
            {
                // Add the new geode to the scene graph
                parent->addChild(createHardwareInstancedNode(node, instances, 0, instances.size()));
            }
            else
            {
                std::cout << "WARNING: Instance size is: " << instances.size() << " hardware limit is: " << maxInstanceMatrices << " - ATTEMPTING TO SPLIT NODES!" << std::endl;

#if 0
                unsigned int numGeodes = std::ceil((uint32_t)instances.size() / maxInstanceMatrices);
                for (uint32_t g = 0; g < numGeodes; ++g)
                {
                    unsigned int start = g * maxInstanceMatrices;
                    unsigned int end = std::min((unsigned int)instances.size(), (start + maxInstanceMatrices));

                    parent->addChild(createHardwareInstancedNode(node, instances, start, end));
                }
#endif
            }
        }

        return true;
    }

    InstancedMapGameState::InstancedMapGameState(Application & app) : GameState(app)
    {
    }

    InstancedMapGameState::~InstancedMapGameState()
    {
    }

    void InstancedMapGameState::onEnter()
    {
        const std::string regionFileName = "town_center.dsmap";

        std::cout << "InstancedMapGameState::onEnter: " << regionFileName << std::endl;

        // add a stats handler so we can see our scene graph stats
        app().addEventHandler(new osgViewer::StatsHandler);

        // load the region
        osg::ref_ptr<osg::Node> regionA = osgDB::readNodeFile(regionFileName, useDSMAPOption());

        std::cout << "Finished load Town Center" << std::endl;

        // Check if we support instance drawing
        if (!app().supportsDrawInstanced())
        {
            // add the switch to the application
            app().setSceneData(regionA.get());

            return;
        }

        // Try to instance the map
        convertGraphToUseDrawInstanced(regionA->asGroup());

        app().setSceneData(regionA.get());
    }
}