"""Merge the four bucket packages into one multifunction package with shared weights.

Run convert_bucket.py for each length first, e.g.

    cd scripts/build
    for L in 64 128 256 512; do
        uv run python convert_bucket.py --length $L --max-heads 4 --max-options 32 --output-dir build
    done
    uv run python build_multifunction.py --build-dir build

Each function (L64, L128, L256, L512) keeps its own fixed-shape graph; identical weight
tensors are stored once (~0.9 GB instead of ~3.5 GB for four separate packages).
"""
import argparse
from pathlib import Path

import coremltools as ct

from runtime import package_name

BUCKETS = (64, 128, 256, 512)
ROOT = Path(__file__).resolve().parents[2]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--build-dir", default="build")
    parser.add_argument("--output", default=str(ROOT / "models" / "GLiNER2.5-Decide-MultiFn-fp16.mlpackage"))
    args = parser.parse_args()
    desc = ct.utils.MultiFunctionDescriptor()
    for length in BUCKETS:
        source = Path(args.build_dir) / package_name("fp16", length, 4, 32)
        desc.add_function(str(source), src_function_name="main", target_function_name=f"L{length}")
    desc.default_function_name = "L128"
    ct.utils.save_multifunction(desc, args.output)
    print("saved", args.output)


if __name__ == "__main__":
    main()