"""第 43 章：MMS-TTS 合成与 Whisper 识别；默认只核对随书记录。"""

from argparse import ArgumentParser
from hashlib import sha256
import json
from pathlib import Path
import platform
import re
from time import perf_counter
import wave


MODELS = {
    "tts": ("facebook/mms-tts-eng", "c71de0fe7204c83f1c10820a7d696d0b450048ba"),
    "asr": ("openai/whisper-tiny.en", "87c7102498dcde7456f24cfd30239ca606ed9063"),
}
TEXTS = [
    "A blue cup is on the table.",
    "The workshop opens at nine and closes at six.",
    "Please pick up order number three seven two.",
    "There are two red circles and one blue square.",
]
CONDITIONS = [("clean", None), ("snr20", 20), ("snr0", 0)]
DEFAULT = Path(__file__).parent / "data/speech/report.json"


def file_hash(path):
    h = sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def words(text):
    # 明确的教学规范化：小写，字母/数字连续串；不改写数字、同义词。
    return re.findall(r"[a-z0-9]+", text.lower())


def edit_counts(reference, hypothesis):
    a, b = words(reference), words(hypothesis)
    cost = [[0] * (len(b) + 1) for _ in range(len(a) + 1)]
    for i in range(len(a) + 1):
        cost[i][0] = i
    for j in range(len(b) + 1):
        cost[0][j] = j
    for i in range(1, len(a) + 1):
        for j in range(1, len(b) + 1):
            cost[i][j] = min(cost[i-1][j-1] + (a[i-1] != b[j-1]),
                             cost[i-1][j] + 1, cost[i][j-1] + 1)
    i, j = len(a), len(b)
    counts = {"substitutions": 0, "deletions": 0, "insertions": 0, "reference_words": i}
    while i or j:
        if i and j and cost[i][j] == cost[i-1][j-1] + (a[i-1] != b[j-1]):
            counts["substitutions"] += int(a[i-1] != b[j-1])
            i, j = i - 1, j - 1
        elif i and cost[i][j] == cost[i-1][j] + 1:
            counts["deletions"] += 1
            i -= 1
        else:
            counts["insertions"] += 1
            j -= 1
    counts["errors"] = cost[-1][-1]
    counts["wer"] = counts["errors"] / len(a) if a else None
    return counts


def write_wav(path, values, sample_rate=16000):
    import numpy as np
    peak = float(np.abs(values).max(initial=0))
    gain = min(1.0, 0.95 / peak) if peak else 1.0
    pcm = np.rint(values * gain * 32767).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(sample_rate)
        output.writeframes(pcm.tobytes())
    return gain


def read_wav(path):
    import numpy as np
    with wave.open(str(path), "rb") as audio:
        assert audio.getnchannels() == 1 and audio.getsampwidth() == 2
        return np.frombuffer(audio.readframes(audio.getnframes()), dtype="<i2").astype("float32") / 32767, audio.getframerate()


def add_noise(clean, snr_db, seed):
    import numpy as np
    generator = np.random.default_rng(seed)
    noise = generator.normal(size=len(clean))
    scale = np.sqrt(np.mean(clean ** 2) / (10 ** (snr_db / 10) * np.mean(noise ** 2)))
    return noise * scale


def validate(path=DEFAULT):
    report = json.loads(path.read_text())
    assert report["source_sha256"] == file_hash(__file__), "脚本与记录版本不同"
    assert report["models"] == {key: list(value) for key, value in MODELS.items()}
    assert report["texts"] == TEXTS and len(report["tts"]) == 4
    assert [row["id"] for row in report["asr"]] == [
        f"{i}-{name}" for i in range(4) for name, _ in CONDITIONS] + ["silence"]
    for row in report["asr"]:
        filename = path.parent / row["audio"]
        assert file_hash(filename) == row["audio_sha256"]
        with wave.open(str(filename), "rb") as audio:
            assert audio.getframerate() == 16000
            assert audio.getnchannels() == 1 and audio.getsampwidth() == 2
            assert audio.getnframes() == row["samples"]
        assert row["score"] == edit_counts(row["reference"], row["text"])
        assert row["eos"] == (row["token_ids"][-1] == report["eos_token_id"])
        assert row["feature_shape"] == [1, 80, 3000]
        assert row["seconds"] > 0
    return report


def run(folder):
    import numpy as np
    import torch
    import transformers
    from huggingface_hub import snapshot_download
    from transformers import AutoTokenizer, VitsModel, WhisperProcessor, WhisperForConditionalGeneration

    if folder.exists() and any(folder.iterdir()):
        raise ValueError("请使用空输出目录，保留已有实验结果")
    folder.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(4)
    torch.use_deterministic_algorithms(True)
    paths = {name: Path(snapshot_download(
        model, revision=revision,
        allow_patterns=["*.json", "*.txt", "model.safetensors", "README.md"]))
             for name, (model, revision) in MODELS.items()}
    tokenizer = AutoTokenizer.from_pretrained(paths["tts"], local_files_only=True)
    tts = VitsModel.from_pretrained(paths["tts"], torch_dtype=torch.float32,
                                    local_files_only=True).eval()
    sample_rate = tts.config.sampling_rate
    assert sample_rate == 16000
    report = {
        "models": {key: list(value) for key, value in MODELS.items()}, "texts": TEXTS,
        "source_sha256": file_hash(__file__), "tts": [], "asr": [],
        "environment": {"python": platform.python_version(), "platform": platform.platform(),
                        "torch": torch.__version__, "numpy": np.__version__,
                        "transformers": transformers.__version__,
                        "device": "cpu", "dtype": "float32", "threads": 4},
        "config": {"sample_rate": 16000, "max_new_tokens": 64, "do_sample": False,
                   "num_beams": 1, "return_timestamps": False,
                   "return_dict_in_generate": True, "vad": False,
                   "tts_seed_start": 2026, "noise_seed_start": 3030},
        "weights": {key: {"bytes": (path / "model.safetensors").stat().st_size,
                          "sha256": file_hash(path / "model.safetensors")}
                    for key, path in paths.items()},
        "tts_config": tts.config.to_dict(),
    }
    cases = []
    for i, text in enumerate(TEXTS):
        torch.manual_seed(2026 + i)
        inputs = tokenizer(text, return_tensors="pt")
        start = perf_counter()
        with torch.inference_mode():
            values = tts(**inputs).waveform[0].numpy()
        elapsed = perf_counter() - start
        gain = write_wav(folder / f"{i}-clean.wav", values, sample_rate)
        clean, _ = read_wav(folder / f"{i}-clean.wav")
        report["tts"].append({"id": i, "text": text, "seed": 2026 + i,
                              "seconds": elapsed, "samples": len(clean), "pcm_gain": gain,
                              "raw_sha256": sha256(values.tobytes()).hexdigest()})
        for name, snr in CONDITIONS:
            filename = f"{i}-{name}.wav"
            mix_gain = 1.0
            if snr is not None:
                noise = add_noise(clean, snr, 3030 + i)
                mix_gain = write_wav(folder / filename, clean + noise, sample_rate)
            cases.append({"id": f"{i}-{name}", "condition": name, "audio": filename,
                          "reference": text, "snr_db": snr, "mix_gain": mix_gain})
        print(f"TTS {i}：{len(clean) / sample_rate:.2f} 秒音频", flush=True)
    del tts
    write_wav(folder / "silence.wav", np.zeros(4 * sample_rate))
    cases.append({"id": "silence", "condition": "silence", "audio": "silence.wav",
                  "reference": "", "snr_db": None, "mix_gain": 1.0})
    processor = WhisperProcessor.from_pretrained(paths["asr"], local_files_only=True)
    asr = WhisperForConditionalGeneration.from_pretrained(
        paths["asr"], torch_dtype=torch.float32, attn_implementation="eager",
        local_files_only=True).eval()
    report["eos_token_id"] = asr.generation_config.eos_token_id
    report["asr_generation_config"] = asr.generation_config.to_dict()
    report["feature_config"] = processor.feature_extractor.to_dict()
    for case in cases:
        values, rate = read_wav(folder / case["audio"])
        start = perf_counter()
        inputs = processor(values, sampling_rate=rate, return_tensors="pt",
                           return_attention_mask=True)
        with torch.inference_mode():
            generated = asr.generate(**inputs, do_sample=False, num_beams=1,
                                     max_new_tokens=64, return_timestamps=False,
                                     return_dict_in_generate=True)
            # 普通返回值会移除 EOS；完整返回值保留开始/结束标记供核对。
            ids = generated.sequences[0]
        text = processor.decode(ids, skip_special_tokens=True)
        elapsed = perf_counter() - start
        row = {**case, "samples": len(values), "duration": len(values) / rate,
               "audio_sha256": file_hash(folder / case["audio"]),
               "feature_shape": list(inputs.input_features.shape),
               "token_ids": ids.tolist(), "text": text,
               "eos": bool(ids[-1] == report["eos_token_id"]),
               "seconds": elapsed, "rtf": elapsed / (len(values) / rate),
               "score": edit_counts(case["reference"], text)}
        report["asr"].append(row)
        (folder / "report.json").write_text(json.dumps(
            report, ensure_ascii=False, indent=2) + "\n")
        print(f"ASR {case['id']}：{text}", flush=True)
    validate(folder / "report.json")
    return report


def show(report):
    print("4 句合成语音；干净、20 dB、0 dB 各 4 段，另有 1 段静音。")
    for name, _ in CONDITIONS:
        rows = [r["score"] for r in report["asr"] if r["condition"] == name]
        errors = sum(r["errors"] for r in rows)
        total = sum(r["reference_words"] for r in rows)
        print(f"{name}：往返词错误 {errors}/{total}，WER={errors / total:.3f}")
    silence = report["asr"][-1]
    print(f"静音输出：{silence['text']!r}；插入词数={silence['score']['insertions']}")
    print("默认模式只核对保存记录；这不是人声基准，也没有测量主观音质。")


if __name__ == "__main__":
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--out", type=Path)
    parser.add_argument("--record", type=Path, default=DEFAULT)
    args = parser.parse_args()
    if args.run and args.out is None:
        parser.error("--run 需要 --out 指定空输出目录")
    show(run(args.out) if args.run else validate(args.record))