// Copyright 2009-2021 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once
#include <sys/platform.h>
#include <sys/sysinfo.h>
#include <sys/alloc.h>

#include <sys/ref.h>
#include <sys/vector.h>
#include <math/vec2.h>
#include <math/vec3.h>
#include <math/vec4.h>
#include <math/bbox.h>
#include <math/lbbox.h>
#include <math/affinespace.h>
#include <sys/filename.h>
#include <sys/estring.h>
#include <lexers/tokenstream.h>
#include <lexers/streamfilters.h>
#include <lexers/parsestream.h>

#include <sstream>
#include <vector>
#include <memory>
#include <map>
#include <set>
#include <deque>

#include "lights.h"
#include "embree4/rtcore.h"
RTC_NAMESPACE_USE

namespace embree {
    struct Material;

    namespace SceneGraph {
        struct Node;
        struct MaterialNode;
        struct TriangleMeshNode;

        struct Statistics {
            Statistics()
                : numTriangleMeshes(0), numTriangles(0), numTriangleBytes(0),
                  numQuadMeshes(0), numQuads(0), numQuadBytes(0),
                  numSubdivMeshes(0), numPatches(0), numSubdivBytes(0),
                  numCurveSets(0), numCurves(0), numCurveBytes(0),
                  numGridMeshNodes(0), numGrids(0), numGridBytes(0),
                  numPointSets(0), numPoints(0), numPointBytes(0),
                  numTransformNodes(0),
                  numTransformedObjects(0),
                  numLights(0),
                  numCameras(0),
                  numMaterials(0) {
            }

            void print();

            size_t numTriangleMeshes;
            size_t numTriangles;
            size_t numTriangleBytes;

            size_t numQuadMeshes;
            size_t numQuads;
            size_t numQuadBytes;

            size_t numSubdivMeshes;
            size_t numPatches;
            size_t numSubdivBytes;

            size_t numCurveSets;
            size_t numCurves;
            size_t numCurveBytes;

            size_t numGridMeshNodes;
            size_t numGrids;
            size_t numGridBytes;

            size_t numPointSets;
            size_t numPoints;
            size_t numPointBytes;

            size_t numTransformNodes;
            size_t numTransformedObjects;

            size_t numLights;
            size_t numCameras;
            size_t numMaterials;
        };

        extern void (*opaque_geometry_destruction)(void*);

        Ref<Node> load(const FileName& fname, bool singleObject = false);

        struct Node : public RefCount {
            Node(bool closed = false)
                : indegree(0), closed(closed), hasLightOrCamera(false), id(-1), geometry(nullptr) {
            }

            Node(const std::string& name)
                : name(name), indegree(0), closed(false), id(-1), geometry(nullptr) {
            }

            ~Node() {
                if (opaque_geometry_destruction)
                    opaque_geometry_destruction(geometry);
            }

            /* prints scenegraph */
            virtual void print(std::ostream& cout, int depth = 0) = 0;

            /* sets material */
            virtual void setMaterial(Ref<MaterialNode> material) {
            };

            /* calculates the number of parent nodes pointing to this node */
            virtual void calculateInDegree();

            /* calculates for each node if its subtree is closed, indegrees have to be calculated first */
            virtual bool calculateClosed(bool group_instancing);

            /* resets the number of parent nodes pointing to this node */
            virtual void resetInDegree();

            /* checks if the node is closed */
            __forceinline bool isClosed() const { return closed; }

            /* calculates bounding box of node */
            virtual BBox3fa bounds() const {
                return empty;
            }

            virtual BBox3fa bounds(size_t i) const {
                return empty;
            }

            /* calculates linear bounding box of node */
            virtual LBBox3fa lbounds() const {
                return empty;
            }

            virtual LBBox3fa lbounds(size_t i) const {
                return empty;
            }

            /* calculates number of primitives */
            virtual size_t numPrimitives() const {
                return 0;
            }

        public:
            std::string fileName; // when set to some filename the exporter references this file
            std::string name; // name of this node
            size_t indegree; // number of nodes pointing to us
            bool closed; // determines if the subtree may represent an instance
            bool hasLightOrCamera;
            unsigned int id;
            void* geometry;
        };

        struct Transformations {
            __forceinline Transformations() {
            }

            __forceinline Transformations(OneTy)
                : time_range(0.0f, 1.0f) {
                spaces.push_back(one);
            }

            __forceinline Transformations(const BBox1f& time_range, size_t N)
                : time_range(time_range), spaces(N) {
            }

            __forceinline Transformations(const AffineSpace3fa& space)
                : time_range(0.0f, 1.0f) {
                spaces.push_back(space);
            }

            __forceinline Transformations(const AffineSpace3fa& space0, const AffineSpace3fa& space1)
                : time_range(0.0f, 1.0f) {
                spaces.push_back(space0);
                spaces.push_back(space1);
            }

            __forceinline Transformations(const avector<AffineSpace3ff>& spaces)
                : time_range(0.0f, 1.0f), spaces(spaces) { assert(spaces.size()); }

            __forceinline size_t size() const {
                return spaces.size();
            }

            __forceinline AffineSpace3ff& operator[](const size_t i) { return spaces[i]; }
            __forceinline const AffineSpace3ff& operator[](const size_t i) const { return spaces[i]; }

            BBox3fa bounds(const BBox3fa& cbounds) const {
                BBox3fa r = empty;
                for (size_t i = 0; i < spaces.size(); i++)
                    r.extend(xfmBounds(spaces[i], cbounds));
                return r;
            }

            LBBox3fa lbounds(const LBBox3fa& cbounds) const {
                assert(spaces.size());
                if (spaces.size() == 1) {
                    return LBBox3fa(xfmBounds(spaces[0], cbounds.bounds0),
                                    xfmBounds(spaces[0], cbounds.bounds1));
                } else {
                    avector<BBox3fa> bounds(spaces.size());
                    for (size_t i = 0; i < spaces.size(); i++) {
                        const float f = float(i) / float(spaces.size() - 1);
                        bounds[i] = xfmBounds(spaces[i], cbounds.interpolate(f));
                    }
                    return LBBox3fa(bounds);
                }
            }

            void add(const Transformations& other) {
                for (size_t i = 0; i < other.size(); i++) spaces.push_back(other[i]);
            }

            static __forceinline bool isIdentity(AffineSpace3ff const& M, bool q) {
                if (M.l.vx.x != 1.f) return false;
                if (M.l.vx.y != 0.f) return false;
                if (M.l.vx.z != 0.f) return false;
                if (q && M.l.vx.w != 0.f) return false;

                if (M.l.vy.x != 0.f) return false;
                if (M.l.vy.y != 1.f) return false;
                if (M.l.vy.z != 0.f) return false;
                if (q && M.l.vy.w != 0.f) return false;

                if (M.l.vz.x != 0.f) return false;
                if (M.l.vz.y != 0.f) return false;
                if (M.l.vz.z != 1.f) return false;
                if (q && M.l.vz.w != 0.f) return false;

                if (M.p.x != 0.f) return false;
                if (M.p.y != 0.f) return false;
                if (M.p.z != 0.f) return false;
                if (q && M.p.w != 1.f) return false;

                return true;
            }

            static __forceinline AffineSpace3ff mul(AffineSpace3ff const& M0, AffineSpace3ff const& M1, bool q0,
                                                    bool q1, bool& q) {
                q = false;
                if (isIdentity(M0, q0)) {
                    q = q1;
                    return M1;
                }
                if (isIdentity(M1, q1)) {
                    q = q0;
                    return M0;
                }

                // simple case non of the transformations is a quaternion
                if (q0 == false && q1 == false) {
                    return M0 * M1;
                } else if (q0 == true && q1 == true) {
                    std::cout <<
                            "warning: cannot multiply two quaternion decompositions. will convert to regular transforms and multiply"
                            << std::endl;
                    return quaternionDecompositionToAffineSpace(M0) * quaternionDecompositionToAffineSpace(M1);
                } else if (q0 == true && q1 == false) {
                    AffineSpace3fa S;
                    Quaternion3f Q;
                    Vec3fa T;
                    quaternionDecomposition(M0, T, Q, S);
                    S = S * AffineSpace3fa(M1);
                    if (S.l.vx.y != 0.f || S.l.vx.z != 0 || S.l.vy.z != 0)
                        std::cout <<
                                "warning: cannot multiply quaternion and general transformation matrix. will ignore lower diagonal"
                                << std::endl;
                    q = true;
                    return quaternionDecomposition(T, Q, S);
                } else {
                    if (M0.l.vx.y != 0.f || M0.l.vx.z != 0 || M0.l.vy.z != 0 || M0.l.vy.x != 0.f || M0.l.vz.x != 0 || M0
                        .l.vz.y != 0)
                        std::cout <<
                                "warning: cannot multiply general transformation matrix and quaternion. will only consider translation and diagonal as scale factors"
                                << std::endl;
                    AffineSpace3ff M = M1;
                    M.l.vx.y += M0.p.x;
                    M.l.vx.z += M0.p.y;
                    M.l.vy.z += M0.p.z;
                    M.l.vx.x *= M0.l.vx.x;
                    M.l.vy.y *= M0.l.vy.y;
                    M.l.vz.z *= M0.l.vz.z;
                    q = true;
                    return M;
                }
            }

            friend __forceinline Transformations operator*(const Transformations& a, const Transformations& b) {
                if (a.size() == 1) {
                    Transformations c(intersect(a.time_range, b.time_range), b.size());
                    for (size_t i = 0; i < b.size(); i++)
                        c[i] = mul(a[0], b[i], a.quaternion, b.quaternion,
                                   c.quaternion);
                    return c;
                } else if (b.size() == 1) {
                    Transformations c(intersect(a.time_range, b.time_range), a.size());
                    c.quaternion = a.quaternion || b.quaternion;
                    for (size_t i = 0; i < a.size(); i++)
                        c[i] = mul(a[i], b[0], a.quaternion, b.quaternion,
                                   c.quaternion);
                    return c;
                } else if (a.size() == b.size()) {
                    Transformations c(intersect(a.time_range, b.time_range), a.size());
                    c.quaternion = a.quaternion || b.quaternion;
                    for (size_t i = 0; i < a.size(); i++)
                        c[i] = mul(a[i], b[i], a.quaternion, b.quaternion,
                                   c.quaternion);
                    return c;
                } else
                    THROW_RUNTIME_ERROR("number of transformations does not match");
            }

            friend __forceinline std::vector<Transformations> operator*(const std::vector<Transformations>& a,
                                                                        const Transformations& b) {
                std::vector<Transformations> result;
                for (size_t i = 0; i < a.size(); ++i) {
                    result.push_back(a[i] * b);
                }
                return result;
            }

            friend __forceinline std::vector<Transformations> operator*(const Transformations& a,
                                                                        const std::vector<Transformations>& b) {
                return b * a;
            }

            friend __forceinline std::vector<Transformations> operator*(const std::vector<Transformations>& a,
                                                                        const std::vector<Transformations>& b) {
                if (a.size() != b.size())
                    THROW_RUNTIME_ERROR("number of transformations does not match");

                std::vector<Transformations> result;
                for (size_t i = 0; i < a.size(); ++i) {
                    result.push_back(a[i] * b[i]);
                }

                return result;
            }

            AffineSpace3ff interpolate(const float gtime) const {
                assert(time_range.lower == 0.0f && time_range.upper == 1.0f);
                if (spaces.size() == 1) return spaces[0];

                /* calculate time segment itime and fractional time ftime */
                const int time_segments = int(spaces.size() - 1);
                const float time = gtime * float(time_segments);
                const int itime = clamp(int(floor(time)), 0, time_segments - 1);
                const float ftime = time - float(itime);
                return lerp(spaces[itime + 0], spaces[itime + 1], ftime);
            }

        public:
            BBox1f time_range;
            avector<AffineSpace3ff> spaces;
            bool quaternion = false;
        };

        inline std::vector<avector<Vec3fa> > transformMSMBlurVec3faBuffer(
            const std::vector<avector<Vec3fa> >& positions_in, const Transformations& spaces) {
            std::vector<avector<Vec3fa> > positions_out;
            const size_t num_time_steps = positions_in.size();
            assert(num_time_steps);
            const size_t num_vertices = positions_in[0].size();

            /* if we have only one set of vertices, use transformation to generate more vertex sets */
            if (num_time_steps == 1) {
                for (size_t i = 0; i < spaces.size(); i++) {
                    avector<Vec3fa> verts(num_vertices);
                    for (size_t j = 0; j < num_vertices; j++) {
                        verts[j] = xfmPoint((AffineSpace3fa) spaces[i], positions_in[0][j]);
                    }
                    positions_out.push_back(std::move(verts));
                }
            }
            /* otherwise transform all vertex sets with interpolated transformation */
            else {
                for (size_t t = 0; t < num_time_steps; t++) {
                    float time = num_time_steps > 1 ? float(t) / float(num_time_steps - 1) : 0.0f;
                    const AffineSpace3ff space = spaces.interpolate(time);
                    avector<Vec3fa> verts(num_vertices);
                    for (size_t i = 0; i < num_vertices; i++) {
                        verts[i] = xfmPoint((AffineSpace3fa) space, positions_in[t][i]);
                    }
                    positions_out.push_back(std::move(verts));
                }
            }
            return positions_out;
        }

        template<typename Vertex>
        std::vector<avector<Vertex> > transformMSMBlurNormalBuffer(const std::vector<avector<Vertex> >& normals_in,
                                                                   const Transformations& spaces) {
            if (normals_in.size() == 0)
                return normals_in;

            std::vector<avector<Vertex> > normals_out;
            const size_t num_time_steps = normals_in.size();
            const size_t num_vertices = normals_in[0].size();

            /* if we have only one set of vertices, use transformation to generate more vertex sets */
            if (num_time_steps == 1) {
                for (size_t i = 0; i < spaces.size(); i++) {
                    avector<Vertex> norms(num_vertices);
                    for (size_t j = 0; j < num_vertices; j++) {
                        norms[j] = xfmNormal((AffineSpace3fa) spaces[i], normals_in[0][j]);
                    }
                    normals_out.push_back(std::move(norms));
                }
            }
            /* otherwise transform all vertex sets with interpolated transformation */
            else {
                for (size_t t = 0; t < num_time_steps; t++) {
                    float time = num_time_steps > 1 ? float(t) / float(num_time_steps - 1) : 0.0f;
                    const AffineSpace3ff space = spaces.interpolate(time);
                    avector<Vertex> norms(num_vertices);
                    for (size_t i = 0; i < num_vertices; i++) {
                        norms[i] = xfmNormal((AffineSpace3fa) space, normals_in[t][i]);
                    }
                    normals_out.push_back(std::move(norms));
                }
            }
            return normals_out;
        }

        struct GroupNode : public Node {
            GroupNode(const size_t N = 0) {
                children.resize(N);
            }

            GroupNode(std::vector<Ref<Node> >& children)
                : children(children) {
            }

            size_t size() const {
                return children.size();
            }

            void add(const Ref<Node>& node) {
                if (node)
                    children.push_back(node);
            }

            void set(const size_t i, const Ref<Node>& node) {
                children[i] = node;
            }

            Ref<Node> child(size_t i) const {
                return children[i];
            }

            virtual BBox3fa bounds() const {
                BBox3fa b = empty;
                for (auto& c: children)
                    b.extend(c->bounds());
                return b;
            }

            virtual LBBox3fa lbounds() const {
                LBBox3fa b = empty;
                for (auto& c: children)
                    b.extend(c->lbounds());
                return b;
            }

            virtual size_t numPrimitives() const {
                size_t n = 0;
                for (auto& child: children)
                    n += child->numPrimitives();
                return n;
            }

            virtual void setMaterial(Ref<MaterialNode> material) {
                for (auto& child: children)
                    child->setMaterial(material);
            }

            virtual void print(std::ostream& cout, int depth);

            virtual void calculateInDegree();

            virtual bool calculateClosed(bool group_instancing);

            virtual void resetInDegree();

        public:
            std::vector<Ref<Node> > children;
        };

        struct MaterialNode : public Node {
            ALIGNED_STRUCT_USM_(16)

            MaterialNode(const std::string& name = "")
                : Node(name) {
            }

            virtual Material* material() = 0;

            virtual void print(std::ostream& cout, int depth);
        };

        /*! Mesh. */
        struct TriangleMeshNode : public Node {
            typedef Vec3fa Vertex;

            struct Triangle {
            public:
                Triangle() {
                }

                Triangle(int v0, int v1, int v2)
                    : v0(v0), v1(v1), v2(v2) {
                }

            public:
                int v0, v1, v2;
            };

        public:
            TriangleMeshNode(const avector<Vertex>& positions_in,
                             const avector<Vertex>& normals_in,
                             const std::vector<Vec2f>& texcoords,
                             const std::vector<Triangle>& triangles,
                             Ref<MaterialNode> material)
                : Node(true), time_range(0.0f, 1.0f), texcoords(texcoords), triangles(triangles), material(material) {
                positions.push_back(positions_in);
                normals.push_back(normals_in);
            }

            TriangleMeshNode(Ref<MaterialNode> material, const BBox1f time_range = BBox1f(0, 1),
                             size_t numTimeSteps = 0)
                : Node(true), time_range(time_range), material(material) {
                for (size_t i = 0; i < numTimeSteps; i++)
                    positions.push_back(avector<Vertex>());
            }

            TriangleMeshNode(Ref<SceneGraph::TriangleMeshNode> imesh, const Transformations& spaces)
                : Node(true),
                  time_range(imesh->time_range),
                  positions(transformMSMBlurVec3faBuffer(imesh->positions, spaces)),
                  normals(transformMSMBlurNormalBuffer(imesh->normals, spaces)),
                  texcoords(imesh->texcoords), triangles(imesh->triangles), material(imesh->material) {
            }

            virtual void setMaterial(Ref<MaterialNode> material) {
                this->material = material;
            }

            virtual BBox3fa bounds() const {
                BBox3fa b = empty;
                for (const auto& p: positions)
                    for (auto& x: p)
                        b.extend(x);
                return b;
            }

            virtual LBBox3fa lbounds() const {
                avector<BBox3fa> bboxes(positions.size());
                for (size_t t = 0; t < positions.size(); t++) {
                    BBox3fa b = empty;
                    for (auto& x: positions[t])
                        b.extend(x);
                    bboxes[t] = b;
                }
                return LBBox3fa(bboxes);
            }

            virtual size_t numPrimitives() const {
                return triangles.size();
            }

            size_t numVertices() const {
                assert(positions.size());
                return positions[0].size();
            }

            size_t numTimeSteps() const {
                return positions.size();
            }

            size_t numBytes() const {
                return numPrimitives() * sizeof(Triangle) + numVertices() * numTimeSteps() * sizeof(Vertex);
            }

            void verify() const;

            virtual void print(std::ostream& cout, int depth);

            virtual void calculateInDegree();

            virtual void resetInDegree();

        public:
            BBox1f time_range;
            std::vector<avector<Vertex> > positions;
            std::vector<avector<Vertex> > normals;
            std::vector<Vec2f> texcoords;
            std::vector<Triangle> triangles;
            Ref<MaterialNode> material;
        };

        struct LightNode : public Node {
            virtual void print(std::ostream& cout, int depth);

            virtual void calculateStatistics(Statistics& stat);

            virtual bool calculateClosed(bool group_instancing);

            virtual LightType getType() const = 0;

            virtual Ref<LightNode> transform(const AffineSpace3fa& space) const = 0;

            virtual Ref<LightNode> lerp(const Ref<LightNode>& light1_in, float f) const = 0;

            virtual Ref<LightNode> get(float time) const = 0;
        };

        template<typename Light>
        struct LightNodeImpl : public LightNode {
            ALIGNED_STRUCT_(16);

            LightNodeImpl(const Light& light)
                : light(light) {
            }

            virtual LightType getType() const {
                return light.getType();
            }

            virtual Ref<LightNode> transform(const AffineSpace3fa& space) const {
                return new LightNodeImpl(light.transform(space));
            }

            virtual Ref<LightNode> get(float time) const {
                return (LightNode *) this;
            }

            virtual Ref<LightNode> lerp(const Ref<LightNode>& light1_in, float f) const {
                const Ref<LightNodeImpl<Light> > light1 = light1_in.dynamicCast<LightNodeImpl<Light> >();
                assert(light1);
                return new LightNodeImpl(Light::lerp(light, light1->light, f));
            }

            Light light;
        };

        enum InstancingMode {
            INSTANCING_NONE, INSTANCING_GEOMETRY, INSTANCING_GROUP, INSTANCING_FLATTENED, INSTANCING_MULTI_LEVEL
        };

        Ref<Node> flatten(Ref<Node> node, InstancingMode mode);

        Ref<GroupNode> flatten(Ref<GroupNode> node, InstancingMode mode);
    }
}

#include "materials.h"