import { assert } from "../../../../../../core/assert.js";
import { UINT32_MAX } from "../../../../../../core/binary/UINT32_MAX.js";
import { array_copy } from "../../../../../../core/collection/array/array_copy.js";
import {
    decode_octahedron_to_unit
} from "../../../../../../core/geom/3d/normal/octahedron/decode_octahedron_to_unit.js";
import { Ray3 } from "../../../../../../core/geom/3d/ray/Ray3.js";
import { v3_distance } from "../../../../../../core/geom/vec3/v3_distance.js";
import { v3_dot_array_array } from "../../../../../../core/geom/vec3/v3_dot_array_array.js";
import { max2 } from "../../../../../../core/math/max2.js";
import { generate_hammersley_jitter } from "../../../../../../core/math/statistics/generate_hammersley_jitter.js";

/**
 * Weights of rays that are aligned with the probe direction get exponentially weighted by this exponent
 * @type {number}
 */
const DEPTH_SHARPNESS = 3;

const scratch_ray = new Ray3();

/**
 *
 * @type {number[]}
 */
const ray_hit = [];


/**
 *
 * @param {number[]} result
 * @param {number} result_offset
 * @param {PathTracedScene} scene
 * @param {number[]|Vector3} position
 * @param {number} position_offset
 * @param {number} resolution
 * @param {number} max_depth
 * @param {number} [sub_sample_count]
 */
export function bake_octahedral_depth_map(
    result, result_offset,
    scene,
    position, position_offset,
    resolution,
    max_depth, sub_sample_count=4
) {

    assert.lessThan(result_offset + resolution * resolution, UINT32_MAX, 'Result overflow');

    scratch_ray.tMax = max_depth;
    array_copy(position, position_offset, scratch_ray, 0, 3);

    const NORMAL_BIAS = 1e-6;

    const ray_direction = scratch_ray.direction;

    const probe_direction = new Float32Array(3);

    const texel_uv_scale = 0.5 / (resolution);

    const jittered_sub_samples = generate_hammersley_jitter(sub_sample_count);


    for (let oct_x = 0; oct_x < resolution; oct_x++) {
        for (let oct_y = 0; oct_y < resolution; oct_y++) {

            let distance_sum = 0;
            let distance2_sum = 0;
            let weight_sum = 0;

            // offset position by half a pixel, as we are storing values for pixel centers
            const u = (oct_x + 0.5) / (resolution);
            const v = (oct_y + 0.5) / (resolution);

            decode_octahedron_to_unit(
                probe_direction, 0,
                u * 2 - 1, v * 2 - 1
            );

            for (let sub_sample_index = 0; sub_sample_index < sub_sample_count; sub_sample_index++) {

                const sample_index2 = sub_sample_index * 2;

                const sample_u = jittered_sub_samples[sample_index2];
                const sample_v = jittered_sub_samples[sample_index2 + 1];

                const ray_u = u + (sample_u * 2 - 1) * texel_uv_scale;
                const ray_v = v + (sample_v * 2 - 1) * texel_uv_scale;

                decode_octahedron_to_unit(
                    ray_direction, 0,
                    ray_u * 2 - 1, ray_v * 2 - 1
                );

                let distance = scene.trace(ray_hit, scratch_ray);

                if (distance < 0) {
                    // no hit
                    distance = max_depth;
                } else {
                    const surface_normal_x = ray_hit[3];
                    const surface_normal_y = ray_hit[4];
                    const surface_normal_z = ray_hit[5];

                    const surface_position_x = ray_hit[0];
                    const surface_position_y = ray_hit[1];
                    const surface_position_z = ray_hit[2];

                    const hit_normal_bias = NORMAL_BIAS;

                    // sink the contact into the surface along the hit normal
                    const biased_hit_x = surface_position_x + surface_normal_x * hit_normal_bias;
                    const biased_hit_y = surface_position_y + surface_normal_y * hit_normal_bias;
                    const biased_hit_z = surface_position_z + surface_normal_z * hit_normal_bias;

                    distance = v3_distance(
                        scratch_ray[0], scratch_ray[1], scratch_ray[2],
                        biased_hit_x, biased_hit_y, biased_hit_z
                    );

                }

                const direction_dot_product = v3_dot_array_array(ray_direction, 0, probe_direction, 0);

                const weight = Math.pow(
                    max2(0,
                        direction_dot_product
                    ),
                    DEPTH_SHARPNESS
                );

                distance_sum += distance * weight;
                distance2_sum += distance * distance * weight;
                weight_sum += weight;
            }

            const pixel_index = oct_y * resolution + oct_x;

            const address = result_offset + pixel_index * 2;

            const mean = distance_sum / weight_sum;
            const mean2 = distance2_sum / weight_sum;

            result[address] = mean;
            result[address + 1] = mean2;
        }
    }

}

