import { assert } from "../../../assert.js";
import { ceilPowerOfTwo } from "../../../binary/operations/ceilPowerOfTwo.js";
import { split_by_2 } from "../../../binary/split_by_2.js";
import { array_copy } from "../../../collection/array/array_copy.js";
import { array_swap } from "../../../collection/array/array_swap.js";
import { SCRATCH_UINT32_TRAVERSAL_STACK } from "../../../collection/SCRATCH_UINT32_TRAVERSAL_STACK.js";
import { aabb2_array_combine } from "../aabb/aabb2_array_combine.js";

const BINARY_NODE_SIZE = 4;
const LEAF_NODE_SIZE = 5;
const BOX_BYTE_SIZE = 4;

/**
 * Assumes data will be normalized to 0...1 value range
 * @param {Float32Array} data
 * @param {number} address
 * @param {number} bounds_x0
 * @param {number} bounds_y0
 * @param {number} multiplier_x
 * @param {number} multiplier_y
 * @returns {number}
 */
function build_normalized_morton_2d(
    data, address,
    bounds_x0, bounds_y0,
    multiplier_x, multiplier_y
) {

    const x0 = data[address];
    const y0 = data[address + 1];
    const x1 = data[address + 3];
    const y1 = data[address + 4];

    const cx = (x0 + x1) * 0.5;
    const cy = (y0 + y1) * 0.5;

    const normalized_x = (cx - bounds_x0) * multiplier_x
    const normalized_y = (cy - bounds_y0) * multiplier_y

    return (split_by_2(normalized_y) << 1)
        | split_by_2(normalized_x)
        ;
}

/**
 *
 * @param {Float32Array} data
 * @param {number} destination
 * @param {number} source
 */
function copy_box_zero_size(data, destination, source) {
    assert.isNonNegativeInteger(destination, 'destination');
    assert.isNonNegativeInteger(source, 'source');

    const x = data[source];
    const y = data[source + 1];

    assert.notNaN(x, 'x');
    assert.notNaN(y, 'y');

    data[destination] = x;
    data[destination + 1] = y;

    data[destination + 2] = x;
    data[destination + 3] = y;
}


/**
 * NOTE: code is based on BinaryUint32BVH
 * NOTE: this code is mostly untested
 */
export class StaticR2Tree {

    /**
     * @type {ArrayBuffer}
     * @private
     */
    __data_buffer;

    /**
     * @readonly
     * @private
     * @type {Uint32Array}
     */
    __data_uint32;

    /**
     *
     * @type {number}
     * @private
     */
    __node_count_binary = 0;

    /**
     *
     * @type {number}
     * @private
     */
    __node_count_leaf = 0;

    constructor() {
        this.setBuffer(new ArrayBuffer(320));
    }


    /**
     *
     * @param {ArrayBuffer} buffer
     */
    setBuffer(buffer) {
        assert.defined(buffer, 'buffer');

        this.__data_buffer = buffer;

        this.__data_float32 = new Float32Array(this.__data_buffer);
        this.__data_uint32 = new Uint32Array(this.__data_buffer);
    }


    initialize_structure() {
        // compute memory requirements
        const word_count = this.__node_count_binary * BINARY_NODE_SIZE + this.__node_count_leaf * LEAF_NODE_SIZE;
        const storage_size = word_count * 4;

        // possibly resize the storage
        if (this.__data_buffer.byteLength < storage_size) {
            this.setBuffer(new ArrayBuffer(storage_size));
        }
    }

    /**
     *
     * @param {number} count
     */
    setLeafCount(count) {
        this.__node_count_leaf = count;

        const twoLeafLimit = ceilPowerOfTwo(count);

        if (count <= 1) {
            // special case
            this.__node_count_binary = twoLeafLimit;
        } else {
            this.__node_count_binary = twoLeafLimit - 1;
        }

    }


    /**
     *
     * @param {number} index
     * @param {number} payload
     * @param {number} x0
     * @param {number} y0
     * @param {number} x1
     * @param {number} y1
     */
    setLeafData(
        index, payload,
        x0, y0,
        x1, y1
    ) {
        assert.notNaN(x0, 'x0');
        assert.notNaN(y0, 'y0');
        assert.notNaN(x1, 'x1');
        assert.notNaN(y1, 'y1');

        assert.isNonNegativeInteger(payload, 'payload');

        const address = index * LEAF_NODE_SIZE + this.__node_count_binary * BINARY_NODE_SIZE;

        const float32 = this.__data_float32;

        float32[address] = x0;
        float32[address + 1] = y0;
        float32[address + 2] = x1;
        float32[address + 3] = y1;

        this.__data_uint32[address + 4] = payload;
    }

    /**
     *
     * @returns {number}
     */
    getLeafBlockAddress() {
        return this.__node_count_binary * BINARY_NODE_SIZE;
    }

    /**
     *
     * @param {number} leaf_index
     * @returns {number}
     */
    readLeafPayload(leaf_index) {
        const block_address = this.getLeafBlockAddress();

        const address = block_address + leaf_index * LEAF_NODE_SIZE + BOX_BYTE_SIZE;

        return this.__data_uint32[address];
    }


    /**
     * Sort leaf nodes according to their morton codes
     * @param {number[]} bounds
     */
    sort_morton(bounds) {

        const leaf_block_address = this.__node_count_binary * BINARY_NODE_SIZE;

        if (this.__node_count_leaf < 2) {
            // no swaps available
            return;
        }

        const stack = SCRATCH_UINT32_TRAVERSAL_STACK;
        const stack_top = stack.pointer;

        let stackPointer = stack_top + 2;
        let iL, iR;


        stack[0] = 0; // first node
        stack[1] = this.__node_count_leaf - 1; // last node

        const data = this.__data_float32;


        const bounds_x0 = bounds[0];
        const bounds_y0 = bounds[1];

        const bounds_span_x = bounds[2] - bounds_x0;
        const bounds_span_y = bounds[3] - bounds_y0;

        const bounds_multiplier_x = 4095 / bounds_span_x;
        const bounds_multiplier_y = 4095 / bounds_span_y;

        while (stackPointer > stack_top) {
            stackPointer -= 2;

            const right = stack[stackPointer + 1];
            const left = stack[stackPointer];

            iL = left;
            iR = right;

            const pivot_index = (left + right) >>> 1;

            const pivot_address = pivot_index * LEAF_NODE_SIZE + leaf_block_address;

            const pivot = build_normalized_morton_2d(data, pivot_address, bounds_x0, bounds_y0, bounds_multiplier_x, bounds_multiplier_y);

            /* partition */
            while (iL <= iR) {
                while (build_normalized_morton_2d(data, iL * LEAF_NODE_SIZE + leaf_block_address, bounds_x0, bounds_y0, bounds_multiplier_x, bounds_multiplier_y) < pivot) {
                    iL++;
                }

                while (build_normalized_morton_2d(data, iR * LEAF_NODE_SIZE + leaf_block_address, bounds_x0, bounds_y0, bounds_multiplier_x, bounds_multiplier_y) > pivot) {
                    iR--;
                }

                if (iL <= iR) {

                    if (iL !== iR) {
                        //do swap
                        this.__swap_leaves(iL, iR);
                    }

                    iL++;
                    iR--;
                }
            }

            /* recursion */
            if (left < iR) {
                stack[stackPointer++] = left;
                stack[stackPointer++] = iR;
            }
            if (iL < right) {
                stack[stackPointer++] = iL;
                stack[stackPointer++] = right;
            }
        }
    }

    /**
     *
     * @param {number} i
     * @param {number} j
     * @private
     */
    __swap_leaves(i, j) {
        const leaf_block_address = this.getLeafBlockAddress();
        const a = i * LEAF_NODE_SIZE + leaf_block_address;
        const b = j * LEAF_NODE_SIZE + leaf_block_address;

        const float32 = this.__data_float32;

        array_swap(
            float32, a,
            float32, b,
            LEAF_NODE_SIZE
        );
    }

    /**
     * Assemble leaf nodes into hierarchy, set binary node bounds iteratively bottom up
     */
    build() {
        const binary_node_count = this.__node_count_binary;

        const leaf_node_block_address = binary_node_count * BINARY_NODE_SIZE;

        let level = Math.floor(Math.log(binary_node_count) / Math.log(2));

        let i, offset, level_node_count;
        //NOTE: building first level separately allows to avoid some switching logic needed to determine what is the type of lower level node
        //build one level above leaf nodes
        level_node_count = Math.pow(2, level);
        offset = (level_node_count - 1) * BINARY_NODE_SIZE;

        let parentIndex;

        const node_count_leaf = this.__node_count_leaf;

        const float32 = this.__data_float32;

        // build bottom-most level, just above the leaves
        for (i = 0; i < level_node_count; i++) {
            const leaf_index_0 = i << 1;
            const leaf_index_1 = leaf_index_0 + 1;

            const leaf_offset_0 = leaf_node_block_address + leaf_index_0 * LEAF_NODE_SIZE;
            const leaf_offset_1 = leaf_node_block_address + leaf_index_1 * LEAF_NODE_SIZE;

            if (leaf_index_1 < node_count_leaf) {
                // both children nodes are valid
                aabb2_array_combine(
                    float32, offset,
                    float32, leaf_offset_0,
                    float32, leaf_offset_1
                );
            } else if (leaf_index_0 < node_count_leaf) {
                // only left child node is valid
                array_copy(float32, leaf_offset_0, float32, offset, 6);
            } else {
                //initialize to 0-size box same position as previous node
                copy_box_zero_size(this.__data_float32, offset, (offset - BINARY_NODE_SIZE));
            }

            offset += BINARY_NODE_SIZE;
        }

        level--;

        //build intermediate nodes
        for (; level >= 0; level--) {
            level_node_count = 1 << level;
            parentIndex = (level_node_count - 1);

            for (i = 0; i < level_node_count; i++) {

                const childIndex0 = (parentIndex << 1) + 1;

                const address_parent = parentIndex * BINARY_NODE_SIZE;

                const address_child_0 = childIndex0 * BINARY_NODE_SIZE;
                const address_child_1 = address_child_0 + BINARY_NODE_SIZE;

                aabb2_array_combine(
                    float32, address_parent,
                    float32, address_child_0,
                    float32, address_child_1
                );

                parentIndex++;
            }
        }
    }
}