"""
Step 4: measure accuracy on the test set (reciters never trained on + held-out Warsh surahs).

    python evaluate.py --config config/pilot.yaml                 # PyTorch, best.pt
    python evaluate.py --config config/pilot.yaml --onnx          # the exported int8 model, via sherpa-onnx
    python evaluate.py --config config/pilot.yaml --baseline      # the app's current NVIDIA model, for comparison

Reports the letter error rate (CER), word error rate (WER) and skeleton error rate (letters
without alif/hamza, which is what tracking relies on) per riwaya and per reciter.
Both the model output and the reference are letters-only, so every model is scored the same way.
"""
import os
import sys

import torch

from quran_asr import audio, config as config_mod, data, features, model as model_mod, text


def main():
    parser = config_mod.from_args(__doc__)
    parser.add_argument("--onnx", action="store_true", help="evaluate runs/<name>/export/model.int8.onnx")
    parser.add_argument("--baseline", action="store_true", help="evaluate the NVIDIA FastConformer instead")
    parser.add_argument("--checkpoint", default="best.pt")
    parser.add_argument("--limit", type=int, default=0, help="only the first N clips")
    args = parser.parse_args()
    cfg = config_mod.load(args.config)
    test = data.read_manifest(os.path.join(config_mod.manifest_dir(cfg), "test.jsonl"))
    if args.limit:
        test = test[:: max(1, len(test) // args.limit)][: args.limit]

    if args.onnx or args.baseline:
        import sherpa_onnx
        if args.baseline:
            folder = os.path.join(cfg.paths.data_dir, "nvidia_fastconformer")
            rec = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(
                model=os.path.join(folder, "model.onnx"), tokens=os.path.join(folder, "tokens.txt"), num_threads=8)
        else:
            out = os.path.join(cfg.paths.run_dir, "export")
            rec = sherpa_onnx.OfflineRecognizer.from_zipformer_ctc(
                model=os.path.join(out, "model.int8.onnx"), tokens=os.path.join(out, "tokens.txt"), num_threads=8)

        def transcribe(wave):
            s = rec.create_stream()
            s.accept_waveform(features.SAMPLE_RATE, wave)
            rec.decode_stream(s)
            return text.normalize(s.result.text.replace(text.WORD, " "))
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        ckpt = torch.load(os.path.join(cfg.paths.run_dir, args.checkpoint), map_location="cpu", weights_only=False)
        model = model_mod.build(config_mod._wrap(ckpt["model_config"]), len(text.VOCAB))
        model.load_state_dict(ckpt["model"])
        model.to(device).eval()

        def transcribe(wave):
            feats = features.fbank(torch.from_numpy(wave)).to(device)
            with torch.no_grad():
                log_probs, lens = model(feats[None], torch.tensor([feats.shape[0]], device=device))
            return text.decode(model_mod.greedy_decode(log_probs, lens)[0])

    totals = {}
    for k, e in enumerate(test, 1):
        hyp = transcribe(audio.load(e["path"]))
        counts = text.error_counts(hyp, e["text"])
        for key in ("all", e["riwaya"], f"{e['riwaya']}/{e['reciter']}"):
            t = totals.setdefault(key, [0] * 6)
            for j in range(6):
                t[j] += counts[j]
        if k % 200 == 0:
            print(f"  {k}/{len(test)}", flush=True)

    name = "NVIDIA FastConformer (baseline)" if args.baseline else ("exported ONNX" if args.onnx else args.checkpoint)
    print(f"\n{name} on {len(test)} test clips")
    print("error rates; 'skeleton' ignores alif and hamza (spelling conventions), like the app's tracker")
    print(f"{'':52s} {'letters':>8s} {'words':>8s} {'skeleton':>9s}")
    for key, (ce, ct, we, wt, se, st) in sorted(totals.items(), key=lambda kv: (kv[0] != "all", kv[0])):
        print(f"{key:52s} {ce / max(1, ct) * 100:7.2f}% {we / max(1, wt) * 100:7.2f}% {se / max(1, st) * 100:8.2f}%")


if __name__ == "__main__":
    sys.exit(main())