#!/usr/bin/env python3
"""Score an Echo suite result on the reproducible public half of AA-WER v2."""

from __future__ import annotations

import hashlib
import json
import pathlib
import platform
import re
import statistics
import subprocess
import sys
from datetime import datetime, timezone

from jiwer import process_words
from whisper_normalizer.english import EnglishTextNormalizer


ROOT = pathlib.Path(__file__).resolve().parent
NORMALIZER = EnglishTextNormalizer()


def file_sha256(path: pathlib.Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def normalized(text: str) -> str:
    return " ".join(NORMALIZER(text).split())


def weighted_mean(rows: list[dict], value_key: str, weight_key: str) -> float:
    total_weight = sum(row[weight_key] for row in rows)
    return sum(row[value_key] * row[weight_key] for row in rows) / total_weight


def main() -> None:
    result_path = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else ROOT / "echo-coreml-results.json")
    timing_path = ROOT / "echo-coreml-results.stderr.txt"
    sources = json.loads((ROOT / "sources.json").read_text(encoding="utf-8"))
    suite = json.loads(result_path.read_text(encoding="utf-8"))
    source_by_id = {row["fixture_id"]: row for row in sources["rows"]}
    scored: list[dict] = []

    for result in suite["results"]:
        source = source_by_id[result["id"]]
        reference = normalized(source["reference"])
        hypothesis = normalized(result["rawTranscript"])
        measurement = process_words(reference, hypothesis)
        scored.append(
            {
                "id": result["id"],
                "dataset": source["dataset"],
                "duration_seconds": source["duration_seconds"],
                "reference_words": len(reference.split()),
                "hypothesis_words": len(hypothesis.split()),
                "substitutions": measurement.substitutions,
                "deletions": measurement.deletions,
                "insertions": measurement.insertions,
                "hits": measurement.hits,
                "wer": measurement.wer,
                "model_ms": result["modelMs"],
                "audio_duration_ms": result["audioDurationMs"],
                "real_time_factor": result["realTimeFactor"],
                "normalized_reference": reference,
                "normalized_hypothesis": hypothesis,
                "raw_transcript": result["rawTranscript"],
            }
        )

    dataset_scores = {}
    for dataset in ("voxpopuli", "earnings22"):
        rows = [row for row in scored if row["dataset"] == dataset]
        dataset_scores[dataset] = {
            "samples": len(rows),
            "duration_seconds": sum(row["duration_seconds"] for row in rows),
            "duration_weighted_wer": weighted_mean(rows, "wer", "duration_seconds"),
            "micro_wer": (
                sum(row["substitutions"] + row["deletions"] + row["insertions"] for row in rows)
                / sum(row["reference_words"] for row in rows)
            ),
            "mean_model_ms": statistics.mean(row["model_ms"] for row in rows),
            "audio_seconds_per_model_second": (
                sum(row["duration_seconds"] for row in rows)
                / (sum(row["model_ms"] for row in rows) / 1000)
            ),
        }

    public_composite = statistics.mean(
        [
            dataset_scores["voxpopuli"]["duration_weighted_wer"],
            dataset_scores["earnings22"]["duration_weighted_wer"],
        ]
    )
    try:
        chip = subprocess.check_output(
            ["sysctl", "-n", "machdep.cpu.brand_string"], text=True
        ).strip()
    except Exception:
        chip = "unknown"
    timing_text = timing_path.read_text(encoding="utf-8") if timing_path.exists() else ""
    wall_match = re.search(r"^real\s+([0-9.]+)$", timing_text, flags=re.MULTILINE)
    whole_run_wall_seconds = float(wall_match.group(1)) if wall_match else None
    source_duration_seconds = sum(row["duration_seconds"] for row in scored)
    report = {
        "title": "Echo Core ML — AA-WER v2 public-subset reconstruction",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "status": "measured-public-subset-only",
        "model": suite["model"],
        "pipeline": suite["pipeline"],
        "hardware": {
            "machine": platform.machine(),
            "processor": chip,
            "macos": platform.mac_ver()[0],
        },
        "protocol": {
            "included": [
                "VoxPopuli-Cleaned-AA v1: all 628 samples",
                "Earnings22-Cleaned-AA v1: all 6 parent recordings",
            ],
            "excluded": [
                "AA-AgentTalk: proprietary, 50% of the official AA-WER score",
            ],
            "aggregation": (
                "Duration-weighted sample WER inside each dataset; equal-weight "
                "mean of the two public dataset scores."
            ),
            "normalization": (
                "Open-source whisper-normalizer EnglishTextNormalizer 0.1.15. "
                "Artificial Analysis' full current custom normalizer is not public, "
                "so this is not an official AA-WER score."
            ),
        },
        "dataset_scores": dataset_scores,
        "public_subset_composite_wer": public_composite,
        "timing": {
            "source_duration_seconds": source_duration_seconds,
            "whole_run_wall_seconds": whole_run_wall_seconds,
            "audio_seconds_per_wall_second": (
                source_duration_seconds / whole_run_wall_seconds
                if whole_run_wall_seconds
                else None
            ),
            "note": (
                "The production runner reported zero audioDurationMs for long MP3 "
                "parents. Speed factors therefore use the dataset's pinned source "
                "durations; transcript and WER scoring are unaffected."
            ),
        },
        "official_aa_wer_eligible": False,
        "result_file": {
            "path": str(result_path),
            "sha256": file_sha256(result_path),
        },
        "rows": scored,
    }
    output_path = ROOT / "echo-coreml-aa-public-scored.json"
    output_path.write_text(
        json.dumps(report, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    summary = {key: value for key, value in report.items() if key != "rows"}
    (ROOT / "echo-coreml-aa-public-summary.json").write_text(
        json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()
