/*

 * Copyright (c) 2024, the SerenityOS developers.

 * Copyright (c) 2021, Idan Horowitz <idan.horowitz@serenityos.org>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#pragma once

#include <AK/Array.h>

#include <AK/BinaryHeap.h>

namespace Compress {

inline void generate_huffman_lengths(Bytes lengths, ReadonlySpan<u16> frequencies, size_t max_bit_length, u16 shift = 0)

{

    VERIFY(lengths.size() == frequencies.size());

    auto const size = lengths.size();

    VERIFY((1u << max_bit_length) >= size);

    u16 heap_keys[size]; // Used for O(n) heap construction

    u16 heap_values[size];

    u16 huffman_links[size * 2];

    size_t non_zero_freqs = 0;

    for (auto frequency : frequencies) {

        if (frequency == 0)

            continue;

        frequency = max(1, frequency >> shift);

        heap_keys[non_zero_freqs] = frequency;               // sort symbols by frequency

        heap_values[non_zero_freqs] = size + non_zero_freqs; // huffman_links "links"

        non_zero_freqs++;

    }

    // special case for only 1 used symbol

    if (non_zero_freqs < 2) {

        for (size_t i = 0; i < size; i++)

            lengths[i] = (frequencies[i] == 0) ? 0 : 1;

        return;

    }

    BinaryHeap<u16, u16, 288> heap { heap_keys, heap_values, non_zero_freqs };

    // build the huffman tree - binary heap is used for efficient frequency comparisons

    while (heap.size() > 1) {

        u16 lowest_frequency = heap.peek_min_key();

        u16 lowest_link = heap.pop_min();

        u16 second_lowest_frequency = heap.peek_min_key();

        u16 second_lowest_link = heap.pop_min();

        u16 new_link = heap.size() + 1;

        u32 sum = lowest_frequency + second_lowest_frequency;

        sum = min(sum, UINT16_MAX);

        heap.insert(sum, new_link);

        huffman_links[lowest_link] = new_link;

        huffman_links[second_lowest_link] = new_link;

    }

    non_zero_freqs = 0;

    for (size_t i = 0; i < size; i++) {

        if (frequencies[i] == 0) {

            lengths[i] = 0;

            continue;

        }

        u16 link = huffman_links[size + non_zero_freqs];

        non_zero_freqs++;

        size_t bit_length = 1;

        while (link != 1) {

            bit_length++;

            link = huffman_links[link];

        }

        if (bit_length > max_bit_length) {

            VERIFY(shift < 15);

            return generate_huffman_lengths(lengths, frequencies, max_bit_length, shift + 1);

        }

        lengths[i] = bit_length;

    }

}

}