ray.cpp

#include <cmath>
#include "shapes.h"
#include "vector.h"
#include "ray.h"
#include "scene.h"
#include "rng.h"
#include "lights.h"
#include <fstream>
#include <iostream>
#include <thread>
#include <vector>
#include <atomic>

static constexpr double inv_pi = 1 / M_PI;

Vector3 cosine_weighted_distribution(Vector3 normal, PCG32& rng) {
    float u = rng.random_float();
    float v = rng.random_float();

    // Cosine-weighted hemisphere sampling
    float r = sqrt(u);
    float theta = 2 * M_PI * v;

    float x = r * cos(theta);
    float y = r * sin(theta);
    float z = sqrt(1.0f - u);

    // Create a coordinate system (Frisvad method)
    Vector3 w = normal;
    Vector3 u_vec, v_vec;
    if (fabs(w.z) < 0.999f) {
        float a = 1.0f / (1.0f + w.z);
        float b = -w.x * w.y * a;
        u_vec = Vector3(1.0f - w.x * w.x * a, b, -w.x);
        v_vec = Vector3(b, 1.0f - w.y * w.y * a, -w.y);
    } else {
        u_vec = Vector3(1, 0, 0);
        v_vec = Vector3(0, 1, 0);
    }

    return (u_vec * x + v_vec * y + w * z).normalize();
}

int Ray::bounce() {
    if (bounces == 0) {
        info = firstHitCollision;
    } else {
        info = scene.intersectRay(origin, direction, inv_dir);
    }
    if (info.hit) {
        // Backface check
        bool is_inside_object = false;
        if (info.normal.dot(direction) > 0) {
            info.normal = -info.normal;
            is_inside_object = true;
        }
        origin = info.point + info.normal * 0.001; // Offset to avoid self-intersection

        // Read texture maps if available
        if (info.material.has_texture || info.material.has_roughness_map || info.material.has_emission_map || info.material.has_specular_map || info.material.has_refraction_map) {
            // Clamp UVs to [0, texture_size] based on texture size
            int u = ((static_cast<int>(floor(info.uv.x * info.material.textures_width)) % info.material.textures_width) + info.material.textures_width) % info.material.textures_width;
            int 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 (info.material.has_texture) {
                info.material.color = info.material.texture[v][u];
            }
            if (info.material.has_roughness_map) {
                info.material.roughness = info.material.roughness_map[v][u].x;
            }
            if (info.material.has_emission_map) {
                info.material.emission_color = info.material.emission_map[v][u];
            }
            if (info.material.has_specular_map) {
                info.material.specular_color = info.material.specular_map[v][u];
            }
            if (info.material.has_refraction_map) {
                info.material.refraction = info.material.refraction_map[v][u].x;
            }
        }

        // Normalize probabilities
        double sum = info.material.roughness + info.material.specular_probability + info.material.refraction;
        if (sum > 1) {
            info.material.roughness /= sum;
            info.material.specular_probability /= sum;
            info.material.refraction /= sum;
        }

        double rand = rng.random_float();
        double current_probability = info.material.roughness;
        bool is_diffuse = rand < current_probability;
        current_probability += info.material.specular_probability;
        bool is_specular = rand < current_probability;
        current_probability += info.material.refraction;
        bool is_refractive = rand < current_probability;
        bool is_reflective = !is_diffuse && !is_specular && !is_refractive;

        is_specular &= !is_diffuse;
        is_reflective &= !is_diffuse && !is_specular;

        if (is_diffuse) {
            direction = cosine_weighted_distribution(info.normal, rng);
            has_hit_diffuse = true;
        } else if (is_specular || is_reflective) {
            // Calculate reflection direction
            direction -= info.normal * 2.0 * direction.dot(info.normal);
        } else if (is_refractive) {
            double refractive_index_ratio;
            if (is_inside_object) {
                refractive_index_ratio = info.material.refractive_index;
            } else {
                refractive_index_ratio = 1 / info.material.refractive_index;
            }
            // If ratio = 1, optimize by skipping calculations
            if (refractive_index_ratio == 1) {
                origin += direction * 0.005; // Enter or exit the object
                return 1;
            }
            double cos_a = (info.normal * -1).dot(direction);
            double cos_b = 1 - refractive_index_ratio * refractive_index_ratio * (1 - cos_a * cos_a);
            if (cos_b < 0) {
                // Total internal reflection
                direction -= info.normal * 2.0 * direction.dot(info.normal);
            } else {
                cos_b = sqrt(cos_b);
                direction *= refractive_index_ratio + info.normal * (refractive_index_ratio * cos_a - cos_b);
                origin += direction * 0.005; // Enter or exit the object
            }
        }
        radiance += throughput * info.material.emission_color * info.material.emissive_strength;
        throughput *= is_specular ? info.material.specular_color : info.material.color;

        // Sample direct lights
        for (const auto light : scene.lights) {
            light->rng = &rng;
            if (is_diffuse) {
                radiance += throughput * light->sample(info.point, info.normal, scene) * inv_pi;
            }
            if ((is_specular || is_reflective) && (has_hit_diffuse || light->is_mirror_visible())) {
                radiance += throughput * light->sample(info.point, info.normal, (info.point - origin).normalize(), scene);
            }
        }

        inv_dir = 1 / direction;

        // Check if ray fully absorbed
        if (throughput.x < 0.001 && throughput.y < 0.001 && throughput.z < 0.001) {
            return 2; // Ray absorbed
        }

        return 1; // Successful bounce
    }

    if (scene.background) {
        radiance += throughput * scene.background->sample(direction);
    }

    return 0; // No collision
}

Vector3 Ray::trace() {
    const int max_bounces = scene.max_bounces;
    for (bounces = 0; bounces < max_bounces; bounces++) {
        int result = bounce();
        if (result == 0) {
            if (bounces == 0) {
                hit_nothing = true;
            }
            break; // No collision
        } else if (result == 2) {
            break; // Ray absorbed
        }
    }
    return radiance;
}

void Ray::reset() {
    origin = original_origin;
    direction = original_direction;
    inv_dir = original_inv_dir;
    bounces = 0;
    radiance = Vector3(0);
    throughput = Vector3(1);
    has_hit_diffuse = false;
}