# SPDX-License-Identifier: MIT
"""Verbatim scoring function and reference generation snippet from the saved v4 recipe.

This documents the protocol; it does not bundle a benchmark dataset.
"""
from jiwer import wer, cer

def score_predictions(references: list[dict], predictions: list[dict]) -> dict:
    ref_by_id = {r["item_id"]: r for r in references}
    by_id = {p["item_id"]: p["hypothesis"] for p in predictions}
    metrics = {}
    
    for bucket in ["mixed", "nepali", "english"]:
        rows = [r for r in references if r.get("language_bucket") == bucket]
        if not rows:
            continue
        refs = [r["text"] for r in rows]
        hyps = [by_id.get(r["item_id"], "") for r in rows]
        
        w = wer(refs, hyps) * 100.0
        c = cer(refs, hyps) * 100.0
        metrics[bucket] = {"wer": round(w, 2), "cer": round(c, 2), "clips": len(rows)}
        
    all_refs = [r["text"] for r in references]
    all_hyps = [by_id.get(r["item_id"], "") for r in references]
    overall_w = wer(all_refs, all_hyps) * 100.0
    overall_c = cer(all_refs, all_hyps) * 100.0
    
    metrics["overall"] = {"wer": round(overall_w, 2), "cer": round(overall_c, 2), "clips": len(references)}
    
    # Target weighted WER (60% mixed, 20% nepali, 20% english if present, else overall)
    if "mixed" in metrics and "nepali" in metrics:
        mix_w = metrics["mixed"]["wer"]
        nep_w = metrics["nepali"]["wer"]
        en_w = metrics.get("english", {}).get("wer", mix_w)
        metrics["target_weighted_wer"] = round(0.6 * mix_w + 0.2 * nep_w + 0.2 * en_w, 2)
    else:
        metrics["target_weighted_wer"] = round(overall_w, 2)
        
    return metrics


REFERENCE_GENERATION_SNIPPET = '            tokens = base.generate(\n                input_features=inputs["input_features"],\n                attention_mask=inputs.get("attention_mask"),\n                language=None,\n                task="transcribe",\n                return_timestamps=False,\n                max_length=self.args.generation_max_length,\n                num_beams=1,\n                use_cache=True,\n            )\n'