# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# flake8: noqa
# pylint: disable=C0115
# pylint: disable=C0116
# pylint: disable=C0301
"""
Estimate Lhotse dynamic-bucketing bins for SALM-style multimodal training.

This script is the speechlm2 counterpart of:
  * scripts/speech_llm/estimate_token_bins.py        (text-only, 1D/2D)
  * scripts/speech_recognition/estimate_duration_bins_2d.py  (audio, 2D + outlier filtering)

Key properties for the speechlm2 SALM recipe (use_multimodal_sampling=True):

  * Audio cuts and text examples share a single integer-token length axis,
    obtained via ``token_equivalent_duration`` (audio frames cast to tokens).
  * 1D output is a flat integer list ``[i1, ..., iB]``.
  * 2D output is a list of integer pairs ``[[itok_max, otok_max], ...]``
    (input_tokens vs output_tokens), with per-bucket Z-score outlier filtering
    on output-tokens-per-input-token (TPT) and skipped-bucket merging when the
    underlying distribution produces duplicate dim-0 bins.
"""

import argparse
import ast
import math
from functools import partial
from itertools import islice
from pathlib import Path
from typing import Callable, Iterable

import numpy as np
import pandas as pd
import yaml
from lhotse.cut import Cut
from omegaconf import OmegaConf

import nemo.collections.speechlm2.data.salm_dataset  # noqa: F401  (registers lhotse_as_conversation)
from nemo.collections.asr.data.audio_to_text_lhotse import TokenizerWrapper
from nemo.collections.common.data.lhotse.cutset import read_cutset_from_config
from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig, tokenize, tokenize_with_prompt
from nemo.collections.common.data.lhotse.sampling import (
    DurationFilter,
    MultimodalFixedBucketBatchSizeConstraint2D,
    MultimodalSamplingConstraint,
    TokenCountFilter,
    TokenPerTokenFilter,
)
from nemo.collections.common.prompts.formatter import PromptFormatter
from nemo.collections.common.tokenizers import AggregateTokenizer, AutoTokenizer, SentencePieceTokenizer


def parse_args():
    parser = argparse.ArgumentParser(
        description="Estimate Lhotse dynamic-bucketing bins for the speechlm2 SALM recipe. "
        "Supports both 1D (input-token) and 2D ((input_tokens, output_tokens)) bucketing for "
        "mixed audio + text data, using MultimodalSamplingConstraint / "
        "MultimodalFixedBucketBatchSizeConstraint2D.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "input",
        help="Path to a data input configuration YAML file with an 'input_cfg' block "
        "(same shape as data.train_ds.input_cfg in a training config).",
    )
    parser.add_argument(
        "-t",
        "--tokenizer",
        nargs="+",
        required=True,
        help="Path(s) to SPE tokenizer(s) or HuggingFace repo id. More than one path requires --langs "
        "and constructs an AggregateTokenizer.",
    )
    parser.add_argument("-a", "--langs", nargs="+", help="Language names for each AggregateTokenizer sub-tokenizer.")
    parser.add_argument(
        "-b",
        "--buckets",
        type=int,
        default=30,
        help="The desired number of buckets (dim0 => covers input sequence length / audio duration in tokens).",
    )
    parser.add_argument(
        "-s",
        "--sub-buckets",
        type=int,
        default=None,
        help="The desired number of sub-buckets (dim1 => covers output sequence length / num_tokens). "
        "If not provided, we'll only perform 1D bucketing.",
    )
    parser.add_argument(
        "-n",
        "--num_examples",
        type=int,
        default=-1,
        help="The number of examples (utterances) to estimate the bins. -1 means use all data "
        "(be careful: it could be iterated over infinitely).",
    )
    parser.add_argument(
        "-l",
        "--min_tokens",
        type=float,
        default=-float("inf"),
        help="If specified, we'll filter out examples with fewer tokens than this number.",
    )
    parser.add_argument(
        "-u",
        "--max_tokens",
        type=float,
        default=float("inf"),
        help="If specified, we'll filter out examples with more tokens than this number.",
    )
    parser.add_argument(
        "--max_tpt",
        type=float,
        default=float("inf"),
        help="If specified, we'll filter out examples with more output tokens per input token than this.",
    )
    parser.add_argument(
        "--min_duration",
        type=float,
        default=-float("inf"),
        help="If specified, we'll filter out audio cuts shorter than this many seconds.",
    )
    parser.add_argument(
        "--max_duration",
        type=float,
        default=float("inf"),
        help="If specified, we'll filter out audio cuts longer than this many seconds.",
    )
    parser.add_argument(
        "--token_equivalent_duration",
        type=float,
        default=0.08,
        help="Audio seconds equivalent to one text token; used to convert audio duration to tokens. "
        "Should match the data.train_ds.token_equivalent_duration of the recipe (default 0.08 = "
        "8x40ms encoder frame, matching nvidia/canary-1b-v2).",
    )
    parser.add_argument(
        "--token_outlier_threshold",
        type=float,
        default=6.0,
        help="(2D mode only) Z-score threshold for output-tokens-per-input-token (TPT) outliers; "
        "per top-level bucket, examples > N sigma above the mean are excluded from sub-bucket "
        "estimation. Lower values are more aggressive.",
    )
    parser.add_argument(
        "--text-field",
        type=str,
        default="text",
        help="The supervision/cut field that holds transcripts. Must match the recipe's "
        "data.train_ds.text_field (e.g. 'answer' for the SALM nemotron-nano-v3 recipe), "
        "otherwise lhotse_as_conversation builds turns whose 'message' slot is None and "
        "the prompt formatter crashes.",
    )
    parser.add_argument(
        "--lang-field",
        type=str,
        default="lang",
        help="The supervision/cut field that holds the language code. Must match the recipe's "
        "data.train_ds.lang_field (e.g. 'target_lang' for the SALM nemotron-nano-v3 recipe).",
    )
    parser.add_argument(
        "--audio-locator-tag",
        type=str,
        default=None,
        help="Audio placeholder token. Propagates to datasets in input_cfg (e.g. lhotse_as_conversation), "
        "so AudioTurns get a non-null message slot. Required for any conversation-style input that "
        "interleaves audio with text.",
    )
    parser.add_argument(
        "-q", "--quiet", type=bool, default=False, help="When specified, only print the estimated bins."
    )
    parser.add_argument(
        "-f",
        "--prompt-format",
        type=str,
        help="When specified, use a prompt formatter in addition to the tokenizer. Required for "
        "accurate measurement of decoder-style models like Nemotron Nano v3 / Canary-1B.",
    )
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        help="Prompt slots provided as a Python list of dicts (used together with --prompt-format). "
        "Example: [{'role':'user','slots':{'source_lang':'en','target_lang':'en','task':'asr','pnc':'yes'}}]",
    )
    parser.add_argument(
        "-m",
        "--measure-total-length",
        type=bool,
        default=False,
        help="If True, measure context+answer length instead of context-only. Set to True for "
        "decoder-only models, False for encoder-decoder.",
    )
    parser.add_argument(
        "--quantize-bins",
        type=str,
        choices=["none", "pow2", "pow2sum"],
        default="none",
        help="Post-quantize the estimated bin caps so they round to model-friendly sizes. "
        "Floor for any non-'none' mode is 2**5 = 32. Modes: "
        "'none' = leave raw integers; "
        "'pow2' = nearest power of 2 (32, 64, 128, 256, ...); "
        "'pow2sum' = nearest single power of 2 OR sum of two distinct powers of 2 with each "
        "exponent >= 5 (e.g. 32, 64, 96, 128, 160, 192, 256, 288, ...). Duplicates produced "
        "by quantization are collapsed.",
    )
    parser.add_argument(
        "--source-config",
        type=str,
        default=None,
        help="If provided together with --output-config, also write a patched copy of this "
        "experiment YAML with data.train_ds.{num_buckets, bucket_duration_bins} updated to "
        "the estimated values (and data.train_ds.bucket_batch_size dropped, since its length "
        "no longer matches and must be re-tuned).",
    )
    parser.add_argument(
        "--output-config",
        type=str,
        default=None,
        help="Path to write the patched copy of --source-config to. Required iff --source-config " "is set.",
    )
    return parser.parse_args()


def find_non_outliers_z_score(data, threshold=4.0):
    # Note: we don't apply abs() here because we only filter the upper end of the distribution.
    # We don't mind low ratios for bucketing purposes.
    z_scores = (data - np.mean(data)) / np.std(data)
    return np.where(z_scores <= threshold)


def estimate_token_buckets_1d(
    cuts: Iterable[Cut],
    num_buckets: int,
    token_equivalent_duration: float,
    measure_total_length: bool,
    quiet: bool,
) -> list[int]:
    """1D bucketing: equal-token-mass bins along a single input-length axis.

    Mirrors estimate_duration_buckets in lhotse but operates in token units via
    MultimodalSamplingConstraint, which converts audio cuts to tokens through
    token_equivalent_duration and (optionally) sums context+answer when
    measure_total_length=True.
    """
    assert num_buckets > 1
    constraint = MultimodalSamplingConstraint(
        token_equivalent_duration=token_equivalent_duration,
        measure_total_length=measure_total_length,
    )

    sizes = []
    for c in cuts:
        sizes.append(constraint.measure_length(c))
    sizes = np.array(sizes, dtype=np.int32)
    sizes.sort()

    size_per_bucket = sizes.sum() / num_buckets

    if not quiet:
        print("Input-token distribution:")
        print(pd.Series(sizes).describe(percentiles=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]))

    bins: list[int] = []
    tot = 0
    for size in sizes:
        if tot > size_per_bucket:
            bins.append(int(size))
            tot = 0
        tot += size
    bins.append(int(sizes[-1]))
    return bins


def estimate_token_buckets_2d(
    cuts: Iterable[Cut],
    num_buckets: int,
    num_subbuckets: int,
    token_equivalent_duration: float,
    measure_total_length: bool,
    token_outlier_threshold: float,
    quiet: bool,
) -> list[tuple[int, int]]:
    """2D bucketing on (input_tokens, output_tokens) with per-bucket TPT outlier filtering.

    Combines:
      * MultimodalFixedBucketBatchSizeConstraint2D from the speech_llm script
        (handles both audio Cuts and text Formattable examples).
      * The outlier filtering and skipped-bucket merging from
        scripts/speech_recognition/estimate_duration_bins_2d.py, adapted from
        seconds-per-token (TPS) to output-tokens-per-input-token (TPT) since
        dim0 is now in tokens, not seconds.
    """
    assert num_buckets > 1
    assert num_subbuckets is not None and num_subbuckets >= 1

    constraint = MultimodalFixedBucketBatchSizeConstraint2D(
        [(0.0, 0.0)],
        [0],
        token_equivalent_duration=token_equivalent_duration,
        measure_total_length=measure_total_length,
    )

    num_input_tokens = []
    num_output_tokens = []
    for c in cuts:
        itoks, otoks = constraint.measure_length(c)
        num_input_tokens.append(itoks)
        num_output_tokens.append(otoks)
    num_input_tokens = np.array(num_input_tokens, dtype=np.int32)
    num_output_tokens = np.array(num_output_tokens, dtype=np.int32)

    # Sort jointly by input length so we can iterate in order and slice the
    # output-token array per top-level bucket.
    joint = np.rec.fromarrays([num_input_tokens, num_output_tokens])
    joint.sort()
    num_input_tokens = joint.f0
    num_output_tokens = joint.f1

    size_per_bucket = num_input_tokens.sum() / num_buckets
    max_input_tokens = int(num_input_tokens[-1])

    if not quiet:
        print("Input-token distribution:")
        print(pd.Series(num_input_tokens).describe(percentiles=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]))
        tpt_all = num_output_tokens / np.maximum(num_input_tokens, 1)
        print("Output tokens per input token (TPT) distribution:")
        print(pd.Series(tpt_all).describe(percentiles=[0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]))

    # First pass: choose dim-0 (input-token) bin edges using equal-mass slicing.
    input_bins: list[int] = []
    bin_indexes: list[int] = [0]
    tot = 0.0
    for binidx, size in enumerate(num_input_tokens):
        if tot > size_per_bucket:
            input_bins.append(int(size))
            bin_indexes.append(binidx)
            tot = 0.0
        tot += size

    if not quiet:
        print(f"Initial input_bins={input_bins}")

    bins: list[tuple[int, int]] = []

    def _estimate_output_token_buckets(max_bucket_input, start_idx, end_idx, corr_subbuckets):
        # Slice this bucket and discard top TPT outliers (Z-score on output/input).
        itoks_bucket_all = num_input_tokens[start_idx:end_idx]
        otoks_bucket_all = num_output_tokens[start_idx:end_idx]
        if len(itoks_bucket_all) == 0:
            return
        tpt_all = otoks_bucket_all / np.maximum(itoks_bucket_all, 1)
        non_outlier_indexes = find_non_outliers_z_score(tpt_all, threshold=token_outlier_threshold)
        otoks_bucket = otoks_bucket_all[non_outlier_indexes]
        itoks_bucket = itoks_bucket_all[non_outlier_indexes]
        if len(otoks_bucket) == 0:
            # Pathological case: all examples in this bucket got flagged. Fall back to the raw slice.
            otoks_bucket = otoks_bucket_all
            itoks_bucket = itoks_bucket_all
        max_tpt_bucket = (otoks_bucket / np.maximum(itoks_bucket, 1)).max()
        # Sort within-bucket by output tokens for sub-bucketing.
        otoks_bucket_sorted = np.sort(otoks_bucket)
        if not quiet:
            outlier_tpt = np.delete(tpt_all, non_outlier_indexes)
            print(
                f"[bucket <= {max_bucket_input} tok] [otoks: {int(otoks_bucket_sorted.min())} - "
                f"{int(otoks_bucket_sorted.max())}] [approx-max-tpt: {max_tpt_bucket:.3f}] "
                f"Discarded {end_idx - start_idx - len(otoks_bucket_sorted)} outliers",
                end=" ",
            )
            if len(outlier_tpt) > 0:
                print(f"(min-outlier: {outlier_tpt.min():.3f}, max-outlier: {outlier_tpt.max():.3f}).", end="")
            print()

        tokens_per_subbucket = otoks_bucket_sorted.sum() / corr_subbuckets
        tot_toks = 0
        for num_toks in otoks_bucket_sorted:
            if tot_toks > tokens_per_subbucket:
                bins.append((max_bucket_input, int(num_toks)))
                tot_toks = 0
            tot_toks += num_toks
        bins.append((max_bucket_input, int(otoks_bucket_sorted[-1])))

    # Second pass: walk the dim-0 bins, merging consecutive bins with identical
    # input-token caps (skewed distributions can produce duplicates) and
    # multiplying their sub-bucket allowance accordingly.
    skipped_buckets = 1
    start_idx = 0
    for i, (input_bin, binidx) in enumerate(zip(input_bins, bin_indexes[1:])):
        is_last = i == len(input_bins) - 1
        if (not is_last and input_bins[i + 1] == input_bin) or (is_last and max_input_tokens == input_bin):
            skipped_buckets += 1
            continue
        _estimate_output_token_buckets(
            max_bucket_input=input_bin,
            start_idx=start_idx,
            end_idx=binidx,
            corr_subbuckets=num_subbuckets * skipped_buckets,
        )
        start_idx = binidx
        skipped_buckets = 1

    # Final bucket carries any remaining skipped sub-buckets up to the global max.
    _estimate_output_token_buckets(
        max_bucket_input=max_input_tokens,
        start_idx=start_idx,
        end_idx=len(num_input_tokens),
        corr_subbuckets=num_subbuckets * skipped_buckets,
    )
    return bins


def load_tokenizer(paths: list[str], langs: list[str] = None) -> TokenizerWrapper:
    if len(paths) == 1:
        (p,) = paths
        if Path(p).exists():
            tok = SentencePieceTokenizer(p)
        else:
            # Assume HuggingFace repo id; trust_remote_code is required for
            # custom tokenizers (e.g. nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16).
            tok = AutoTokenizer(p, use_fast=True, trust_remote_code=True)
    else:
        assert langs is not None and len(paths) == len(langs), (
            f"Cannot create AggregateTokenizer; each tokenizer must have a language assigned via "
            f"--langs (got --tokenizer={paths} and --langs={langs})"
        )
        tok = AggregateTokenizer({lang: SentencePieceTokenizer(p) for lang, p in zip(langs, paths)})
    return TokenizerWrapper(tok)


def apply_tokenizer(cut, tokenizer=None, prompt: PromptFormatter = None):
    if prompt is not None:
        cut = tokenize_with_prompt(cut, tokenizer, prompt)
    elif tokenizer is not None:
        cut = tokenize(cut, tokenizer)
    return cut


_POW2_FLOOR_EXP = 5  # 2**5 = 32


def _quantize_value(v: int, mode: str) -> int:
    """Round a single integer bin cap to the nearest mode-allowed value (>= 32)."""
    if mode == "none":
        return int(v)
    floor = 1 << _POW2_FLOOR_EXP
    v = max(int(v), floor)
    if mode == "pow2":
        log2v = math.log2(v)
        lo = 1 << int(math.floor(log2v))
        hi = 1 << int(math.ceil(log2v))
        return lo if (v - lo) <= (hi - v) else hi
    if mode == "pow2sum":
        # Allowed values: single powers 2^a (a >= 5) and sums 2^a + 2^b (a > b >= 5).
        # Generate candidates densely enough to bracket v.
        max_exp = int(math.ceil(math.log2(2 * v))) + 1
        candidates = set()
        for a in range(_POW2_FLOOR_EXP, max_exp + 1):
            candidates.add(1 << a)
            for b in range(_POW2_FLOOR_EXP, a):
                candidates.add((1 << a) + (1 << b))
        return min(candidates, key=lambda c: (abs(c - v), c))
    raise ValueError(f"Unknown --quantize-bins mode: {mode!r}")


def quantize_bins(bins, mode: str):
    """Quantize each bin cap (or pair) and collapse duplicates while preserving order."""
    if mode == "none":
        return bins
    out = []
    seen = set()
    for b in bins:
        if isinstance(b, (tuple, list)):
            qb = tuple(_quantize_value(x, mode) for x in b)
        else:
            qb = _quantize_value(b, mode)
        if qb not in seen:
            out.append(qb)
            seen.add(qb)
    return out


def maybe_patch_config(source_config, output_config, bins, num_buckets_total):
    """Write a copy of ``source_config`` with the bucket fields updated.

    Drops ``data.train_ds.bucket_batch_size`` since its length must match
    ``num_buckets`` and the new bin layout almost certainly invalidates the old
    batch-size schedule -- the user is expected to re-tune it (e.g. via
    ``scripts/speech_recognition/oomptimizer.py``).
    """
    if source_config is None and output_config is None:
        return
    if source_config is None or output_config is None:
        raise SystemExit("--source-config and --output-config must be provided together")
    with open(source_config) as f:
        cfg = yaml.safe_load(f)
    train_ds = cfg.get("data", {}).get("train_ds") if isinstance(cfg, dict) else None
    if train_ds is None:
        print(f"WARNING: {source_config} has no data.train_ds; skipping patch.")
        return
    # Force list-of-lists for 2D bins so YAML doesn't render Python tuples.
    if bins and isinstance(bins[0], (tuple, list)):
        bins_yaml = [list(b) for b in bins]
    else:
        bins_yaml = list(bins)
    train_ds["num_buckets"] = int(num_buckets_total)
    train_ds["bucket_duration_bins"] = bins_yaml
    train_ds.pop("bucket_batch_size", None)
    with open(output_config, "w") as f:
        yaml.safe_dump(cfg, f, sort_keys=False)
    print(f"Wrote patched config to {output_config}")
    print("Note: bucket_batch_size was dropped -- re-tune it (e.g. via oomptimizer.py) before training.")


class RejectionsCounter:
    def __init__(self, predicate: Callable, message: str):
        self.predicate = predicate
        self.message = message
        self.total = 0
        self.rejected = 0

    def __call__(self, example) -> bool:
        ans = self.predicate(example)
        self.total += 1
        if not ans:
            self.rejected += 1
        return ans

    def print_report(self) -> None:
        if self.rejected:
            print(f"{self.message} | Rejected {self.rejected}/{self.total} examples.")


def main():
    args = parse_args()

    if not args.quiet:
        pd.set_option("display.float_format", lambda x: "%.2f" % x)

    tokenizer = None
    prompt = None
    if args.tokenizer is not None:
        tokenizer = load_tokenizer(args.tokenizer, args.langs)
        if args.prompt_format is not None:
            prompt_defaults = ast.literal_eval(args.prompt) if args.prompt is not None else None
            prompt = PromptFormatter.resolve(args.prompt_format)(tokenizer._tokenizer, defaults=prompt_defaults)

    assert args.input.endswith(".yaml"), f"Expected a YAML input config, got: {args.input}"
    dotlist = [
        f"input_cfg={args.input}",
        "force_finite=True",
        "metadata_only=True",
        f"text_field={args.text_field}",
        f"lang_field={args.lang_field}",
        # Propagate to NeMoMultimodalConversation so its total_length / input_length
        # can convert audio turns to token equivalents (otherwise the constraint hits
        # the "token_equivalent_duration must be set" assert in _compute_num_audio_tokens).
        f"token_equivalent_duration={args.token_equivalent_duration}",
    ]
    if args.audio_locator_tag is not None:
        dotlist.append(f"audio_locator_tag={args.audio_locator_tag}")
    config = OmegaConf.merge(
        OmegaConf.structured(LhotseDataLoadingConfig),
        OmegaConf.from_dotlist(dotlist),
    )
    cuts, _ = read_cutset_from_config(config)

    duration_filter = RejectionsCounter(DurationFilter(args.min_duration, args.max_duration), "Duration filtering")
    cuts = cuts.filter(duration_filter)
    cuts = cuts.map(partial(apply_tokenizer, tokenizer=tokenizer, prompt=prompt), apply_fn=None)
    token_filter = RejectionsCounter(
        TokenCountFilter(args.min_tokens, args.max_tokens, args.measure_total_length), "Token count filtering"
    )
    cuts = cuts.filter(token_filter)
    tpt_filter = RejectionsCounter(TokenPerTokenFilter(-1, args.max_tpt), "Output tokens per input token filtering")
    cuts = cuts.filter(tpt_filter)
    if (N := args.num_examples) > 0:
        cuts = islice(cuts, N)

    is_2d = args.sub_buckets is not None
    if is_2d:
        bins = estimate_token_buckets_2d(
            cuts,
            num_buckets=args.buckets,
            num_subbuckets=args.sub_buckets,
            token_equivalent_duration=args.token_equivalent_duration,
            measure_total_length=args.measure_total_length,
            token_outlier_threshold=args.token_outlier_threshold,
            quiet=args.quiet,
        )
    else:
        bins = estimate_token_buckets_1d(
            cuts,
            num_buckets=args.buckets,
            token_equivalent_duration=args.token_equivalent_duration,
            measure_total_length=args.measure_total_length,
            quiet=args.quiet,
        )

    if args.quantize_bins != "none":
        before = len(bins)
        bins = quantize_bins(bins, args.quantize_bins)
        if not args.quiet:
            print(f"Quantization '{args.quantize_bins}': {before} -> {len(bins)} bins after dedupe.")

    if is_2d:
        bins_str = "[" + ",".join(f"[{b:d},{sb:d}]" for b, sb in bins) + "]"
    else:
        bins_str = "[" + ",".join(f"{b:d}" for b in bins) + "]"
    num_buckets_total = len(bins)

    if args.quiet:
        print(bins_str)
        maybe_patch_config(args.source_config, args.output_config, bins, num_buckets_total)
        return

    duration_filter.print_report()
    token_filter.print_report()
    tpt_filter.print_report()
    print("Use the following options in your config:")
    print(f"\tuse_bucketing=1")
    print(f"\tnum_buckets={num_buckets_total}")
    print(f"\tbucket_duration_bins={bins_str}")
    maybe_patch_config(args.source_config, args.output_config, bins, num_buckets_total)


if __name__ == "__main__":
    main()