import { Metis } from "./metis.js";
import { metis_options } from "./metis_options.js";

/**
 * How many times to retry partitioning before giving up
 * @type {number}
 */
const MAX_STEPS = 50;

/**
 * @deprecated Calls the WASM-backed `Metis.INSTANCE`, which is deprecated. Build
 * the equivalent on top of `metis_partition_kway` from
 * `./native/metis_partition_kway.js` — the binary-search-on-k pattern here is
 * independent of the partitioner, so the same loop body works against either
 * backend.
 *
 * The native partitioner does roughly the same balance enforcement internally
 * via FM with `ub_factor = 1.03`, so the binary search typically converges in
 * 1-2 passes; for many use cases a single `metis_partition_kway` call with
 * `k = Math.ceil(node_count / patch_size)` is sufficient and the BS wrapper is
 * unnecessary.
 *
 * @param {Uint32Array} result
 * @param {number} node_count
 * @param {Uint32Array} edge_addresses
 * @param {Uint32Array} adjacency_uint32
 * @param {Uint32Array} edge_weights_uint32
 * @param {number} patch_size
 * @returns {Promise<number>}
 */
export async function metis_cluster_bs(
    result,
    node_count,
    edge_addresses,
    adjacency_uint32,
    edge_weights_uint32,
    patch_size
) {

    const metis = Metis.INSTANCE;

    let partition_count_bound_lower = Math.ceil(node_count / patch_size);
    let partition_count_bound_upper = partition_count_bound_lower;

    let partition_count = partition_count_bound_lower;

    let counts = new Uint32Array(partition_count);

    let passes = 0;
    let p;
    let max_count;

    const metisOptions = new metis_options();
    metisOptions.seed = 0xCAFE;

    for (; ;) {
        passes++;

        if (passes > MAX_STEPS) {
            throw new Error(`Maximum number of passes exceeded (${MAX_STEPS}). Current cluster count= ${partition_count}, initial=${partition_count_bound_lower}. Max elements per cluster = ${max_count}`);
        }

        p = await metis.partition(
            node_count,
            partition_count,
            edge_addresses,
            adjacency_uint32,
            edge_weights_uint32,
            metisOptions
        );

        max_count = 0;
        for (let i = 0; i < node_count; i++) {
            const part = p[i];

            counts[part]++;

            if (counts[part] > max_count) {
                max_count = counts[part];
            }
        }

        const overflow = max_count - patch_size;

        // perform binary search for the optimal partition count
        if (overflow > 0) {
            // update lower bound
            partition_count_bound_lower = partition_count + 1;

            // part overflow
            if (partition_count === partition_count_bound_upper) {
                partition_count_bound_upper = Math.ceil(partition_count * 1.04);
                partition_count = partition_count_bound_upper;
            } else {
                // try half-way mark
                partition_count = Math.ceil((partition_count_bound_lower + partition_count_bound_upper) / 2);
            }

        } else if (partition_count > partition_count_bound_lower && (partition_count_bound_upper - partition_count_bound_lower) > 1) {
            // no overflow, but we're probably overshooting
            partition_count_bound_upper = partition_count;
            partition_count = Math.ceil((partition_count_bound_lower + partition_count_bound_upper) / 2);

        } else {

            // done
            break;
        }


        if (partition_count > counts.length) {
            counts = new Uint32Array(partition_count)
        } else {
            counts.fill(0);
        }

        // console.log(`${partition_count}`);
    }

    // remap P values. Metis assigns unique unsigned integer values to each cluster, but they don't necessarily start at 0, to fix this we remap values
    const mapping = [];
    let index_counter = 0;

    for (let i = 0; i < p.length; i++) {
        const source_index = p[i];

        let target_index = mapping[source_index];
        if (target_index === undefined) {
            target_index = index_counter++;
            mapping[source_index] = target_index;
        }

        p[i] = target_index;
    }

    result.set(p);
    return partition_count;
}