#include "psdr/nrf.h"
#include "enoki/dynamic.h"
#include "psdr/scene/scene.h"
#include "psdr/core/intersection.h"
#include "psdr/shape/mesh.h"

namespace psdr
{
using namespace enoki;

void NRF::get_textures(Scene &scene, size_t resolution, size_t multisample_factor, Vector3fD& diffuse, Vector3fD& specular, FloatD& roughness) {
    using FloatX = DynamicArray<Array<float>>;
    const size_t num_samples = multisample_factor * multisample_factor;
    float spacing = 1.0 / resolution;
    float subspacing = spacing / multisample_factor / 2.f;
    FloatX x_ = linspace<FloatX>(0.f, 1.f - spacing, resolution);
    FloatX y_ = linspace<FloatX>(1.f - spacing, 0.f, resolution);
    for (size_t i = 0; i < multisample_factor; i++) {
        float x_offset = (i * 2 + 1) * subspacing;
        for (size_t j = 0; j < multisample_factor; j++) { 
            float y_offset = (j * 2 + 1) * subspacing;
            Array<FloatX, 2> grid_ = meshgrid(x_ +  x_offset, y_ + y_offset);

            size_t size = sqr(resolution);
            FloatC x = FloatC::copy(grid_[0].data(), size);
            FloatC y = FloatC::copy(grid_[1].data(), size);
            FloatC z = zero<FloatC>(size) + 5.f;
            Vector3fC ray_origin(x, y, z);
            Vector3fC ray_dir = zero<Vector3fC>(resolution * resolution);
            ray_dir[2] -= 1.f;
            RayD ray(ray_origin, ray_dir);

            MaskD active = true;
            IntersectionD its = scene.ray_intersect(ray, active);
            active &= its.is_valid();
            diffuse[active] += detach(its.diffuse);
            specular[active] += detach(its.specular);
            roughness[active] += detach(its.roughness);
        
        }
    }
    
    diffuse /= num_samples;
    specular /= num_samples;
    roughness /= num_samples;
}

};