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

#include "scenegraph.h"
// #include "xml_loader.h"
// #include "xml_writer.h"
#include "obj_loader.h"
// #include "ply_loader.h"
// #include "corona_loader.h"

namespace embree {
    // extern "C" RTCDevice g_device;

    void (*SceneGraph::opaque_geometry_destruction)(void *) = nullptr;

    Ref<SceneGraph::Node> SceneGraph::load(const FileName &filename, const bool singleObject) {
        if (toLowerCase(filename.ext()) == std::string("obj"))
            return loadOBJ(filename, false, singleObject);
            // else if (toLowerCase(filename.ext()) == std::string("ply" )) return loadPLY(filename);
            // else if (toLowerCase(filename.ext()) == std::string("xml" )) return loadXML(filename);
            // else if (toLowerCase(filename.ext()) == std::string("scn" )) return loadCorona(filename);
        else
            throw std::runtime_error("unknown scene format: " + filename.ext());
    }

    void SceneGraph::LightNode::calculateStatistics(Statistics& stat)
    {
        indegree++;
        if (indegree == 1) stat.numLights++;
    }

    void SceneGraph::GroupNode::print(std::ostream &cout, int depth) {
        cout << "GroupNode @ " << this << " { " << std::endl;
        tab(cout, depth + 1);
        cout << "closed = " << closed << std::endl;
        for (size_t i = 0; i < children.size(); i++) {
            tab(cout, depth + 1);
            cout << "child" << i << " = ";
            children[i]->print(cout, depth + 1);
        }
        tab(cout, depth);
        cout << "}" << std::endl;
    }

    void SceneGraph::MaterialNode::print(std::ostream &cout, int depth) {
        cout << "MaterialNode @ " << this << " { closed = " << closed << " }" << std::endl;
    }

    void SceneGraph::TriangleMeshNode::print(std::ostream &cout, int depth) {
        cout << "TriangleMeshNode @ " << this << " { closed = " << closed << " }" << std::endl;
    }

    void SceneGraph::LightNode::print(std::ostream& cout, int depth) {
        cout << "LightNode @ " << this << " { closed = " << closed << " }" << std::endl;
    }

    void SceneGraph::Node::calculateInDegree() {
        indegree++;
    }

    void SceneGraph::TriangleMeshNode::calculateInDegree() {
        indegree++;
        if (indegree == 1 && material != 0)
            material->calculateInDegree();
    }

    void SceneGraph::GroupNode::calculateInDegree() {
        indegree++;
        if (indegree == 1) {
            for (auto &c: children)
                c->calculateInDegree();
        }
    }

    bool SceneGraph::Node::calculateClosed(bool group_instancing) {
        assert(indegree);
        closed = true;
        hasLightOrCamera = false;
        return indegree == 1;
    }

    bool SceneGraph::GroupNode::calculateClosed(bool group_instancing) {
        assert(indegree);
        if (!closed) {
            closed = group_instancing;
            hasLightOrCamera = false;
            for (auto &c: children) {
                closed &= c->calculateClosed(group_instancing);
                hasLightOrCamera |= c->hasLightOrCamera;
            }
        }
        return closed && (indegree == 1);
    }

    bool SceneGraph::LightNode::calculateClosed(bool group_instancing)
    {
        assert(indegree);
        closed = true;
        hasLightOrCamera = true;
        return indegree == 1;
    }

    void SceneGraph::Node::resetInDegree() {
        closed = false;
        indegree--;
    }

    void SceneGraph::TriangleMeshNode::resetInDegree() {
        closed = false;
        if (indegree == 1 && material != 0)
            material->resetInDegree();
        indegree--;
    }

    void SceneGraph::GroupNode::resetInDegree() {
        closed = false;
        if (indegree == 1) {
            for (auto &c: children)
                c->resetInDegree();
        }
        indegree--;
    }

    void SceneGraph::TriangleMeshNode::verify() const {
        const size_t N = numVertices();
        if (normals.size() && normals.size() != positions.size())
            THROW_RUNTIME_ERROR("incompatible number of time steps");
        for (const auto &p: positions)
            if (p.size() != N)
                THROW_RUNTIME_ERROR("incompatible vertex array sizes");
        for (const auto &n: normals)
            if (n.size() && n.size() != N)
                THROW_RUNTIME_ERROR("incompatible vertex array sizes");
        if (texcoords.size() && texcoords.size() != N)
            THROW_RUNTIME_ERROR("incompatible vertex array sizes");
        for (auto &tri: triangles) {
            if (size_t(tri.v0) >= N || size_t(tri.v1) >= N || size_t(tri.v2) >= N)
                THROW_RUNTIME_ERROR("invalid triangle");
        }
    }

    bool test_location(const std::vector<avector<Vec3ff>> &in, ssize_t ipos, std::vector<avector<Vec3ff>> &out,
                       ssize_t opos) {
        if (opos < 0)
            return false;

        for (ssize_t i = ipos, j = opos; i < ipos + 4 && j < (ssize_t) out[0].size(); i++, j++) {
            for (size_t k = 0; k < in.size(); k++) {
                if (any(abs((vfloat4) in[k][i].m128 - (vfloat4) out[k][j].m128) > 0.01f * (vfloat4) max(
                                abs(in[k][i]), abs(out[k][j])).m128))
                    return false;
            }
        }
        return true;
    }

    std::pair<int, int> quad_index2(int p, int a0, int a1, int b0, int b1) {
        if (b0 == a0)
            return std::make_pair(p - 1, b1);
        else if (b0 == a1)
            return std::make_pair(p + 0, b1);
        else if (b1 == a0)
            return std::make_pair(p - 1, b0);
        else if (b1 == a1)
            return std::make_pair(p + 0, b0);
        else
            return std::make_pair(0, -1);
    }

    std::pair<int, int> quad_index3(int a0, int a1, int a2, int b0, int b1, int b2) {
        if (b0 == a0)
            return quad_index2(0, a2, a1, b1, b2);
        else if (b0 == a1)
            return quad_index2(1, a0, a2, b1, b2);
        else if (b0 == a2)
            return quad_index2(2, a1, a0, b1, b2);
        else if (b1 == a0)
            return quad_index2(0, a2, a1, b0, b2);
        else if (b1 == a1)
            return quad_index2(1, a0, a2, b0, b2);
        else if (b1 == a2)
            return quad_index2(2, a1, a0, b0, b2);
        else if (b2 == a0)
            return quad_index2(0, a2, a1, b0, b1);
        else if (b2 == a1)
            return quad_index2(1, a0, a2, b0, b1);
        else if (b2 == a2)
            return quad_index2(2, a1, a0, b0, b1);
        else
            return std::make_pair(0, -1);
    }

    struct SceneGraphFlattener {
        Ref<SceneGraph::Node> node;
        std::map<Ref<SceneGraph::Node>, Ref<SceneGraph::Node>> object_mapping;
        std::map<std::string, int> unique_id;

        SceneGraphFlattener(Ref<SceneGraph::Node> in, SceneGraph::InstancingMode instancing) {
            in->calculateInDegree();
            in->calculateClosed(instancing == SceneGraph::INSTANCING_GROUP);

            std::vector<Ref<SceneGraph::Node>> geometries;
            if (instancing != SceneGraph::INSTANCING_NONE) {
                if (instancing == SceneGraph::INSTANCING_FLATTENED)
                    convertFlattenedInstances(geometries, in);
                else if (instancing == SceneGraph::INSTANCING_MULTI_LEVEL)
                    convertMultiLevelInstances(geometries, in);
                else
                    convertInstances(geometries, in, one);
                convertLightsAndCameras(geometries, in, one);
            } else {
                convertGeometries(geometries, in, one);
                convertLightsAndCameras(geometries, in, one);
            }
            in->resetInDegree();

            node = new SceneGraph::GroupNode(geometries);
        }

        std::string makeUniqueID(std::string id) {
            if (id == "")
                id = "camera";
            std::map<std::string, int>::iterator i = unique_id.find(id);
            if (i == unique_id.end()) {
                unique_id[id] = 0;
                return id;
            } else {
                int n = ++unique_id[id];
                return id + "_" + toString(n);
            }
        }

        void convertLightsAndCameras(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node,
                                     const SceneGraph::Transformations &spaces) {
            if (!node->hasLightOrCamera)
                return;

            /*if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
              convertLightsAndCameras(group,xfmNode->child, spaces*xfmNode->spaces);
            }
            else */
            if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertLightsAndCameras(group, child, spaces);
            }
            // else if (Ref<SceneGraph::AnimatedLightNode> lightNode = node.dynamicCast<SceneGraph::AnimatedLightNode>()) {
            //   if (spaces.size() != 1) throw std::runtime_error("animated lights cannot get instantiated with a transform animation");
            //   group.push_back(lightNode->transform(spaces[0]).dynamicCast<SceneGraph::Node>());
            // }
            else if (Ref<SceneGraph::LightNode> lightNode = node.dynamicCast<SceneGraph::LightNode>()) {
                if (spaces.size() == 1)
                    group.push_back(lightNode->transform(spaces[0]).dynamicCast<SceneGraph::Node>());
                else {
                    // std::vector<Ref<SceneGraph::LightNode>> lights(spaces.size());
                    // for (size_t i=0; i<spaces.size(); i++)
                    //   lights[i] = lightNode->transform(spaces[i]);
                    //
                    // group.push_back(new SceneGraph::AnimatedLightNode(std::move(lights),spaces.time_range));
                }
            }
            // else if (Ref<SceneGraph::AnimatedPerspectiveCameraNode> cameraNode = node.dynamicCast<SceneGraph::AnimatedPerspectiveCameraNode>())
            // {
            //   if (spaces.size() != 1) throw std::runtime_error("animated cameras cannot get instantiated with a transform animation");
            //   group.push_back(new SceneGraph::AnimatedPerspectiveCameraNode(cameraNode,spaces[0],makeUniqueID(cameraNode->name)));
            // }
            // else if (Ref<SceneGraph::PerspectiveCameraNode> cameraNode = node.dynamicCast<SceneGraph::PerspectiveCameraNode>())
            // {
            //   if (spaces.size() == 1)
            //     group.push_back(new SceneGraph::PerspectiveCameraNode(cameraNode,spaces[0],makeUniqueID(cameraNode->name)));
            //   else
            //   {
            //     std::vector<Ref<SceneGraph::PerspectiveCameraNode>> cameras(spaces.size());
            //     for (size_t i=0; i<spaces.size(); i++)
            //       cameras[i] = new SceneGraph::PerspectiveCameraNode(cameraNode,spaces[i]);
            //
            //     group.push_back(new SceneGraph::AnimatedPerspectiveCameraNode(std::move(cameras),spaces.time_range,makeUniqueID(cameraNode->name)));
            //   }
            // }
        }

        void convertGeometries(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node,
                               const SceneGraph::Transformations &spaces) {
            /*if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
                convertGeometries(group, xfmNode->child, spaces * xfmNode->spaces);
            } else if (Ref<SceneGraph::MultiTransformNode> xfmNode = node.dynamicCast<
                SceneGraph::MultiTransformNode>()) {
                for (size_t i = 0; i < xfmNode->spaces.size(); ++i)
                    convertGeometries(group, xfmNode->child, spaces * xfmNode->spaces[i]);
            } else */if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertGeometries(group, child, spaces);
            } else if (Ref<SceneGraph::TriangleMeshNode> mesh = node.dynamicCast<SceneGraph::TriangleMeshNode>()) {
                group.push_back(new SceneGraph::TriangleMeshNode(mesh, spaces));
            } /*else if (Ref<SceneGraph::QuadMeshNode> mesh = node.dynamicCast<SceneGraph::QuadMeshNode>()) {
                group.push_back(new SceneGraph::QuadMeshNode(mesh, spaces));
            } else if (Ref<SceneGraph::GridMeshNode> mesh = node.dynamicCast<SceneGraph::GridMeshNode>()) {
                group.push_back(new SceneGraph::GridMeshNode(mesh, spaces));
            } else if (Ref<SceneGraph::SubdivMeshNode> mesh = node.dynamicCast<SceneGraph::SubdivMeshNode>()) {
                group.push_back(new SceneGraph::SubdivMeshNode(mesh, spaces));
            } else if (Ref<SceneGraph::HairSetNode> mesh = node.dynamicCast<SceneGraph::HairSetNode>()) {
                group.push_back(new SceneGraph::HairSetNode(mesh, spaces));
            } else if (Ref<SceneGraph::PointSetNode> mesh = node.dynamicCast<SceneGraph::PointSetNode>()) {
                group.push_back(new SceneGraph::PointSetNode(mesh, spaces));
            }*/
        }

        Ref<SceneGraph::Node> lookupGeometries(Ref<SceneGraph::Node> node) {
            if (object_mapping.find(node) == object_mapping.end()) {
                std::vector<Ref<SceneGraph::Node>> geometries;
                convertGeometries(geometries, node, one);
                object_mapping[node] = new SceneGraph::GroupNode(geometries);
            }

            return object_mapping[node];
        }

        void convertInstances(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node,
                              const std::vector<SceneGraph::Transformations> &spaces) {
           /*if (node->isClosed()) {
                //if (group.size() % 10000 == 0) std::cout << "." << std::flush;
                group.push_back(new SceneGraph::MultiTransformNode(spaces, lookupGeometries(node)));
            } else if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
                for (size_t i = 0; i < spaces.size(); ++i)
                    convertInstances(group, xfmNode->child, spaces[i] * xfmNode->spaces);
            } else if (Ref<SceneGraph::MultiTransformNode> xfmNode = node.dynamicCast<
                SceneGraph::MultiTransformNode>()) {
                convertInstances(group, xfmNode->child, spaces * xfmNode->spaces);
            } else */if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertInstances(group, child, spaces);
            }
        }

        void convertInstances(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node,
                              const SceneGraph::Transformations &spaces) {
            /*if (node->isClosed()) {
                //if (group.size() % 10000 == 0) std::cout << "." << std::flush;
                group.push_back(new SceneGraph::TransformNode(spaces, lookupGeometries(node)));
            } else if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
                convertInstances(group, xfmNode->child, spaces * xfmNode->spaces);
            } else if (Ref<SceneGraph::MultiTransformNode> xfmNode = node.dynamicCast<
                SceneGraph::MultiTransformNode>()) {
                convertInstances(group, xfmNode->child, spaces * xfmNode->spaces);
            } else */if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertInstances(group, child, spaces);
            }
        }

        void convertMultiLevelInstances(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node) {
            if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertMultiLevelInstances(group, child);
            } else if (node.dynamicCast<SceneGraph::TriangleMeshNode>()) {
                group.push_back(node);
            } /*else if (node.dynamicCast<SceneGraph::QuadMeshNode>()) {
                group.push_back(node);
            } else if (node.dynamicCast<SceneGraph::GridMeshNode>()) {
                group.push_back(node);
            } else if (node.dynamicCast<SceneGraph::SubdivMeshNode>()) {
                group.push_back(node);
            } else if (node.dynamicCast<SceneGraph::HairSetNode>()) {
                group.push_back(node);
            } else if (node.dynamicCast<SceneGraph::PointSetNode>()) {
                group.push_back(node);
            }*/ else if (object_mapping.find(node) != object_mapping.end()) {
                group.push_back(object_mapping[node]);
            } /*else if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
                auto new_node = new SceneGraph::TransformNode(xfmNode->spaces,
                                                              convertMultiLevelInstances(xfmNode->child));
                object_mapping[node] = new_node;
                group.push_back(new_node);
            } else if (Ref<SceneGraph::MultiTransformNode> xfmNode = node.dynamicCast<
                SceneGraph::MultiTransformNode>()) {
                auto new_node = new SceneGraph::MultiTransformNode(xfmNode->spaces,
                                                                   convertMultiLevelInstances(xfmNode->child));
                object_mapping[node] = new_node;
                group.push_back(new_node);
            }*/
        }

        Ref<SceneGraph::Node> convertMultiLevelInstances(const Ref<SceneGraph::Node> &node) {
            if (object_mapping.find(node) != object_mapping.end()) {
                return object_mapping[node];
            }

            std::vector<Ref<SceneGraph::Node>> group;
            convertMultiLevelInstances(group, node);
            auto new_node = new SceneGraph::GroupNode(group);
            object_mapping[node] = new_node;
            return new_node;
        }

        void convertFlattenedInstances(std::vector<Ref<SceneGraph::Node>> &group, const Ref<SceneGraph::Node> &node) {
            /*if (Ref<SceneGraph::TransformNode> xfmNode = node.dynamicCast<SceneGraph::TransformNode>()) {
                group.push_back(node);
            } else */if (Ref<SceneGraph::GroupNode> groupNode = node.dynamicCast<SceneGraph::GroupNode>()) {
                for (const auto &child: groupNode->children)
                    convertFlattenedInstances(group, child);
            }
        }
    };

    Ref<SceneGraph::Node> SceneGraph::flatten(Ref<Node> node, InstancingMode mode) {
        return SceneGraphFlattener(node, mode).node;
    }

    Ref<SceneGraph::GroupNode> SceneGraph::flatten(Ref<SceneGraph::GroupNode> node, SceneGraph::InstancingMode mode) {
        return flatten(node.dynamicCast<SceneGraph::Node>(), mode).dynamicCast<SceneGraph::GroupNode>();
    }
}