#!/usr/bin/env python3
"""Reproduce the article's sensitivity analyses, downloadable data and figures.

Uses saved transcripts; it never submits reference text to a recognizer.
Install jiwer==4.0.0 transformers==5.17.0 matplotlib==3.10.8 num2words==0.5.14.

Standalone: download this file and benchmark-data.json into one directory, then:
    python reproduce.py --published-data benchmark-data.json --destination reproduced
This rescoring mode checks the published counts and recreates all figures;
it does not rerun inference or establish recognition accuracy on new audio.
"""

import argparse
import csv
import importlib.metadata
import json
import re
import statistics
import textwrap
import unicodedata
from pathlib import Path

try:
    from benchmark import basic_normalize, word_score, write_json
except ModuleNotFoundError:
    # The downloadable copy also runs beside benchmark-data.json, without the
    # private app repository. Keep these rules identical to benchmark.py.
    def basic_normalize(text):
        text = unicodedata.normalize("NFC", text).replace("’", "'").replace("‘", "'")
        return " ".join(re.sub(r"[^\w\s']", " ", text.lower()).split())

    def word_score(reference, hypothesis, normalize):
        import jiwer
        score = jiwer.process_words(normalize(reference), normalize(hypothesis))
        return dict(referenceWords=score.hits + score.substitutions + score.deletions,
                    substitutions=score.substitutions, deletions=score.deletions,
                    insertions=score.insertions, wer=score.wer)

    def write_json(path, value):
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n")

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2] if len(HERE.parents) >= 3 else HERE
DEFAULT_DESTINATION = ROOT / "website/blog/local-speech-benchmark/assets"
ENGINES = ["apple-speech-analyzer", "qwen3-asr-1.7b-q8_0-transcribe-cpp", "phonon-2-mlx", "parakeet-ultra-mlx"]
LABELS = ["Apple SpeechAnalyzer", "Qwen3-ASR 1.7B", "Phonon-2", "Parakeet Ultra"]
COLORS = ["#4869b4", "#bc571f", "#24735a", "#8061a8"]
INK, MUTED, PAPER, GRID = "#17251f", "#596760", "#fffdf8", "#e5e9e2"


def italian_format_normalize(text):
    """Symmetric, reference-independent rules, deliberately limited in scope."""
    from num2words import num2words

    text = basic_normalize(text)
    # Repair fused boundaries such as del2011, without repairing the number.
    text = re.sub(r"(?<=[^\W\d_])(?=\d)|(?<=\d)(?=[^\W\d_])", " ", text)
    text = re.sub(r"\b(\d+)\s+m\b", r"\1 metri", text)
    text = re.sub(r"\b\d+\b", lambda m: num2words(int(m[0]), lang="it"), text)
    return basic_normalize(text)


def total_scores(measurements, normalize):
    scores = [word_score(m["reference"], m["texts"][0], normalize) for m in measurements]
    totals = {k: sum(s[k] for s in scores)
              for k in ["referenceWords", "substitutions", "deletions", "insertions"]}
    totals["errors"] = sum(totals[k] for k in ["substitutions", "deletions", "insertions"])
    totals["wer"] = totals["errors"] / totals["referenceWords"]
    return totals


def analyze(baseline, ablation):
    from transformers.models.whisper.english_normalizer import EnglishTextNormalizer

    engines = {e["engine"]: e for e in baseline["engines"]}
    formatting = {}
    for key in ENGINES:
        subset = [m for m in engines[key]["measurements"] if m["locale"] == "it-IT"]
        formatting[key] = total_scores(subset, italian_format_normalize)
    clips = [m["id"] for m in engines[ENGINES[0]]["measurements"] if m["locale"] == "en-US"]
    leave_one_out = []
    for omitted in clips:
        row = {"omitted": omitted, "engines": {}}
        for key in ENGINES:
            subset = [m for m in engines[key]["measurements"]
                      if m["locale"] == "en-US" and m["id"] != omitted]
            words = sum(m["scores"][0]["referenceWords"] for m in subset)
            errors = sum(sum(m["scores"][0][k] for k in ["substitutions", "deletions", "insertions"])
                         for m in subset)
            row["engines"][key] = {"errors": errors, "referenceWords": words, "wer": errors / words}
        row["winner"] = min(row["engines"], key=lambda k: row["engines"][k]["wer"])
        leave_one_out.append(row)
    assert ablation["qwenAppleAlignmentEnabled"] is False
    assert len(ablation["engines"]) == 1 and ablation["engines"][0]["engine"] == ENGINES[1]
    expected_hashes = {f["id"]: f["audioSHA256"] for f in baseline["fixtureManifest"]}
    import hashlib
    assert {f["id"] for f in ablation["fixtures"]} == set(expected_hashes)
    for f in ablation["fixtures"]:
        assert hashlib.sha256(Path(f["audio"]).read_bytes()).hexdigest() == expected_hashes[f["id"]]
    references = {m["id"]: m["reference"] for m in engines[ENGINES[1]]["measurements"]}
    measured = [dict(m, reference=references[m["id"]]) for m in ablation["engines"][0]["measurements"]]
    no_alignment = {}
    for locale in ["en-US", "it-IT"]:
        subset = [m for m in measured if m["locale"] == locale]
        normalize = EnglishTextNormalizer({}) if locale == "en-US" else basic_normalize
        no_alignment[locale] = dict(
            total_scores(subset, normalize),
            warmTranscriptionSeconds=sum(statistics.median(m["processingSeconds"]) for m in subset),
            repeatWordScoresIdentical=all(
                len({json.dumps(word_score(m["reference"], t, normalize), sort_keys=True)
                     for t in m["texts"]}) == 1 for m in subset
            ),
        )
    # These invariants catch accidental forgiveness of real number errors.
    assert italian_format_normalize("15 metri del2011") == italian_format_normalize("quindici metri del 2011")
    assert italian_format_normalize("marzo207") != italian_format_normalize("marzo 2017")
    return dict(
        baseline=baseline,
        italianFormattingSensitivity=dict(
            normalization="NFC/case/punctuation; split letter-digit boundaries; numeric m -> metri; integer digits -> Italian cardinal words; same rules on references and hypotheses",
            package=f"num2words=={importlib.metadata.version('num2words')}", engines=formatting,
            limitations="Post-hoc sensitivity check, not a replacement leaderboard. Limited cardinal/unit rules do not implement full Italian inverse text normalization.",
        ),
        englishLeaveOneClipOut=leave_one_out,
        qwenWithoutAppleEvidence=dict(
            createdAt=ablation["createdAt"], macOS=ablation["macOS"],
            physicalMemoryBytes=ablation["physicalMemoryBytes"],
            fixtureIDs=list(expected_hashes), measurements=measured, summaries=no_alignment,
            modelPreparationSeconds=ablation["engines"][0]["modelPreparationSeconds"],
            warmupSecondsByLocale=ablation["engines"][0]["warmupSecondsByLocale"],
            protocol="Same signed Debug app and verified audio; Qwen-only; Apple evidence disabled; two warm measured passes per fixture after a locale warmup; automatic diarization recomputed once per fixture",
            caveat="Changes speech-window segmentation, evidence-based text correction and fallback together; not an isolated timestamp ablation. Fixed order and a later process make timing directional.",
        ),
    )


def figures(data, destination):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np

    plt.rcParams.update({"font.family": ["DejaVu Sans", "Arial", "sans-serif"], "font.size": 12,
                         "text.color": INK, "axes.labelcolor": MUTED,
                         "xtick.color": MUTED, "ytick.color": INK,
                         "figure.facecolor": PAPER, "axes.facecolor": PAPER,
                         "svg.fonttype": "none", "svg.hashsalt": "meetingstow-transcription-article",
                         "savefig.facecolor": PAPER})
    engines = {e["engine"]: e for e in data["baseline"]["engines"]}
    def polish(ax, xlabel, limit=None):
        ax.set_xlabel(xlabel, labelpad=12)
        ax.set_axisbelow(True)
        ax.grid(axis="x", color=GRID)
        ax.tick_params(axis="both", length=0, pad=9)
        for spine in ax.spines.values(): spine.set_visible(False)
        if limit: ax.set_xlim(*limit)
    def save(fig, name):
        fig.savefig(destination / f"{name}.svg", bbox_inches="tight", pad_inches=.25, metadata={"Date": None})
        fig.savefig(destination / f"{name}.png", dpi=200, bbox_inches="tight", pad_inches=.25)
        # Stack the two language panels for readable phone-sized figures.
        if len(fig.axes) == 2:
            fig.set_size_inches(7.4, 9.4)
            grid = fig.add_gridspec(2, 1)
            for index, ax in enumerate(fig.axes): ax.set_subplotspec(grid[index])
        else:
            fig.set_size_inches(7.4, 5.8)
        fig._suptitle.set_text(textwrap.fill(fig._suptitle.get_text(), width=44))
        fig._suptitle.set_fontsize(18)
        for ax in fig.axes:
            ax.set_xlabel(textwrap.fill(ax.get_xlabel(), width=48))
            ax.xaxis.label.set_fontsize(12)
            ax.tick_params(labelsize=13)
        fig.savefig(destination / f"{name}-mobile.svg", bbox_inches="tight", pad_inches=.2, metadata={"Date": None})
        for filename in [f"{name}.svg", f"{name}-mobile.svg"]:
            path = destination / filename
            path.write_text("\n".join(line.rstrip() for line in path.read_text().splitlines()) + "\n")
        plt.close(fig)
    def two_panels(title):
        fig, axes = plt.subplots(1, 2, figsize=(13.2, 4.7), layout="constrained")
        fig.suptitle(title, fontsize=20, fontweight="bold", x=.01, ha="left")
        return fig, axes

    fig, axes = two_panels("Word error rates for four transcription engines")
    for ax, locale, title in zip(axes, ["en-US", "it-IT"], ["English · 6 meeting excerpts", "Italian · 20 read-speech clips"]):
        values = [engines[k]["summaries"][locale]["wer"] * 100 for k in ENGINES]
        ax.barh(LABELS, values, color=COLORS, height=.55)
        for i, value in enumerate(values): ax.text(value + .35, i, f"{value:.2f}%", va="center", fontweight="bold")
        ax.invert_yaxis(); ax.set_title(title, loc="left", pad=18, fontsize=14)
        polish(ax, "Word error rate (%) · lower is better", (0, 20))
    save(fig, "word-error-rate")

    fig, axes = two_panels("Transcription time, including speaker separation")
    for ax, locale, title in zip(axes, ["en-US", "it-IT"], ["English · 360 seconds of audio", "Italian · 275.16 seconds of audio"]):
        for i, key in enumerate(ENGINES):
            s = engines[key]["summaries"][locale]
            shared, alignment = s["sharedDiarizationSeconds"], s["appleAlignmentSeconds"]
            core = s["warmTranscriptionSeconds"] - alignment
            ax.barh(i, shared, color="#adb8ad", height=.55, label="Shared speaker analysis" if i==0 else None)
            if alignment: ax.barh(i, alignment, left=shared, color="#a7b9df", height=.55, label="Apple evidence for Qwen")
            ax.barh(i, core, left=shared+alignment, color=COLORS[i], height=.55)
            ax.text(s["warmPipelineSeconds"]+.35, i, f'{s["warmPipelineSeconds"]:.2f}s', va="center", fontweight="bold")
        ax.set_yticks(range(len(ENGINES)), LABELS); ax.invert_yaxis()
        ax.set_title(title, loc="left", pad=18, fontsize=14)
        polish(ax, "Warm pipeline time (seconds) · lower is better", (0, 27))
        ax.legend(frameon=False, fontsize=9, loc="upper left", bbox_to_anchor=(0, -.26))
    save(fig, "pipeline-time")

    fig, axes = two_panels("Missing, changed, and added words")
    for ax, locale, title in zip(axes, ["en-US", "it-IT"], ["English · 775 reference words", "Italian · 493 reference words"]):
        left = np.zeros(len(ENGINES))
        for field, label, color in [("deletions", "Missing", "#24735a"), ("substitutions", "Changed", "#bc571f"), ("insertions", "Added", "#4869b4")]:
            values = np.array([engines[k]["summaries"][locale][field] for k in ENGINES])
            ax.barh(LABELS, values, left=left, color=color, label=label, height=.55)
            for i, value in enumerate(values):
                if value >= 6: ax.text(left[i]+value/2, i, str(value), color="white", ha="center", va="center", fontsize=11, fontweight="bold")
            left += values
        ax.invert_yaxis(); ax.set_title(title, loc="left", pad=18, fontsize=14)
        polish(ax, "Word errors (count)", (0, 145))
        ax.legend(frameon=False, fontsize=10, loc="upper left", bbox_to_anchor=(0, -.26), ncol=3)
    save(fig, "error-types")

    fig, ax = plt.subplots(figsize=(10, 4.5), layout="constrained")
    fig.suptitle("Italian scores with and without number normalization", fontsize=18, fontweight="bold", x=.01, ha="left")
    y = np.arange(len(ENGINES))
    strict = [engines[k]["summaries"]["it-IT"]["wer"] * 100 for k in ENGINES]
    relaxed = [data["italianFormattingSensitivity"]["engines"][k]["wer"] * 100 for k in ENGINES]
    for i in range(len(ENGINES)):
        ax.barh(i-.16, strict[i], height=.28, color=COLORS[i], alpha=.3)
        ax.barh(i+.16, relaxed[i], height=.28, color=COLORS[i])
        ax.text(strict[i]+.2, i-.16, f"{strict[i]:.2f}% original", va="center", fontsize=11)
        ax.text(relaxed[i]+.2, i+.16, f"{relaxed[i]:.2f}% formatting-aware", va="center", fontsize=11, fontweight="bold")
    ax.set_yticks(y, LABELS); ax.invert_yaxis()
    polish(ax, "Word error rate (%) · identical rules applied to both sides", (0, 17))
    save(fig, "italian-formatting")

    fig, axes = two_panels("Qwen scores with and without Apple evidence")
    for ax, locale, title in zip(axes, ["en-US", "it-IT"], ["English", "Italian"]):
        current = engines[ENGINES[1]]["summaries"][locale]
        alt = data["qwenWithoutAppleEvidence"]["summaries"][locale]
        values = [current["wer"]*100, alt["wer"]*100]
        ax.barh(["Apple-assisted Qwen", "Without Apple evidence"], values, color=[COLORS[1], "#d8a17b"], height=.5)
        for i, s in enumerate([current, alt]):
            ax.text(values[i]+.3, i, f'{values[i]:.2f}% · {s["warmTranscriptionSeconds"]:.2f}s', va="center", fontsize=12, fontweight="bold")
        ax.invert_yaxis(); ax.set_title(title, loc="left", pad=18, fontsize=14)
        polish(ax, "Word error rate (%) · lower is better", (0, 25))
    save(fig, "qwen-pipeline-experiment")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--baseline", type=Path, default=HERE/"results/2026-10-19/results.json")
    parser.add_argument("--ablation", type=Path)
    parser.add_argument("--published-data", type=Path,
                        help="Rescore downloadable benchmark-data.json without the app repository")
    parser.add_argument("--destination", type=Path)
    args = parser.parse_args()
    if args.published_data:
        from transformers.models.whisper.english_normalizer import EnglishTextNormalizer
        data = json.loads(args.published_data.read_text())
        for engine in data["baseline"]["engines"]:
            for measurement in engine["measurements"]:
                normalize = EnglishTextNormalizer({}) if measurement["locale"] == "en-US" else basic_normalize
                scores = [word_score(measurement["reference"], t, normalize) for t in measurement["texts"]]
                assert scores == measurement["scores"], (engine["engine"], measurement["id"])
            for locale in ["en-US", "it-IT"]:
                subset = [m for m in engine["measurements"] if m["locale"] == locale]
                normalize = EnglishTextNormalizer({}) if locale == "en-US" else basic_normalize
                totals = total_scores(subset, normalize)
                assert all(engine["summaries"][locale][k] == v for k, v in totals.items())
                seconds = sum(statistics.median(m["processingSeconds"]) for m in subset)
                assert seconds == engine["summaries"][locale]["warmTranscriptionSeconds"]
            italian = [m for m in engine["measurements"] if m["locale"] == "it-IT"]
            assert total_scores(italian, italian_format_normalize) == data["italianFormattingSensitivity"]["engines"][engine["engine"]]
        for locale in ["en-US", "it-IT"]:
            alternative = data["qwenWithoutAppleEvidence"]
            subset = [m for m in alternative["measurements"] if m["locale"] == locale]
            normalize = EnglishTextNormalizer({}) if locale == "en-US" else basic_normalize
            assert all(alternative["summaries"][locale][k] == v
                       for k, v in total_scores(subset, normalize).items())
        args.destination = args.destination or Path.cwd()/"reproduced"
    else:
        if not args.ablation:
            parser.error("--ablation is required unless --published-data is provided")
        data = analyze(json.loads(args.baseline.read_text()), json.loads(args.ablation.read_text()))
        args.destination = args.destination or DEFAULT_DESTINATION
    args.destination.mkdir(parents=True, exist_ok=True)
    if not args.published_data:
        write_json(HERE/"article-results.json", data)
    write_json(args.destination/"benchmark-data.json", data)
    with (args.destination/"benchmark-summary.csv").open("w", newline="") as f:
        writer = csv.writer(f, lineterminator="\n")
        writer.writerow(["language", "engine", "word_error_rate_percent", "rating_out_of_10", "word_errors", "reference_words", "warm_transcription_seconds", "shared_diarization_seconds", "warm_pipeline_seconds"])
        by_engine = {e["engine"]: e for e in data["baseline"]["engines"]}
        for locale in ["en-US", "it-IT"]:
            for key, label in zip(ENGINES, LABELS):
                s = by_engine[key]["summaries"][locale]
                writer.writerow([locale, label, s["wer"]*100, s["accuracyScoreOutOf10"], s["errors"], s["referenceWords"], s["warmTranscriptionSeconds"], s["sharedDiarizationSeconds"], s["warmPipelineSeconds"]])
    figures(data, args.destination)
    samples = {}
    for locale, scope in [("en-US", "6 meeting excerpts · 360 seconds · 775 reference words"),
                          ("it-IT", "20 read-speech clips · 275.16 seconds · 493 reference words")]:
        rows = []
        for key, label, color in zip(ENGINES, LABELS, COLORS):
            s = by_engine[key]["summaries"][locale]
            rows.append(dict(name=label, color=color, wer=s["wer"] * 100,
                             time=s["warmTranscriptionSeconds"], errors=s["errors"]))
        samples[locale] = dict(scope=scope, rows=rows)
    (args.destination/"score-data.js").write_text(
        "// Generated by article_analysis.py from the scored native measurements.\n"
        + "globalThis.meetingstowBenchmarkSamples = "
        + json.dumps(samples, ensure_ascii=False, indent=2) + ";\n"
    )
    (args.destination/"reproduce.py").write_text(Path(__file__).read_text())
    print(json.dumps({"formatting": data["italianFormattingSensitivity"]["engines"],
                      "qwen_without_apple": data["qwenWithoutAppleEvidence"]["summaries"],
                      "leave_one_out": data["englishLeaveOneClipOut"]}, indent=2))


if __name__ == "__main__":
    main()