#!/usr/bin/env python3
"""Multilingual LibriSpeech (8 languages) WER for a CAFE .nemo model.

One (model, language) per invocation; prints WER% / RTFx and optionally writes a JSON result.
Decoding and scoring mirror the Open ASR Leaderboard multilingual runner exactly — greedy batched
decoding with the CUDA-graph decoder off, the official English normalizer for `en` and
`ml_normalizer` otherwise, empty references dropped, compound-word pairing before the WER call —
so the numbers are comparable to the leaderboard and to the model card.

  export PYTHONPATH=/path/to/open_asr_leaderboard:$PYTHONPATH
  python eval_mls8.py --model_id ../CAFE-25lang.nemo --lang de \
      --adapters ../adapters.safetensors

Dependencies: nemo_toolkit[asr], datasets, evaluate, jiwer, soundfile, and the `normalizer`
package from https://github.com/huggingface/open_asr_leaderboard (on PYTHONPATH).
"""
from __future__ import annotations

import argparse
import io
import json
import os
import time

os.environ.setdefault("DATASETS_USE_TORCHCODEC", "0")

import evaluate
import numpy as np
import soundfile
import torch
from datasets import Audio, load_dataset
from nemo.collections.asr.models import ASRModel
from normalizer import data_utils
from normalizer.eval_utils import normalize_compound_pairs
from omegaconf import OmegaConf

# The 8 MLS languages, each from the source the leaderboard uses for it.
LANG_SOURCE = {
    "en": ("parler-tts/mls_eng", "default"),
    "de": ("facebook/multilingual_librispeech", "german"),
    "nl": ("nithinraok/asr-leaderboard-datasets", "mls_nl"),
    "pl": ("nithinraok/asr-leaderboard-datasets", "mls_pl"),
    "es": ("nithinraok/asr-leaderboard-datasets", "mls_es"),
    "fr": ("nithinraok/asr-leaderboard-datasets", "mls_fr"),
    "it": ("nithinraok/asr-leaderboard-datasets", "mls_it"),
    "pt": ("nithinraok/asr-leaderboard-datasets", "mls_pt"),
}
TEXT_COLS = ("text", "transcript", "sentence")


def pick(keys, candidates):
    for c in candidates:
        if c in keys:
            return c
    return None


def main(args):
    lang = args.lang
    dataset_name, config_name = LANG_SOURCE[lang]
    wer_metric = evaluate.load("wer")
    cache_dir = os.path.join(args.work_dir, "audio_cache", f"{config_name}_{lang}")
    os.makedirs(cache_dir, exist_ok=True)

    device = torch.device(f"cuda:{args.device}") if args.device >= 0 else torch.device("cpu")
    dtype = torch.bfloat16 if args.device >= 0 else torch.float32
    model = (ASRModel.restore_from(args.model_id, map_location=device)
             if args.model_id.endswith(".nemo")
             else ASRModel.from_pretrained(args.model_id, map_location=device))
    model.to(dtype).eval()

    if args.adapters:                     # 25-language model, non-English path
        from adapter_wrapper import attach_adapters
        attach_adapters(model, args.adapters, device, dtype)

    if model.cfg.decoding.strategy != "beam":
        model.cfg.decoding.strategy = "greedy_batch"
        if hasattr(model.cfg.decoding, "greedy"):
            OmegaConf.update(model.cfg.decoding, "greedy.use_cuda_graph_decoder", False,
                             force_add=True)
        model.change_decoding_strategy(model.cfg.decoding)

    print(f"Loading {dataset_name} [{config_name}] split={args.split} (lang={lang})")
    ds = load_dataset(dataset_name, config_name, split=args.split, streaming=False)
    ds = ds.cast_column("audio", Audio(decode=False))
    text_col = pick(ds.column_names, TEXT_COLS)
    if text_col is None:
        raise ValueError(f"no text column in {ds.column_names}")
    if args.max_eval_samples and args.max_eval_samples > 0:
        ds = ds.select(range(min(args.max_eval_samples, len(ds))))

    def prepare(batch, idx):
        paths, durs = [], []
        for i, sample in zip(idx, batch["audio"]):
            path = os.path.join(cache_dir, f"{config_name}_{i}.wav")
            if "array" in sample:
                arr, sr = np.float32(sample["array"]), sample.get("sampling_rate", 16000)
            elif sample.get("bytes") is not None:
                with io.BytesIO(sample["bytes"]) as fh:
                    arr, sr = soundfile.read(fh, dtype="float32")
            elif sample.get("path"):
                arr, sr = soundfile.read(sample["path"], dtype="float32")
            else:
                raise ValueError("sample lacks array/bytes/path")
            if getattr(arr, "ndim", 1) > 1:
                arr = arr.mean(axis=1)
            if not os.path.exists(path):
                soundfile.write(path, arr, sr)
            paths.append(path)
            durs.append(len(arr) / sr)
        return {"audio_filepaths": paths, "durations": durs,
                "references": list(batch[text_col])}

    print("Materialising audio cache...")
    ds = ds.map(prepare, with_indices=True, batched=True, batch_size=args.batch_size,
                remove_columns=ds.column_names)
    files, durs, refs0 = list(ds["audio_filepaths"]), list(ds["durations"]), list(ds["references"])
    order = sorted(range(len(durs)), key=lambda k: durs[k], reverse=True)   # longest first
    files = [files[i] for i in order]
    refs0 = [refs0[i] for i in order]
    durs = [durs[i] for i in order]

    total_time, hyps = 0.0, []
    for warm in range(2):                 # first pass warms up, second pass is timed
        batch = files[: args.batch_size * 4] if warm == 0 else files
        t0 = time.time()
        with torch.inference_mode(), torch.no_grad():
            hyps = model.transcribe(batch, batch_size=args.batch_size, verbose=False,
                                    num_workers=1)
        if warm == 1:
            total_time = time.time() - t0
    if isinstance(hyps, tuple) and len(hyps) == 2:
        hyps = hyps[0]

    if lang == "en":
        refs = [data_utils.normalizer(r) for r in refs0]
        preds = [data_utils.normalizer(h.text) for h in hyps]
    else:
        refs = [data_utils.ml_normalizer(r, lang=lang) for r in refs0]
        preds = [data_utils.ml_normalizer(h.text, lang=lang) for h in hyps]
    pairs = [(r, p) for r, p in zip(refs, preds) if data_utils.is_target_text_in_range(r)]
    refs, preds = (list(x) for x in zip(*pairs)) if pairs else ([], [])
    wr, wp = normalize_compound_pairs(refs, preds)
    wer = round(100 * wer_metric.compute(references=wr, predictions=wp), 2)
    rtfx = round(sum(durs) / total_time, 2) if total_time > 0 else 0.0

    result = {"model_name": args.model_name or args.model_id, "lang": lang,
              "dataset": dataset_name, "config": config_name, "split": args.split,
              "adapters": bool(args.adapters), "num_samples": len(refs),
              "wer": wer, "rtfx": rtfx}
    print("RESULT_JSON " + json.dumps(result))
    print(f"WER: {wer}%  RTFX: {rtfx}  (lang={lang}, n={len(refs)})")
    if args.result_json:
        os.makedirs(os.path.dirname(os.path.abspath(args.result_json)), exist_ok=True)
        with open(args.result_json, "w") as f:
            json.dump(result, f, indent=2)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model_id", required=True, help="path to .nemo, or a Hub model id")
    p.add_argument("--model_name", default=None)
    p.add_argument("--lang", required=True, choices=list(LANG_SOURCE))
    p.add_argument("--adapters", default="", help="25-language model: adapters.safetensors "
                                                 "(use for every language except en)")
    p.add_argument("--split", default="test")
    p.add_argument("--device", type=int, default=0, help="-1 for CPU")
    p.add_argument("--batch_size", type=int, default=32)
    p.add_argument("--max_eval_samples", type=int, default=-1)
    p.add_argument("--work_dir", default=os.getcwd())
    p.add_argument("--result_json", default=None)
    main(p.parse_args())