#include "Scene.h"
#include "../thirdparty/OBJ_Loader.h"
#include "Sphere.h"
#include <QApplication>
#include <QMatrix4x4>
#include <QQuaternion>

raytry::Scene::Scene() {
    background.emplace_back(ground);
    for (int i = 0; i < spheres.size(); ++i) {
        float x = -5.f + static_cast<float>(i);
        float spec = static_cast<float>(i) / static_cast<float>(spheres.size() - 1);
        auto mat = RenderMaterial{{197, 52, 27}};
        mat.specularIntensity = spec;
        spheres[i] = Sphere(camera.pos() + camera.fwd() * 5 + QVector3D(x, 0.f, 0.f), 0.4f, mat);
        background.emplace_back(spheres[i]);
    }
    const auto &triangleTransformation = camera.pos() + camera.fwd() * 10;

    const auto &triangleBase = camera.pos() + camera.fwd() * 10;
    triangles[0] = Triangle(QVector3D{triangleBase[0] - 3, triangleBase[1], 0.0},
                            QVector3D{triangleBase[0] + 3, triangleBase[1], 0.0},
                            QVector3D{triangleBase[0], triangleBase[1], 6.0});
    background.emplace_back(triangles[0]);

    QMatrix4x4 transformation{};
    transformation.translate(camera.pos() + (camera.fwd() * 5));
    transformation.rotate(QQuaternion::fromAxisAndAngle(1.0f, 0.0f, 0.0f, 90.0f));
    transformation.rotate(QQuaternion::fromAxisAndAngle(0.0f, 1.0f, 0.0f, 90.0f));

    std::setlocale(LC_ALL, "C");// the obj parser uses locale-specific number format...
    objl::Loader objLoader{};
    if (!objLoader.LoadFile("newell_teaset/teapot.obj")) {
        qWarning() << "Could not load teapot..";
    } else {
        auto mat = RenderMaterial{{37, 96, 185}};
        mat.transparency = 0.2f;
        mat.reflectivity = 0.5f;
        teapotMesh.reserve(teapotMesh.size() + (objLoader.LoadedIndices.size() / 3));

        auto iter = objLoader.LoadedIndices.begin();
        while (iter != objLoader.LoadedIndices.end()) {
            auto a = objLoader.LoadedVertices[*iter++];
            auto b = objLoader.LoadedVertices[*iter++];
            auto c = objLoader.LoadedVertices[*iter++];
            QVector3D qta{a.Position.X, a.Position.Y, a.Position.Z};
            QVector3D qtb{b.Position.X, b.Position.Y, b.Position.Z};
            QVector3D qtc{c.Position.X, c.Position.Y, c.Position.Z};
            Triangle& tri = teapotMesh.emplace_back(transformation.map(qta), transformation.map(qtb), transformation.map(qtc));
            tri.setMaterial(mat);
        }
    }
    qInfo() << "Loaded teapot with" << teapotMesh.size() << "triangles.";

    foregroundTree = KD::Tree::build(teapotMesh);
}

std::list<std::pair<std::reference_wrapper<const raytry::RenderObject>, float>> raytry::Scene::findIntersections(const raytry::Ray &ray) {
    std::list<std::pair<std::reference_wrapper<const RenderObject>, float>> matches{};
    if (!foregroundTree.getNodes().empty()) {
        const QVector3D invDirection{1 / ray.directionVec.x(), 1 / ray.directionVec.y(), 1 / ray.directionVec.z()};
        const std::optional<std::pair<QVector3D, QVector3D>> &treeHit = foregroundTree.getBounds().fastslabtest(ray, invDirection);
        if (treeHit) {
            // KD!
            const KD::Node *node = foregroundTree.getNodes().data();
            float tMin = treeHit->first[node->getAxis()];
            float tMax = treeHit->second[node->getAxis()];
            unsigned int currentNodeIdx = 0;
            std::list<std::tuple<unsigned int, float, float>> nodes{};
            while (node != nullptr) {
                if (node->isLeaf()) {
                    assert(node->farIdx() >= 0);
                    assert(node->farIdx() < foregroundTree.getLeafs().size());
                    for (const auto &item: foregroundTree.getLeafs()[node->farIdx()].triangles) {
                        const auto distance = item->intersects(ray);
                        ++triangleIntersections;
                        if (distance) {
                            bool emplaced = false;
                            const RenderObject &triangle = *item;
                            for (auto mpos = matches.begin(); !emplaced && mpos != matches.end(); ++mpos) {
                                if ((*mpos).second > distance.value()) {
                                    matches.emplace(mpos, triangle, distance.value());
                                    emplaced = true;
                                }
                            }
                            if (!emplaced) {
                                matches.emplace_back(triangle, distance.value());
                            }
                        }
                    }

                    if (nodes.empty()) {
                        node = nullptr;
                    } else {
                        std::tie(currentNodeIdx, tMin, tMax) = nodes.front();
                        assert(currentNodeIdx < foregroundTree.getNodes().size());
                        node = &foregroundTree.getNodes()[currentNodeIdx];
                        nodes.pop_front();
                    }
                } else {
                    const auto tPlane = (node->val() - ray.originVec[node->getAxis()]) * invDirection[node->getAxis()];

                    unsigned int firstChild, secondChild;
                    int belowFirst = (ray.originVec[node->getAxis()] < node->val()) ||
                                     (ray.originVec[node->getAxis()] == node->val() && ray.directionVec[node->getAxis()] <= 0);
                    if (belowFirst) {
                        firstChild = currentNodeIdx + 1;
                        secondChild = currentNodeIdx + node->farIdx();
                    } else {
                        firstChild = currentNodeIdx + node->farIdx();
                        secondChild = currentNodeIdx + 1;
                    }

                    if (tPlane > tMax || tPlane <= 0) {
                        currentNodeIdx = firstChild;
                    } else if (tPlane < tMin) {
                        currentNodeIdx = secondChild;
                    } else {
                        nodes.emplace_back(secondChild, tPlane, tMax);
                        currentNodeIdx = firstChild;
                        tMax = tPlane;
                    }
                    node = &foregroundTree.getNodes()[currentNodeIdx];
                }
            }
        }
    }

    for (RenderObject &obj: background) {
        const auto distance = obj.intersects(ray);
        if (distance) {
            bool emplaced = false;
            for (auto mpos = matches.begin(); !emplaced && mpos != matches.end(); ++mpos) {
                if ((*mpos).second > distance.value()) {
                    matches.emplace(mpos, obj, distance.value());
                    emplaced = true;
                }
            }
            if (!emplaced) {
                matches.emplace_back(obj, distance.value());
            }
        }
    }
    return matches;
}

void raytry::Scene::rayTrace(const raytry::Ray &ray, unsigned int depth, QColor& color) {
    if (depth > maxDepth) {
        color = QColor{0, 0, 0};
    } else {
        auto matches = findIntersections(ray);
        if (matches.empty()) {
            color = backgroundColor;
        } else {
            auto [object, distance] = matches.front();
            QVector3D hitpoint = ray.originVec + (ray.directionVec * distance);
            const QVector3D &normal = object.get().calculateNormal(hitpoint);
            assert(normal.normalized() == normal);
            QVector3D toLight = pointLight - hitpoint;
            float lightDistance = toLight.length();
            toLight.normalize();
            const Ray &shadowRay{hitpoint, toLight};
            auto shadowMatches = findIntersections(shadowRay);
            float lightSourceIntensity = 1.0f;
            if (!shadowMatches.empty()) {
                for (const auto &occlusion: shadowMatches) {
                    if (occlusion.second > lightDistance) {
                        break;
                    }
                    lightSourceIntensity *= occlusion.first.get().material().transparency;
                }
            }
            const RenderMaterial &material = object.get().material();
            assert(globalIllumination >= 0 && globalIllumination <= 1);
            color = mixColor(material.ambientColor, material.ambientIntensity * globalIllumination);

            const float &diffusionModifier = std::max(0.0f, QVector3D::dotProduct(normal, toLight));
            color = addColors(color, mixColor(mixColor(mixColor(material.diffuseColor, diffusionModifier), material.diffuseIntensity), lightSourceIntensity));

            if (material.reflectivity > 0) {
                const Ray &reflectionRay = material.reflectionRay(ray, hitpoint, normal);
                QColor reflectionColor{};
                rayTrace(reflectionRay, depth + 1u, reflectionColor);
                color = addColors(mixColor(color, 1 - material.reflectivity), mixColor(reflectionColor, material.reflectivity));
            }

            const QVector3D &fromLight = (hitpoint - pointLight).normalized();
            const QVector3D &lightReflection = fromLight - (2 * QVector3D::dotProduct(fromLight, normal) * normal);
            const float &specularModifier = std::pow(std::max(0.0f, QVector3D::dotProduct(lightReflection, -ray.directionVec)), material.specularFalloff);
            color = addColors(color, mixColor(mixColor(mixColor(material.specularColor, specularModifier), material.specularIntensity), lightSourceIntensity));

            if (material.transparency > 0) {
                const std::optional<Ray> &optionalRefractionRay = material.refractionRay(ray, hitpoint, normal);
                QColor refractionColor{backgroundColor};
                if (optionalRefractionRay) {
                    rayTrace(optionalRefractionRay.value(), depth + 1u, refractionColor);
                }
                color = addColors(mixColor(color, 1 - material.transparency), mixColor(refractionColor, material.transparency));
            }
        }
    }
}

void raytry::Scene::render(raytry::ViewerWindow &window, std::function<void()> updateImage) {
    size_t rayC = 0;
    QImage &img = window.getDisplayImage();
    unsigned int pixels = img.width() * img.height();
    camera.for_each_ray(img, [&](int x, int y, Ray r) {
        ++rayC;
        if (rayC % (pixels / 20) == 0) {
            qInfo() << "Raytracing #" << rayC << "/" << pixels << "(" << rayC * 100 / pixels << "%)";
        }

        QColor color{};
        rayTrace(r, 0, color);
        img.setPixelColor(x, y, color);

        if (rayC % (pixels / 1000) == 0) {
            updateImage();
        }
    });
    qInfo() << "All pixels computed. Tested" << triangleIntersections << "ray-triangle intersections.";
}