import { assert } from "../../../../assert.js";

/**
 * Max-heap of (uint32 id, float32 score) entries with an external id→slot index
 * for O(log n) `update` and `remove` by id (rather than O(n) linear scan).
 *
 * Used by the k-way Fiduccia–Mattheyses refinement, where each vertex move
 * requires updating the gain entries of every neighbour's boundary vertex.
 * Without the index, that's O(degree · |boundary|) per move; with it,
 * O(degree · log |boundary|).
 *
 * Behaviour notes:
 *  - This is a MAX-heap (`pop_max` returns the entry with the largest score).
 *  - Ids must be in `[0, id_capacity)`; the index table is sized at construction.
 *  - Duplicate ids are not allowed. Calling `insert` with an existing id throws.
 *
 * Storage layout mirrors meep's `Uint32Heap`: each heap slot is 8 bytes packed
 * as `(float32 score, uint32 id)`. The id→slot map is a separate
 * `Uint32Array(id_capacity)` where `SLOT_NOT_PRESENT` means "not in heap".
 *
 * We measured a JS `Map<number, number>` as the index instead — it removes the
 * eager `id_capacity` allocation but added 35-60% to FM-dominated benchmarks
 * because `Map.get/set/delete` per `__swap_slots` call (O(log n) per pop) is
 * meaningfully slower than typed-array indexing. Kept the typed-array form.
 */

const SLOT_NOT_PRESENT = 0xFFFFFFFF;
const BYTES_PER_SLOT = 8;

export class IndexedFloatMaxHeap {
    /**
     * @param {number} id_capacity ids must be in [0, id_capacity)
     * @param {number} [initial_slot_capacity] initial heap slot count (grows on demand)
     */
    constructor(
        id_capacity,
        initial_slot_capacity = Math.min(64, id_capacity)
    ) {
        assert.isNonNegativeInteger(id_capacity, 'id_capacity');
        assert.isNonNegativeInteger(initial_slot_capacity, 'initial_slot_capacity');

        this.__slot_capacity = initial_slot_capacity;
        this.__size = 0;

        this.__data_buffer = new ArrayBuffer(initial_slot_capacity * BYTES_PER_SLOT);
        this.__data_uint32 = new Uint32Array(this.__data_buffer);
        this.__data_float32 = new Float32Array(this.__data_buffer);

        /**
         * Maps id → slot index in [0, size), or SLOT_NOT_PRESENT.
         * @type {Uint32Array}
         */
        this.__slot_of_id = new Uint32Array(id_capacity);
        this.__slot_of_id.fill(SLOT_NOT_PRESENT);
    }

    get size() {
        return this.__size;
    }

    /**
     * Empty the heap. Only resets the index entries that are actually populated,
     * so this is O(size), not O(id_capacity).
     */
    clear() {
        const data_uint32 = this.__data_uint32;
        const size = this.__size;
        const slot_of_id = this.__slot_of_id;
        for (let slot = 0; slot < size; slot++) {
            slot_of_id[data_uint32[slot * 2 + 1]] = SLOT_NOT_PRESENT;
        }
        this.__size = 0;
    }

    contains(id) {
        return this.__slot_of_id[id] !== SLOT_NOT_PRESENT;
    }

    /**
     * Swap the entries at two slot indices, keeping the id→slot index in sync.
     * @private
     */
    __swap_slots(slot_a, slot_b) {
        const data_uint32 = this.__data_uint32;
        const word_a = slot_a << 1;
        const word_b = slot_b << 1;

        const score_a = data_uint32[word_a];
        const id_a    = data_uint32[word_a + 1];
        const score_b = data_uint32[word_b];
        const id_b    = data_uint32[word_b + 1];

        data_uint32[word_a]     = score_b;
        data_uint32[word_a + 1] = id_b;
        data_uint32[word_b]     = score_a;
        data_uint32[word_b + 1] = id_a;

        const slot_of_id = this.__slot_of_id;
        slot_of_id[id_a] = slot_b;
        slot_of_id[id_b] = slot_a;
    }

    /**
     * Restore heap order downward starting at `start_slot`. Used when an entry's
     * score has just decreased.
     * @private
     */
    __sift_down(start_slot) {
        const size = this.__size;
        const data_float32 = this.__data_float32;
        let slot = start_slot;

        for (;;) {
            const left_child = (slot << 1) + 1;
            const right_child = left_child + 1;
            let largest_slot = slot;

            if (left_child < size && data_float32[left_child * 2] > data_float32[largest_slot * 2]) {
                largest_slot = left_child;
            }
            if (right_child < size && data_float32[right_child * 2] > data_float32[largest_slot * 2]) {
                largest_slot = right_child;
            }
            if (largest_slot === slot) {
                break;
            }
            this.__swap_slots(slot, largest_slot);
            slot = largest_slot;
        }
    }

    /**
     * Restore heap order upward starting at `start_slot`. Used when an entry's
     * score has just increased.
     * @private
     */
    __sift_up(start_slot) {
        const data_float32 = this.__data_float32;
        let slot = start_slot;

        while (slot > 0) {
            const parent_slot = (slot - 1) >> 1;
            if (data_float32[parent_slot * 2] >= data_float32[slot * 2]) {
                break;
            }
            this.__swap_slots(parent_slot, slot);
            slot = parent_slot;
        }
    }

    __grow_capacity() {
        const old_capacity = this.__slot_capacity;
        const new_capacity = Math.max(old_capacity + 16, Math.ceil(old_capacity * 1.5));
        const new_buffer = new ArrayBuffer(new_capacity * BYTES_PER_SLOT);
        const new_uint32 = new Uint32Array(new_buffer);
        new_uint32.set(this.__data_uint32);
        this.__data_buffer = new_buffer;
        this.__data_uint32 = new_uint32;
        this.__data_float32 = new Float32Array(new_buffer);
        this.__slot_capacity = new_capacity;
    }

    /**
     * @param {number} id
     * @param {number} score
     */
    insert(id, score) {
        if (this.__slot_of_id[id] !== SLOT_NOT_PRESENT) {
            throw new Error(`id ${id} already in heap`);
        }
        if (this.__size >= this.__slot_capacity) {
            this.__grow_capacity();
        }
        const new_slot = this.__size;
        this.__data_float32[new_slot * 2] = score;
        this.__data_uint32[new_slot * 2 + 1] = id;
        this.__slot_of_id[id] = new_slot;
        this.__size = new_slot + 1;
        this.__sift_up(new_slot);
    }

    /**
     * @returns {number} id of the entry with the largest score, or -1 if empty
     */
    pop_max() {
        if (this.__size === 0) {
            return -1;
        }
        const data_uint32 = this.__data_uint32;
        const top_id = data_uint32[1];
        const new_size = this.__size - 1;
        this.__size = new_size;
        this.__slot_of_id[top_id] = SLOT_NOT_PRESENT;

        if (new_size > 0) {
            const last_slot_word = new_size << 1;
            const moved_score = this.__data_float32[last_slot_word];
            const moved_id = data_uint32[last_slot_word + 1];
            this.__data_float32[0] = moved_score;
            data_uint32[1] = moved_id;
            this.__slot_of_id[moved_id] = 0;
            this.__sift_down(0);
        }

        return top_id;
    }

    /**
     * @param {number} id
     * @param {number} new_score
     */
    update(id, new_score) {
        const slot = this.__slot_of_id[id];
        if (slot === SLOT_NOT_PRESENT) {
            throw new Error(`id ${id} not in heap`);
        }
        const slot_score_address = slot * 2;
        const old_score = this.__data_float32[slot_score_address];
        this.__data_float32[slot_score_address] = new_score;
        if (new_score > old_score) {
            this.__sift_up(slot);
        } else if (new_score < old_score) {
            this.__sift_down(slot);
        }
    }

    /**
     * @param {number} id
     * @returns {boolean} true if the id was present and removed
     */
    remove(id) {
        const slot = this.__slot_of_id[id];
        if (slot === SLOT_NOT_PRESENT) {
            return false;
        }
        this.__slot_of_id[id] = SLOT_NOT_PRESENT;
        const new_size = this.__size - 1;
        this.__size = new_size;

        if (slot === new_size) {
            return true;
        }

        // Move the last entry into the vacated slot, then restore heap order.
        const last_slot_word = new_size << 1;
        const moved_score = this.__data_float32[last_slot_word];
        const moved_id = this.__data_uint32[last_slot_word + 1];
        const slot_word = slot * 2;
        const old_slot_score = this.__data_float32[slot_word];
        this.__data_float32[slot_word] = moved_score;
        this.__data_uint32[slot_word + 1] = moved_id;
        this.__slot_of_id[moved_id] = slot;

        if (moved_score > old_slot_score) {
            this.__sift_up(slot);
        } else {
            this.__sift_down(slot);
        }
        return true;
    }
}
