# 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.

import argparse
import ast
import math
import warnings
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
from lhotse.cut import Cut
from omegaconf import OmegaConf

from nemo.collections.common.data import apply_prompt_format_fn
from nemo.collections.common.data.lhotse.cutset import read_cutset_from_config
from nemo.collections.common.data.lhotse.dataloader import LhotseDataLoadingConfig, tokenize
from nemo.collections.common.data.lhotse.sampling import DurationFilter, FixedBucketBatchSizeConstraint2D
from nemo.collections.common.prompts.formatter import PromptFormatter
from nemo.collections.common.tokenizers import (
    AggregateTokenizer,
    CanaryTokenizer,
    SentencePieceTokenizer,
    TokenizerSpec,
)
from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper


def parse_args():
    parser = argparse.ArgumentParser(
        description="Estimate duration bins for Lhotse dynamic bucketing using a sample of the input dataset. "
        "The dataset is read either from one or more manifest files and supports data weighting. "
        "Unlike estimate_duration_bins.py, this script prepares the setup for 2D bucketing. "
        "This means that each main bucket for audio duration is sub-divided into sub-buckets "
        "for the number of output tokens (supporting BPE and Aggregated tokenizers). "
        "2D bucketing is especially useful for encoder-decoder models where input audio duration is often "
        "not sufficient to stratify the sampling with an optimal GPU utilization.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "input",
        help='Data input. Options: '
        '1) "path.json" - any single NeMo manifest; '
        '2) "[[path1.json],[path2.json],...]" - any collection of NeMo manifests; '
        '3) "[[path1.json,weight1],[path2.json,weight2],...]" - any collection of weighted NeMo manifests; '
        '4) "input_cfg.yaml" - a new option supporting input configs, same as in model training \'input_cfg\' arg; '
        '5) "path/to/shar_data" - a path to Lhotse Shar data directory; '
        '6) "key=val" - in case none of the previous variants cover your case: "key" is the key you\'d use in NeMo training config with its corresponding value ',
    )
    parser.add_argument(
        "-t",
        "--tokenizer",
        nargs="+",
        required=True,
        help="Path to one or more SPE tokenizers. More than one means we'll use AggregateTokenizer and --langs argument must also be used. When provided, we'll estimate a 2D distribution for input and output sequence lengths.",
    )
    parser.add_argument(
        "-a", "--langs", nargs="+", help="Language names for each of AggregateTokenizer sub-tokenizers."
    )
    parser.add_argument(
        "-b",
        "--buckets",
        type=int,
        default=30,
        help="The desired number of buckets (dim0 => covers input sequence length / audio duration).",
    )
    parser.add_argument(
        "-s",
        "--sub-buckets",
        type=int,
        default=2,
        help="The desired number of sub-buckets (dim1 => covers output sequence length / num_tokens).",
    )
    parser.add_argument("--text-field", default="text", help="The key in manifests to read transcripts from.")
    parser.add_argument("--lang-field", default="lang", help="The key in manifests to read language from.")
    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_duration",
        type=float,
        default=-float("inf"),
        help="If specified, we'll filter out utterances shorter than this.",
    )
    parser.add_argument(
        "-u",
        "--max_duration",
        type=float,
        default=float("inf"),
        help="If specified, we'll filter out utterances longer than this.",
    )
    parser.add_argument(
        "--max_tps", type=float, default=None, help="Deprecated. TPS is automatically determined per bucket."
    )
    parser.add_argument(
        "--token_outlier_threshold",
        type=float,
        default=4.0,
        help="The lower this is, the more outliers in transcript token count will be filtered out. "
        "By default allow token counts at 4 sigma away from distribution mean, computed separately for every bucket.",
    )
    parser.add_argument(
        "-q", "--quiet", type=bool, default=False, help="When specified, only print the estimated duration bins."
    )
    parser.add_argument(
        "-f",
        "--prompt-format",
        type=str,
        help="When specified, we'll use a prompt formatter in addition to the tokenizer for the purpose of estimating token count bins. "
        "This is useful for accurate 2D bucket estimation with models such as EncDecMultiTaskModel (Canary-1B), "
        "or any model where the label sequence consists of a user prompt and a model's response.",
    )
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        help="Prompt slots provided as a Python list of dicts. It is used together with --prompt-format option."
        "For example, with Canary-1B you may use: [{'role':'user','slots':{'source_lang':'en','target_lang':'en','task':'asr','pnc':'yes'}]",
    )
    return parser.parse_args()


def sort_two_arrays(A, B):
    joint = np.rec.fromarrays([A, B])
    joint.sort()
    return joint.f0, joint.f1


def estimate_duration_buckets(
    cuts: Iterable[Cut],
    num_buckets: int,
    num_subbuckets: int,
    max_tps: float,
    max_duration: float,
    token_outlier_threshold: float,
    quiet: bool,
) -> list[tuple[float, float]]:
    """
    This function is based on lhotse.dataset.sampling.dynamic_bucketing.estimate_duration_buckets.
    It extends it to a 2D bucketing case.
    """
    assert num_buckets > 1

    constraint = FixedBucketBatchSizeConstraint2D([(0.0, 0.0)], [0])

    # Gather the duration and token count statistics for the dataset.
    sizes = []
    num_tokens = []
    for c in cuts:
        dur, toks = constraint.measure_length(c)
        sizes.append(dur)
        num_tokens.append(toks)
    sizes = np.array(sizes, dtype=np.float32)
    num_tokens = np.array(num_tokens, dtype=np.int32)
    sizes, num_tokens = sort_two_arrays(sizes, num_tokens)

    # We are building buckets with equal duration (empirically leads to more even bucket exhaustion over time).
    # We need to determine how much duration to allocate per bucket.
    size_per_bucket = sizes.sum() / num_buckets

    if not quiet:
        print("Duration 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, 0.995, 0.999]))
    if math.isinf(max_duration):
        max_duration = round(sizes[-1], 3)  # Round to 3 decimal places to be consistent for the output format.

    bins = []
    tps_thresholds = []
    bin_indexes = [0]
    tot = 0.0

    def _estimate_token_buckets(max_bucket_duration, start_idx, end_idx, corr_subbuckets=None):
        # Since this is 2D bucketing, apply the same bin creation logic
        # for the second dimension (i.e. token count) as for the first dimension (duration).
        # That means we aim to have each bucket contain roughly the same number of tokens.
        # Note that this estimation is biased towards more padding if you have
        # a lot of zero-token examples (e.g. non-speech).
        nonlocal bins

        if not corr_subbuckets:
            corr_subbuckets = num_subbuckets

        # Start by discarding outlier examples as defined by token-per-second (TPS) attribute.
        # We empirically determined high TPS examples to cause severe OOMs limiting batch sizes.
        # We cap the TPS for each top-level bucket at 4 standard deviations of TPS.
        # Examples exceeding that TPS value will be discarded during sampling at training time.
        num_tokens_bucket_all = num_tokens[start_idx:end_idx]
        sizes_bucket_all = sizes[start_idx:end_idx]
        non_outlier_indexes = find_non_outliers_z_score(
            num_tokens_bucket_all / sizes_bucket_all, threshold=token_outlier_threshold
        )
        num_tokens_bucket = num_tokens_bucket_all[non_outlier_indexes]
        sizes_bucket = sizes_bucket_all[non_outlier_indexes]
        max_tps_bucket = (num_tokens_bucket / sizes_bucket).max()
        num_tokens_bucket, sizes_bucket = sort_two_arrays(num_tokens_bucket, sizes_bucket)
        if not quiet:
            outlier_tps = np.delete(num_tokens_bucket_all / sizes_bucket_all, non_outlier_indexes)
            print(
                f"[bucket <= {max_bucket_duration:.2f}s] [{num_tokens_bucket.min()} - {num_tokens_bucket.max()}] [approx-max-tps: {max_tps_bucket:.2f}] Discarded {end_idx - start_idx - len(num_tokens_bucket)} max token outliers",
                end=" ",
            )
            if len(outlier_tps) > 0:
                print(f"min-outlier: {outlier_tps.min():.2f}, max-outlier: {outlier_tps.max():.2f}).", end="")
            print()

        tokens_per_subbucket = num_tokens_bucket.sum() / corr_subbuckets
        tot_toks = 0
        # Iterate over token counts, and whenever we hit tokens_per_subbucket, create a new 2D bucket bin.
        for num_toks, size in zip(num_tokens_bucket, sizes_bucket):
            # Threshold hit: we are creating a new (max_duration, max_num_tokens) bin.
            if tot_toks > tokens_per_subbucket:
                bins.append((max_bucket_duration, num_toks))
                tps_thresholds.append(max_tps_bucket)
                tot_toks = 0
            tot_toks += num_toks
        bins.append((max_bucket_duration, num_toks))
        tps_thresholds.append(max_tps_bucket)

    duration_bins = []

    # Iterate over data, and whenever we hit size_per_bucket, register it as a new duration bucket.
    for binidx, size in enumerate(sizes):
        if tot > size_per_bucket:
            size = round(size, 3)  # Round to 3 decimal places to be consistent for the output format.
            duration_bins.append(size)
            bin_indexes.append(binidx)
            tot = 0.0
        tot += size

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

    skipped_buckets = 1
    start_idx = 0

    # Iterate over newly created duration bins to handle cases where some bins have the same value —
    # this usually happens when the data is skewed.
    # If we detect such bins, we skip estimating token buckets for that particular bin.
    # Instead, we keep track of how many bins got skipped because they had the same duration.
    # Then, when we finally hit a bin with a different duration, we treat all those skipped bins as one "combined" bin.
    # For that combined bin, we create more subbuckets — specifically, the number of skipped bins × `num_subbuckets` (set by the user).
    #
    # Example of durations bins created from skewed duration distribution: [5, 20, 30, 30, 30, 40]
    # Here, we'd end up making token subbuckets for: [5, 20, 40]
    # where [20, 40] bucket will have 4 times more subbuckets (as we combined 4 buckets into 1) than usual bucket in that settings.

    for i, (duration_bin, binidx) in enumerate(zip(duration_bins, bin_indexes[1:])):
        if (i != len(duration_bins) - 1 and duration_bins[i + 1] == duration_bin) or (
            i == len(duration_bins) - 1 and max_duration == duration_bin
        ):
            skipped_buckets += 1
            continue
        _estimate_token_buckets(
            max_bucket_duration=duration_bin,
            start_idx=start_idx,
            end_idx=binidx,
            corr_subbuckets=num_subbuckets * skipped_buckets,
        )
        start_idx = binidx
        skipped_buckets = 1

    # Estimate an extra 2D bin set for global max duration.
    # Also, if the last value in duration_bins is equal to max_duration,
    # we need to make sure we properly handle any previously "skipped" buckets that ended at this max value.
    _estimate_token_buckets(
        max_bucket_duration=max_duration,
        start_idx=start_idx,
        end_idx=len(sizes),
        corr_subbuckets=num_subbuckets * skipped_buckets,
    )
    return bins, tps_thresholds


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


def load_tokenizer(paths: list[str], langs: list[str] = None, is_canary: bool = True) -> TokenizerSpec:
    if len(paths) == 1:
        tok = SentencePieceTokenizer(paths[0])
    else:
        assert langs is not None and len(paths) == len(
            langs
        ), f"Cannot create AggregateTokenizer; each tokenizer must have assigned a language via --langs option (we got --tokenizers={paths} and --langs={langs})"
        if is_canary:
            tokcls = CanaryTokenizer
        else:
            tokcls = AggregateTokenizer
        tok = tokcls({lang: SentencePieceTokenizer(p) for lang, p in zip(langs, paths)})
    return tok


def apply_tokenizer(cut, tokenizer=None, prompt: PromptFormatter = None):
    if prompt is not None:
        encoded = apply_prompt_format_fn(cut, prompt)
        cut.supervisions[0].tokens = encoded["input_ids"]

    elif tokenizer is not None:
        cut = tokenize(cut, TokenizerWrapper(tokenizer))

    return cut


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)

    if args.max_tps is not None:
        warnings.warn(
            "The option --max_tps has been deprecated in favor of "
            "automatic TPS determination that's variable across buckets."
        )

    tokenizer = None
    prompt = None
    if args.tokenizer is not None:
        tokenizer = load_tokenizer(
            paths=args.tokenizer,
            langs=args.langs,
            is_canary=args.prompt_format is not None and 'canary' in args.prompt_format,
        )
        if args.prompt_format is not None:
            prompt_defaults = None
            if args.prompt is not None:
                prompt_defaults = ast.literal_eval(args.prompt)
            prompt = PromptFormatter.resolve(args.prompt_format)(tokenizer, defaults=prompt_defaults)

    if '=' in args.input:
        inp_arg = args.input
    elif args.input.endswith(".yaml"):
        inp_arg = f"input_cfg={args.input}"
    elif Path(args.input).is_dir():
        inp_arg = f"shar_path={args.input}"
    else:
        inp_arg = f"manifest_filepath={args.input}"
    config = OmegaConf.merge(
        OmegaConf.structured(LhotseDataLoadingConfig),
        OmegaConf.from_dotlist(
            [inp_arg, "metadata_only=true", f"text_field={args.text_field}", f"lang_field={args.lang_field}"]
        ),
    )
    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))
    if (N := args.num_examples) > 0:
        cuts = islice(cuts, N)

    duration_bins, tps_thresholds = estimate_duration_buckets(
        cuts,
        num_buckets=args.buckets,
        num_subbuckets=args.sub_buckets,
        max_duration=args.max_duration,
        max_tps=args.max_tps,
        token_outlier_threshold=args.token_outlier_threshold,
        quiet=args.quiet,
    )
    duration_bins = "[" + ','.join(f"[{b:.3f},{sb:d}]" for b, sb in duration_bins) + "]"
    tps_thresholds = "[" + ",".join(f"{t:.2f}" for t in tps_thresholds) + "]"
    if not args.quiet:
        duration_filter.print_report()
    print("Use the following options in your config:")
    print(f"\tuse_bucketing=1")
    print(f"\tnum_buckets={args.buckets}")
    print(f"\tbucket_duration_bins={duration_bins}")
    print(f"The max_tps setting below is optional, use it if your data has low quality long transcript outliers:")
    print(f"\tmax_tps={tps_thresholds}")


if __name__ == "__main__":
    main()