scene.cpp

#include "scene.h"
#include "ray.h"
#include "rng.h"
#include "vector.h"
#include <thread>
#include <vector>
#include <atomic>

static const float supersampling[4] = {0.125f, 0.375f, 0.625f, 0.875f};

CollisionInfo Scene::intersectRay(Vector3& origin, Vector3& direction, Vector3& inv_dir) const {
    CollisionInfo closestInfo;
    closestInfo.distance = INFINITY;

    for (const auto& shape : shapes) {
        CollisionInfo info = shape->collide_ray(origin, direction, inv_dir);
        if (info.hit && info.distance < closestInfo.distance) {
            closestInfo = info;
            closestInfo.object = shape;
        }
    }

    return closestInfo;
}

void Scene::rotateCamera(Vector3 euler) {
    // Convert degrees to radians
    euler = euler * (M_PI / 180.0);

    // Rotation matrices
    Vector3 right = camera_cframe[0];
    Vector3 up = camera_cframe[1];
    Vector3 forward = camera_cframe[2];

    // Yaw (around up)
    if (euler.y != 0) {
        double cosY = std::cos(euler.y);
        double sinY = std::sin(euler.y);
        Vector3 new_forward = forward * cosY + right * sinY;
        Vector3 new_right = right * cosY - forward * sinY;
        camera_cframe[0] = new_right.normalize();
        camera_cframe[2] = new_forward.normalize();
    }

    // Pitch (around right)
    if (euler.x != 0) {
        double cosX = std::cos(euler.x);
        double sinX = std::sin(euler.x);
        Vector3 new_forward = forward * cosX - up * sinX;
        Vector3 new_up = up * cosX + forward * sinX;
        camera_cframe[1] = new_up.normalize();
        camera_cframe[2] = new_forward.normalize();
    }

    // Roll (around forward)
    if (euler.z != 0) {
        double cosZ = std::cos(euler.z);
        double sinZ = std::sin(euler.z);
        Vector3 new_right = right * cosZ + up * sinZ;
        Vector3 new_up = up * cosZ - right * sinZ;
        camera_cframe[0] = new_right.normalize();
        camera_cframe[1] = new_up.normalize();
    }
}

void Scene::setCameraRotation(Vector3 euler) {
    // Reset to default orientation
    camera_cframe[0] = Vector3(1, 0, 0);
    camera_cframe[1] = Vector3(0, 1, 0);
    camera_cframe[2] = Vector3(0, 0, -1);

    // Apply rotations
    rotateCamera(euler);
}

std::vector<Vector3> Scene::render() {
    std::vector<Vector3> framebuffer(width * height, Vector3(0, 0, 0));
    PCG32 rng(rand(), rand());

    double fovtan = std::tan((fov*(M_PI/180))/2);
    double aspect_ratio = (double)width / (double)height;
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            Vector3 pixel_color(0, 0, 0);
            for (int ss = 0; ss < 16; ss++) {
                bool no_hit = false;
                Vector3 ndc(
                    (x + supersampling[ss / 4]) / width,
                    (y + supersampling[ss % 4]) / height,
                    0
                );
                Vector3 direction = (
                    camera_cframe[0] * (2 * ndc.x - 1) * aspect_ratio * fovtan +
                    camera_cframe[1] * (1 - 2 * ndc.y) * fovtan +
                    camera_cframe[2]
                ).normalize();
                Vector3 inv_dir = 1 / direction;
                CollisionInfo firstHitCollision = intersectRay(camera_position, direction, inv_dir);

                Ray ray(camera_position, direction, rng, *this);
                ray.firstHitCollision = firstHitCollision;
                for (int s = 0; s < samples_per_pixel / 16; s++) {
                    pixel_color = pixel_color + ray.trace();
                    if (ray.hit_nothing) {
                        no_hit = true;
                        break;
                    }
                    ray.reset();
                }
                if (no_hit) {
                    if (background) pixel_color += background->sample(direction) * (samples_per_pixel / 16.0 - 1);
                    continue;
                }
            }
            pixel_color = pixel_color / samples_per_pixel;
            if (tone_map == TONEMAP_REINHARD) {
                pixel_color = pixel_color.reinhard_tone_map();
            } else if (tone_map == TONEMAP_ACES) {
                pixel_color = pixel_color.aces_tone_map();
            }
            framebuffer[y * width + x] = pixel_color.linear_to_srgb();
        }
    }

    return framebuffer;
}

void Scene::renderTile(int start_x, int start_y, int tile_width, int tile_height, std::vector<Vector3>& framebuffer) {
    double fovtan = std::tan((fov*(M_PI/180))/2);
    double aspect_ratio = (double)width / (double)height;
    PCG32 rng(rand(), rand());

    for (int y = start_y; y < start_y + tile_height && y < height; y++) {
        for (int x = start_x; x < start_x + tile_width && x < width; x++) {
            Vector3 pixel_color(0, 0, 0);
            for (int ss = 0; ss < 16; ss++) {
                bool no_hit = false;

                double aspect_ratio = (double)width / (double)height;
                Vector3 ndc(
                    (x + supersampling[ss / 4]) / width,
                    (y + supersampling[ss % 4]) / height,
                    0
                );
                Vector3 direction = (
                    camera_cframe[0] * (2 * ndc.x - 1) * aspect_ratio * fovtan +
                    camera_cframe[1] * (1 - 2 * ndc.y) * fovtan +
                    camera_cframe[2]
                ).normalize();
                Vector3 inv_dir = 1 / direction;
                CollisionInfo firstHitCollision = intersectRay(camera_position, direction, inv_dir);

                Ray ray(camera_position, direction, rng, *this);
                ray.firstHitCollision = firstHitCollision;

                for (int s = 0; s < samples_per_pixel / 16; s++) {
                    pixel_color = pixel_color + ray.trace();
                    if (ray.hit_nothing) {
                        no_hit = true;
                        break;
                    }
                    ray.reset();
                }

                if (no_hit) {
                    if (background) pixel_color += background->sample(direction) * (samples_per_pixel / 16.0 - 1);
                    continue;
                }
            }

            pixel_color = pixel_color / samples_per_pixel;
            if (tone_map == TONEMAP_REINHARD) {
                pixel_color = pixel_color.reinhard_tone_map();
            } else if (tone_map == TONEMAP_ACES) {
                pixel_color = pixel_color.aces_tone_map();
            }
            framebuffer[y * width + x] = pixel_color.linear_to_srgb();
        }
    }
}

std::vector<Vector3> Scene::renderThreaded(int num_threads, int tile_size) {
    std::vector<Vector3> framebuffer(width * height, Vector3(0, 0, 0));
    std::vector<std::thread> threads;
    std::atomic<int> next_tile(0);

    int tiles_x = (width + tile_size - 1) / tile_size;
    int tiles_y = (height + tile_size - 1) / tile_size;
    int total_tiles = tiles_x * tiles_y;

    auto worker = [&]() {
        while (true) {
            int tile_index = next_tile.fetch_add(1);
            if (tile_index >= total_tiles) break;

            int tx = tile_index % tiles_x;
            int ty = tile_index / tiles_x;

            int start_x = tx * tile_size;
            int start_y = ty * tile_size;

            renderTile(start_x, start_y, tile_size, tile_size, framebuffer);
        }
    };

    for (int i = 0; i < num_threads; i++) {
        threads.emplace_back(worker);
    }

    for (auto& t : threads) t.join();
    return framebuffer;
}

std::vector<CollisionInfo> Scene::renderCollisionInfo() {
    double fovtan = std::tan((fov*(M_PI/180))/2);
    double aspect_ratio = (double)width / (double)height;
    std::vector<CollisionInfo> output(width * height);

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            Vector3 pixel_color(0, 0, 0);
            Vector3 ndc(
                (x + 0.5) / width,
                (y + 0.5) / height,
                0
            );
            Vector3 direction = (
                camera_cframe[0] * (2 * ndc.x - 1) * aspect_ratio * fovtan +
                camera_cframe[1] * (1 - 2 * ndc.y) * fovtan +
                camera_cframe[2]
            ).normalize();
            Vector3 inv_dir = 1 / direction;

            output[y * width + x] = intersectRay(camera_position, direction, inv_dir);
        }
    }

    return output;
}

void Scene::renderDebug(std::vector<Vector3>* normals, std::vector<double>* depth, std::vector<Vector3>* albedo) {
    if (!normals && !depth && !albedo) return; // Early exit

    std::vector<CollisionInfo> collision = Scene::renderCollisionInfo();

    if (normals) {
        normals->clear();
        normals->resize(width * height);
    }
    if (depth) {
        depth->clear();
        depth->resize(width * height);
    }
    if (albedo) {
        albedo->clear();
        albedo->resize(width * height);
    }

    for (int i = 0; i < width * height; i++) {
        CollisionInfo& info = collision[i];
        if (normals) (*normals)[i] = info.normal;
        if (depth) (*depth)[i] = info.distance;
        int u;
        int v;
        if (info.material.has_texture) {
            u = ((static_cast<int>(floor(info.uv.x * info.material.textures_width)) % info.material.textures_width) + info.material.textures_width) % info.material.textures_width;
            v = ((static_cast<int>(floor(info.uv.y * info.material.textures_height)) % info.material.textures_height) + info.material.textures_height) % info.material.textures_height;
        }
        if (albedo) (*albedo)[i] = info.material.has_texture ? info.material.texture[v][u] : info.material.color;
    }
}