// Copyright 2023 Intel Corporation
// SPDX-License-Identifier: BSD-3-Clause

// This example is based on Syoyo Fujita's aobench:
// https://github.com/syoyo/aobench also see the Google code page:
// https://code.google.com/archive/p/aobench/

#include "aobench_stub.ih"

#define NUM_AO_SAMPLES 2

#ifndef AOBENCH_TARGET_SYCL
struct Parameters {
    int width;
    int height;
    int y_offset;
    int n_samples;
    Scene *uniform scene;
    int *uniform rng_seeds;
    float *uniform image;
};
#else
struct Parameters {
    int width = 256;
    int height = 256;
    int y_offset = 0;
    int n_samples = 4;
    Scene *scene = nullptr;
    int *rng_seeds = nullptr;
    float *image = nullptr;
};

// ISPC standard library's random number generator implementation
// ported to SYCL to match the ISPC version
struct RNGState {
    unsigned int z1, z2, z3, z4;
};

inline unsigned int random(RNGState *state) {
    unsigned int b;

    b = ((state->z1 << 6) ^ state->z1) >> 13;
    state->z1 = ((state->z1 & 4294967294U) << 18) ^ b;
    b = ((state->z2 << 2) ^ state->z2) >> 27;
    state->z2 = ((state->z2 & 4294967288U) << 2) ^ b;
    b = ((state->z3 << 13) ^ state->z3) >> 21;
    state->z3 = ((state->z3 & 4294967280U) << 7) ^ b;
    b = ((state->z4 << 3) ^ state->z4) >> 12;
    state->z4 = ((state->z4 & 4294967168U) << 13) ^ b;
    return (state->z1 ^ state->z2 ^ state->z3 ^ state->z4);
}

inline float frandom(RNGState *state) {
    unsigned int irand = random(state);
    irand &= (1ul << 23) - 1;
    return floatbits(0x3F800000 | irand) - 1.0f;
}

inline void seed_rng(RNGState *state, unsigned int seed) {
    state->z1 = seed;
    state->z2 = seed ^ 0xbeeff00d;
    state->z3 = ((seed & 0xfffful) << 16) | (seed >> 16);
    state->z4 =
        (((seed & 0xfful) << 24) | ((seed & 0xff00ul) << 8) | ((seed & 0xff0000ul) >> 8) | (seed & 0xff000000ul) >> 24);
}

#endif

inline void ortho_basis(vec3f basis[3], const vec3f n) {
    basis[2] = n;
    basis[1] = 0.f;

    if (n.x < 0.6f && n.x > -0.6f) {
        basis[1].x = 1.f;
    } else if (n.y < 0.6f && n.y > -0.6f) {
        basis[1].y = 1.f;
    } else if (n.z < 0.6f && n.z > -0.6f) {
        basis[1].z = 1.f;
    } else {
        basis[1].x = 1.f;
    }
    basis[0] = normalized(cross(basis[1], basis[2]));
    basis[1] = normalized(cross(basis[2], basis[0]));
}

float ambient_occlusion(const RayHit &rayhit, Scene *uniform scene, RNGState &rng_state) {
    const vec3f ray_org = {rayhit.ray.org_x, rayhit.ray.org_y, rayhit.ray.org_z};
    const vec3f ray_dir = {rayhit.ray.dir_x, rayhit.ray.dir_y, rayhit.ray.dir_z};
    const vec3f hitp = ray_org + rayhit.ray.tfar * ray_dir;
    const vec3f n = {rayhit.hit.Ng_x, rayhit.hit.Ng_y, rayhit.hit.Ng_z};
    vec3f basis[3];
    ortho_basis(basis, n);
    float occlusion = 0.f;

    for (int j = 0; j < NUM_AO_SAMPLES; ++j) {
        for (int i = 0; i < NUM_AO_SAMPLES; ++i) {
            const float theta = sqrt(frandom(&rng_state));
            const float phi = 2.0f * PI * frandom(&rng_state);

            const float x = cos(phi) * theta;
            const float y = sin(phi) * theta;
            const float z = sqrt(1.f - theta * theta);

            Ray ao_ray;
            // Transform from object space to world space
            ao_ray.org_x = hitp.x;
            ao_ray.org_y = hitp.y;
            ao_ray.org_z = hitp.z;
            ao_ray.dir_x = x * basis[0].x + y * basis[1].x + z * basis[2].x;
            ao_ray.dir_y = x * basis[0].y + y * basis[1].y + z * basis[2].y;
            ao_ray.dir_z = x * basis[0].z + y * basis[1].z + z * basis[2].z;

            ao_ray.time = 0;
            ao_ray.tnear = 0.001f;
            ao_ray.tfar = 1e20f;

            occludedV(scene, &ao_ray);
            if (ao_ray.tfar < 0.f) {
                occlusion += 1.f;
            }
        }
    }
    occlusion = (NUM_AO_SAMPLES * NUM_AO_SAMPLES - occlusion) / (NUM_AO_SAMPLES * NUM_AO_SAMPLES);
    return occlusion;
}

void compute_tile(Parameters *uniform p, const uniform int tile_x, const uniform int tile_y) {
    const uniform float inv_samples = 1.f / p->n_samples;
    // Now SIMD_WIDTH = programCount so we don't need the foreach
    const int ix = tile_x * programCount + programIndex;
    if (ix >= p->width) {
        return;
    }

    RNGState rng_state;
    seed_rng(&rng_state, p->rng_seeds[tile_y * p->width + ix]);

    p->image[tile_y * p->width + ix] = 0;

    for (uniform int v = 0; v < p->n_samples; ++v) {
        for (uniform int u = 0; u < p->n_samples; ++u) {
            const float px = (ix + (u * inv_samples) - p->width / 2.f) / (p->width / 2.f);
            const float py = -(tile_y + (v * inv_samples) - p->height / 2.f) / (p->height / 2.f);

            RayHit rayhit;
            rayhit.ray.org_x = 0.f;
            rayhit.ray.org_y = 0.f;
            rayhit.ray.org_z = 0.f;

            vec3f dir = {px, py, -1.f};
            dir = normalized(dir);
            rayhit.ray.dir_x = dir.x;
            rayhit.ray.dir_y = dir.y;
            rayhit.ray.dir_z = dir.z;

            rayhit.ray.time = 0.f;
            rayhit.ray.tnear = 0.f;
            rayhit.ray.tfar = 1e20f;

            rayhit.hit.primID = INVALID_GEOMETRY_ID;
            rayhit.hit.geomID = INVALID_GEOMETRY_ID;
            rayhit.hit.instID[0] = INVALID_GEOMETRY_ID;

            intersectV(p->scene, &rayhit);
            if (rayhit.hit.primID != INVALID_GEOMETRY_ID) {
                const float col = ambient_occlusion(rayhit, p->scene, rng_state);
                // It's just a grayscale image
                p->image[tile_y * p->width + ix] += col;
            }
        }
    }
    p->image[tile_y * p->width + ix] *= inv_samples * inv_samples;
}

#ifndef AOBENCH_TARGET_SYCL
task void compute_ao_tile(void *uniform _p) {
    Parameters *uniform p = (Parameters * uniform) _p;
    const uniform int tile_x = taskIndex0;
    const uniform int tile_y = taskIndex1 + p->y_offset;
    compute_tile(p, tile_x, tile_y);
}
#endif
