#!/usr/bin/env python3
"""Prepare real Vosk feeds, quantize acoustic ONNX, and render A/B WAVs.

The frontend mirrors VoskSpeechEngine.kt: number normalization, WordPiece
positions, dictionary lookup and all five phone feature channels.
"""

from __future__ import annotations

import argparse
import bisect
import json
import math
import re
import sys
import time
import wave
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator

import numpy as np

SAMPLE_RATE = 22_050
SPEAKER_ID = 4
BERT_DIMENSIONS = 768
MAX_WORD_PIECE_CHARS = 100
CACHE_SIZE = 4_096
MULTISTREAM_PUNCTUATION = set("!(),-.:;?")
BERT_EXCLUDED_PUNCTUATION = {"-", ",", ".", "?", "!", ";", ":", '"'}
PHONEMES = [
    "_", "^", "$", " ", "!", "'", "(", ")", ",", "-", ".", "...", ":", ";", "?",
    "a0", "a1", "b", "bj", "c", "ch", "d", "dj", "e0", "e1", "f", "fj", "g", "gj",
    "h", "hj", "i0", "i1", "j", "k", "kj", "l", "lj", "m", "mj", "n", "nj", "o0",
    "o1", "p", "pj", "r", "rj", "s", "sch", "sh", "sj", "t", "tj", "u0", "u1", "v",
    "vj", "y0", "y1", "z", "zh", "zj",
]
PHONEME_IDS = {phone: index for index, phone in enumerate(PHONEMES)}

DIGITS = ["ноль", "один", "два", "три", "четыре", "пять", "шесть", "семь", "восемь", "девять"]
TEENS = [
    "десять", "одиннадцать", "двенадцать", "тринадцать", "четырнадцать",
    "пятнадцать", "шестнадцать", "семнадцать", "восемнадцать", "девятнадцать",
]
TENS = ["", "десять", "двадцать", "тридцать", "сорок", "пятьдесят", "шестьдесят", "семьдесят", "восемьдесят", "девяносто"]
HUNDREDS = ["", "сто", "двести", "триста", "четыреста", "пятьсот", "шестьсот", "семьсот", "восемьсот", "девятьсот"]
SCALES = [
    None,
    ("тысяча", "тысячи", "тысяч"),
    ("миллион", "миллиона", "миллионов"),
    ("миллиард", "миллиарда", "миллиардов"),
    ("триллион", "триллиона", "триллионов"),
    ("квадриллион", "квадриллиона", "квадриллионов"),
]
NUMBER_PATTERN = re.compile(r"(?<![^\W_])([+-]?)(\d{1,18})(?:[,.](\d+))?(?![^\W_])", re.UNICODE)
YEAR_IN_PREPOSITIONAL_PATTERN = re.compile(r"(?<![^\W_])(\d{4})\s+(году|г\.)", re.IGNORECASE | re.UNICODE)

SOFT_LETTERS = set("яёюиье")
SYLLABLE_STARTS = set("#ъьаёуюэеиы-")
SOFT_HARD_CONSONANTS = {
    "б": "b", "в": "v", "г": "g", "д": "d", "з": "z", "к": "k", "л": "l",
    "м": "m", "н": "n", "п": "p", "р": "r", "с": "s", "т": "t", "ф": "f", "х": "h",
}
OTHER_CONSONANTS = {"ж": "zh", "ц": "c", "ч": "ch", "ш": "sh", "щ": "sch", "й": "j"}
VOWELS = {"а": "a", "я": "a", "у": "u", "ю": "u", "о": "o", "ё": "o", "э": "e", "е": "e", "и": "i", "ы": "y"}


def scale_form(value: int, forms: tuple[str, str, str]) -> str:
    if 11 <= value % 100 <= 14:
        return forms[2]
    return forms[0] if value % 10 == 1 else forms[1] if value % 10 in (2, 3, 4) else forms[2]


def under_thousand_to_words(value: int, feminine: bool) -> str:
    words: list[str] = []
    if value // 100:
        words.append(HUNDREDS[value // 100])
    tail = value % 100
    if 10 <= tail <= 19:
        words.append(TEENS[tail - 10])
    else:
        if tail // 10:
            words.append(TENS[tail // 10])
        units = tail % 10
        if units:
            words.append("одна" if feminine and units == 1 else "две" if feminine and units == 2 else DIGITS[units])
    return " ".join(words)


def integer_to_words(value: int) -> str:
    if value == 0:
        return DIGITS[0]
    groups: list[int] = []
    while value:
        groups.append(value % 1_000)
        value //= 1_000
    result: list[str] = []
    for group_index in range(len(groups) - 1, -1, -1):
        group = groups[group_index]
        if not group:
            continue
        words = under_thousand_to_words(group, group_index == 1)
        forms = SCALES[group_index] if group_index < len(SCALES) else None
        result.append(words if forms is None else f"{words} {scale_form(group, forms)}")
    return " ".join(result)


def normalize_numbers(text: str) -> str:
    prepared = re.sub(r"(?<=\d)[\u00a0\u202f](?=\d{3}(?:\D|$))", "", text)
    prepared = re.sub(r"№\s*(?=\d)", "номер ", prepared)

    def replace(match: re.Match[str]) -> str:
        sign = "минус " if match.group(1) == "-" else "плюс " if match.group(1) == "+" else ""
        integer_digits = match.group(2)
        if len(integer_digits) > 1 and integer_digits.startswith("0"):
            integer = " ".join(DIGITS[int(char)] for char in integer_digits)
        else:
            integer = integer_to_words(int(integer_digits))
        fraction = match.group(3)
        return f"{sign}{integer}" if not fraction else f"{sign}{integer} запятая {' '.join(DIGITS[int(char)] for char in fraction)}"

    prepared = YEAR_IN_PREPOSITIONAL_PATTERN.sub(
        lambda match: f"{year_in_prepositional_words(int(match.group(1)))} {match.group(2)}",
        prepared,
    )
    return NUMBER_PATTERN.sub(replace, prepared)


ORDINAL_UNITS_PREPOSITIONAL = ("", "первом", "втором", "третьем", "четвёртом", "пятом", "шестом", "седьмом", "восьмом", "девятом")
ORDINAL_TEENS_PREPOSITIONAL = ("десятом", "одиннадцатом", "двенадцатом", "тринадцатом", "четырнадцатом", "пятнадцатом", "шестнадцатом", "семнадцатом", "восемнадцатом", "девятнадцатом")
ORDINAL_TENS_PREPOSITIONAL = ("", "десятом", "двадцатом", "тридцатом", "сороковом", "пятидесятом", "шестидесятом", "семидесятом", "восьмидесятом", "девяностом")
ORDINAL_HUNDREDS_PREPOSITIONAL = ("", "сотом", "двухсотом", "трёхсотом", "четырёхсотом", "пятисотом", "шестисотом", "семисотом", "восьмисотом", "девятисотом")
EXACT_THOUSANDTH_PREPOSITIONAL = ("", "тысячном", "двухтысячном", "трёхтысячном", "четырёхтысячном", "пятитысячном", "шеститысячном", "семитысячном", "восьмитысячном", "девятитысячном")


def ordinal_under_thousand_prepositional(value: int) -> str:
    last_two = value % 100
    units = value % 10
    if units and not 11 <= last_two <= 19:
        return f"{integer_to_words(value - units) if value > units else ''} {ORDINAL_UNITS_PREPOSITIONAL[units]}".strip()
    if 11 <= last_two <= 19:
        return f"{integer_to_words(value - last_two) if value > last_two else ''} {ORDINAL_TEENS_PREPOSITIONAL[last_two - 10]}".strip()
    if last_two >= 20:
        return f"{integer_to_words(value - last_two) if value > last_two else ''} {ORDINAL_TENS_PREPOSITIONAL[last_two // 10]}".strip()
    return ORDINAL_HUNDREDS_PREPOSITIONAL[value // 100]


def year_in_prepositional_words(value: int) -> str:
    if not 1000 <= value <= 9999:
        raise ValueError(value)
    if value % 1000 == 0:
        return EXACT_THOUSANDTH_PREPOSITIONAL[value // 1000]
    thousands = value // 1000
    thousands_words = under_thousand_to_words(thousands, True)
    thousands_prefix = "" if thousands_words == "одна" else thousands_words + " "
    return f"{thousands_prefix}{scale_form(thousands, SCALES[1])} {ordinal_under_thousand_prepositional(value % 1000)}"


def normalize_text(text: str) -> str:
    stressed = re.sub(r"([А-Яа-яЁё])\u0301", lambda match: "+" + match.group(1), text)
    translated = normalize_numbers(stressed).lower().replace("—", "-").replace("–", "-").replace("…", "...")
    translated = translated.translate(str.maketrans({"«": '"', "»": '"', "“": '"', "”": '"'}))
    return re.sub(r"\s+", " ", translated).strip()


class WordPieceTokenizer:
    def __init__(self, path: Path):
        self.tokens = path.read_text(encoding="utf-8").splitlines()
        self.vocabulary = {token: index for index, token in enumerate(self.tokens)}
        for special in ("[CLS]", "[SEP]", "[UNK]"):
            if special not in self.vocabulary:
                raise ValueError(f"Missing {special} in {path}")

    def encode(self, text: str) -> tuple[np.ndarray, list[int]]:
        ids = [self.vocabulary["[CLS]"]]
        positions = [0]
        for token in self._basic_tokens(text):
            punctuation = token in BERT_EXCLUDED_PUNCTUATION
            pieces = [self.vocabulary.get(token, self.vocabulary["[UNK]"])] if punctuation else self._word_pieces(token)
            if not punctuation:
                positions.append(len(ids))
            ids.extend(pieces)
        ids.append(self.vocabulary["[SEP]"])
        positions.append(len(ids) - 1)
        return np.asarray(ids, dtype=np.int64), positions

    @staticmethod
    def _basic_tokens(text: str) -> list[str]:
        result: list[str] = []
        word: list[str] = []
        for char in text:
            if char.isspace():
                if word:
                    result.append("".join(word)); word.clear()
            elif char.isalnum():
                word.append(char.lower())
            elif char == "+":
                continue
            else:
                if word:
                    result.append("".join(word)); word.clear()
                result.append(char)
        if word:
            result.append("".join(word))
        return result

    def _word_pieces(self, token: str) -> list[int]:
        if len(token) > MAX_WORD_PIECE_CHARS:
            return [self.vocabulary["[UNK]"]]
        result: list[int] = []
        start = 0
        while start < len(token):
            found: int | None = None
            found_end = start
            for end in range(len(token), start, -1):
                piece = token[start:end] if start == 0 else f"##{token[start:end]}"
                if piece in self.vocabulary:
                    found = self.vocabulary[piece]; found_end = end; break
            if found is None:
                return [self.vocabulary["[UNK]"]]
            result.append(found); start = found_end
        return result


class PronunciationDictionary:
    def __init__(self, dictionary_path: Path, index_path: Path):
        self.source = dictionary_path.open("rb")
        entries = []
        for line in index_path.read_text(encoding="utf-8").splitlines():
            word, offset = line.rsplit("\t", 1)
            entries.append((word, int(offset)))
        self.words = [entry[0] for entry in entries]
        self.offsets = [entry[1] for entry in entries]
        self.length = dictionary_path.stat().st_size
        self.cache: OrderedDict[str, str | None] = OrderedDict()

    def close(self) -> None:
        self.source.close()

    def find(self, word: str) -> str | None:
        if word in self.cache:
            value = self.cache.pop(word); self.cache[word] = value; return value
        index = max(0, bisect.bisect_right(self.words, word) - 1)
        end = self.offsets[index + 1] if index + 1 < len(self.offsets) else self.length
        self.source.seek(self.offsets[index])
        value: str | None = None
        while self.source.tell() < end:
            raw = self.source.readline()
            if not raw:
                break
            entry_word, _, pronunciation = raw.decode("utf-8").rstrip("\r\n").partition("\t")
            if entry_word < word:
                continue
            if entry_word == word:
                value = pronunciation.strip() or None
            break
        self.cache[word] = value
        if len(self.cache) > CACHE_SIZE:
            self.cache.popitem(last=False)
        return value


def russian_g2p(word: str) -> str:
    marked: list[tuple[str, int]] = []
    stress = 0
    for char in f"#{word}#":
        if char == "+":
            stress = 1
        else:
            marked.append((char, stress)); stress = 0
    phones: list[str] = []
    for index, (char, accent) in enumerate(marked):
        previous = marked[index - 1][0] if index else "#"
        following = marked[index + 1][0] if index + 1 < len(marked) else None
        if char in SOFT_HARD_CONSONANTS:
            phones.append(SOFT_HARD_CONSONANTS[char] + ("j" if following in SOFT_LETTERS else ""))
        elif char in OTHER_CONSONANTS:
            phones.append(OTHER_CONSONANTS[char])
        elif char in VOWELS:
            if previous in SYLLABLE_STARTS and char in set("яюеё"):
                phones.append("j")
            phones.append(f"{VOWELS[char]}{accent}")
    return " ".join(phones)


@dataclass
class RawPhone:
    phone: str
    punctuation: list[str]
    in_quote: int
    bert_word_index: int


def build_multistream_phones(text: str, dictionary: PronunciationDictionary) -> tuple[np.ndarray, list[int]]:
    raw = [RawPhone("^", [], 0, 0)]
    word: list[str] = []
    pending: list[str] = []
    in_quote = 0
    bert_word_index = 1

    def flush_word() -> None:
        nonlocal bert_word_index
        if not word:
            return
        value = "".join(word)
        pronunciation = dictionary.find(value) or russian_g2p(value)
        raw.extend(RawPhone(phone, [], in_quote, bert_word_index) for phone in pronunciation.split() if phone)
        word.clear(); bert_word_index += 1

    def append_space() -> None:
        raw.append(RawPhone(" ", list(pending), in_quote, bert_word_index)); pending.clear()

    index = 0
    while index < len(text):
        if text.startswith("...", index):
            flush_word(); pending.append("..."); index += 3; continue
        char = text[index]
        if char in {'"', "«", "»", "“", "”"}:
            flush_word(); in_quote = 1 - in_quote
        elif char.isspace():
            flush_word(); append_space()
        elif char == "-" and index > 0 and index + 1 < len(text) and text[index - 1].isalnum() and text[index + 1].isalnum():
            flush_word()
        elif char in MULTISTREAM_PUNCTUATION:
            flush_word(); pending.append(char)
        elif char.isalnum() or char == "+":
            word.append(char)
        index += 1
    flush_word(); append_space(); raw.append(RawPhone("$", [], 0, bert_word_index))

    last_punctuation = " "
    last_sentence_punctuation = " "
    reversed_features: list[list[int]] = []
    reversed_positions: list[int] = []
    for phone in reversed(raw):
        for candidate in ("...", ".", "!", "?", "-"):
            if candidate in phone.punctuation:
                last_sentence_punctuation = candidate; break
        if phone.punctuation:
            last_punctuation = phone.punctuation[0]
        current = phone.punctuation[0] if phone.punctuation else "_"
        values = [PHONEME_IDS[phone.phone], PHONEME_IDS[current], phone.in_quote, PHONEME_IDS[last_punctuation], PHONEME_IDS[last_sentence_punctuation]]
        reversed_features.append(values); reversed_positions.append(phone.bert_word_index)
    features = np.asarray(list(reversed(reversed_features)), dtype=np.int64).T[np.newaxis, :, :]
    return features, list(reversed(reversed_positions))


class Frontend:
    def __init__(self, assets: Path, bert_model: Path | None = None):
        import onnxruntime as ort
        self.tokenizer = WordPieceTokenizer(assets / "vocab.txt")
        self.dictionary = PronunciationDictionary(assets / "dict.tsv", assets / "dictionary.index")
        self.bert = ort.InferenceSession(str(bert_model or assets / "bert.int8.onnx"), providers=["CPUExecutionProvider"])

    def close(self) -> None:
        self.dictionary.close()

    def prepare(self, text: str, articulation: float = 1.0) -> dict[str, np.ndarray]:
        normalized = normalize_text(text)
        ids, embedding_positions = self.tokenizer.encode(normalized)
        shape = (1, ids.size)
        token_embeddings = self.bert.run(None, {
            "input_ids": ids.reshape(shape),
            "attention_mask": np.ones(shape, dtype=np.int64),
            "token_type_ids": np.zeros(shape, dtype=np.int64),
        })[0]
        selected = token_embeddings[np.asarray(embedding_positions, dtype=np.int64)]
        features, phone_positions = build_multistream_phones(normalized, self.dictionary)
        phone_embeddings = selected[np.clip(np.asarray(phone_positions), 0, len(selected) - 1)]
        bert = phone_embeddings.T[np.newaxis, :, :].astype(np.float32, copy=False)
        time_steps = features.shape[2]
        if bert.shape != (1, BERT_DIMENSIONS, time_steps):
            raise ValueError(f"Unexpected BERT shape {bert.shape} for {text!r}")
        return {
            "input": features,
            "input_lengths": np.asarray([time_steps], dtype=np.int64),
            "scales": np.asarray([0.8, articulation, 0.8], dtype=np.float32),
            "sid": np.asarray([SPEAKER_ID], dtype=np.int64),
            "bert": bert,
        }


def load_corpus(path: Path) -> list[str]:
    lines = [line.strip() for line in path.read_text(encoding="utf-8-sig").splitlines()]
    result = [line for line in lines if line and not line.startswith("#")]
    if not result:
        raise ValueError(f"No phrases in {path}")
    return result


def command_prepare(args: argparse.Namespace) -> None:
    args.output.mkdir(parents=True, exist_ok=False)
    frontend = Frontend(args.assets, args.bert_model)
    manifest = []
    try:
        for index, text in enumerate(load_corpus(args.corpus)):
            feeds = frontend.prepare(text, args.articulation)
            name = f"sample-{index:04d}.npz"
            np.savez_compressed(args.output / name, **feeds)
            manifest.append({"file": name, "text": text, "time_steps": int(feeds["input_lengths"][0])})
            print(f"prepared {index + 1}: phones={manifest[-1]['time_steps']} text={text}")
    finally:
        frontend.close()
    (args.output / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"samples={len(manifest)} output={args.output}")


class NpzCalibrationReader:
    def __init__(self, directory: Path):
        self.files = sorted(directory.glob("sample-*.npz"))
        self.iterator: Iterator[Path] | None = None
        if not self.files:
            raise ValueError(f"No sample-*.npz in {directory}")

    def get_next(self) -> dict[str, np.ndarray] | None:
        if self.iterator is None:
            self.iterator = iter(self.files)
        try:
            path = next(self.iterator)
        except StopIteration:
            return None
        with np.load(path) as data:
            return {name: data[name] for name in data.files}

    def rewind(self) -> None:
        self.iterator = iter(self.files)


def command_quantize(args: argparse.Namespace) -> None:
    from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
    if args.output.exists():
        raise FileExistsError(f"Refusing to overwrite {args.output}")
    started = time.time()
    quantize_static(
        str(args.model), str(args.output), NpzCalibrationReader(args.calibration),
        quant_format=QuantFormat.QDQ, op_types_to_quantize=["Conv"], per_channel=True,
        activation_type=QuantType.QInt8, weight_type=QuantType.QInt8,
    )
    source_size = args.model.stat().st_size
    target_size = args.output.stat().st_size
    print(json.dumps({
        "source_bytes": source_size, "target_bytes": target_size,
        "saving_bytes": source_size - target_size,
        "saving_percent": round((source_size - target_size) * 100 / source_size, 4),
        "seconds": round(time.time() - started, 3),
    }, indent=2))


def write_wav(path: Path, samples: np.ndarray) -> None:
    pcm = (np.clip(samples.reshape(-1), -1.0, 1.0) * 32767.0).astype("<i2").tobytes()
    with wave.open(str(path), "wb") as output:
        output.setnchannels(1); output.setsampwidth(2); output.setframerate(SAMPLE_RATE); output.writeframes(pcm)


def command_compare(args: argparse.Namespace) -> None:
    import onnxruntime as ort
    args.output.mkdir(parents=True, exist_ok=False)
    options = ort.SessionOptions(); options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    baseline = ort.InferenceSession(str(args.baseline), sess_options=options, providers=["CPUExecutionProvider"])
    candidate = ort.InferenceSession(str(args.candidate), sess_options=options, providers=["CPUExecutionProvider"])
    manifest = json.loads((args.calibration / "manifest.json").read_text(encoding="utf-8"))
    report = []
    for item in manifest[: args.limit]:
        with np.load(args.calibration / item["file"]) as data:
            feeds = {name: data[name] for name in data.files}
        feeds["scales"] = np.asarray([0.0, float(feeds["scales"][1]), 0.0], dtype=np.float32)
        base_wav, base_length = baseline.run(None, feeds)
        test_wav, test_length = candidate.run(None, feeds)
        common = min(base_wav.size, test_wav.size)
        difference = base_wav.reshape(-1)[:common] - test_wav.reshape(-1)[:common]
        entry = {
            "text": item["text"], "baseline_samples": int(base_wav.size), "candidate_samples": int(test_wav.size),
            "baseline_length": np.asarray(base_length).tolist(), "candidate_length": np.asarray(test_length).tolist(),
            "mae": float(np.mean(np.abs(difference))), "rmse": float(math.sqrt(np.mean(difference * difference))),
        }
        report.append(entry)
        number = len(report)
        write_wav(args.output / f"{number:02d}-baseline.wav", base_wav)
        write_wav(args.output / f"{number:02d}-candidate.wav", test_wav)
        print(f"compared {number}: mae={entry['mae']:.6f} rmse={entry['rmse']:.6f}")
    (args.output / "report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def command_compare_feeds(args: argparse.Namespace) -> None:
    import onnxruntime as ort
    args.output.mkdir(parents=True, exist_ok=args.resume)
    options = ort.SessionOptions(); options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    acoustic = ort.InferenceSession(str(args.model), sess_options=options, providers=["CPUExecutionProvider"])
    baseline_manifest = json.loads((args.baseline_calibration / "manifest.json").read_text(encoding="utf-8"))
    candidate_manifest = json.loads((args.candidate_calibration / "manifest.json").read_text(encoding="utf-8"))
    if [item["text"] for item in baseline_manifest] != [item["text"] for item in candidate_manifest]:
        raise ValueError("Baseline and candidate manifests contain different texts")
    report = []
    for baseline_item, candidate_item in zip(baseline_manifest[: args.limit or None], candidate_manifest[: args.limit or None]):
        with np.load(args.baseline_calibration / baseline_item["file"]) as data:
            baseline_feeds = {name: data[name] for name in data.files}
        with np.load(args.candidate_calibration / candidate_item["file"]) as data:
            candidate_feeds = {name: data[name] for name in data.files}
        baseline_feeds["scales"] = np.asarray([0.0, float(baseline_feeds["scales"][1]), 0.0], dtype=np.float32)
        candidate_feeds["scales"] = np.asarray([0.0, float(candidate_feeds["scales"][1]), 0.0], dtype=np.float32)
        baseline_wav, baseline_length = acoustic.run(None, baseline_feeds)
        candidate_wav, candidate_length = acoustic.run(None, candidate_feeds)
        common = min(baseline_wav.size, candidate_wav.size)
        difference = baseline_wav.reshape(-1)[:common] - candidate_wav.reshape(-1)[:common]
        entry = {
            "text": baseline_item["text"],
            "baseline_samples": int(baseline_wav.size), "candidate_samples": int(candidate_wav.size),
            "baseline_length": np.asarray(baseline_length).tolist(),
            "candidate_length": np.asarray(candidate_length).tolist(),
            "mae": float(np.mean(np.abs(difference))),
            "rmse": float(math.sqrt(np.mean(difference * difference))),
        }
        report.append(entry)
        number = len(report)
        baseline_path = args.output / f"{number:03d}-baseline.wav"
        candidate_path = args.output / f"{number:03d}-student-int8.wav"
        if not baseline_path.exists():
            write_wav(baseline_path, baseline_wav)
        if not candidate_path.exists():
            write_wav(candidate_path, candidate_wav)
        temporary_report = args.output / "report.json.tmp"
        temporary_report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        temporary_report.replace(args.output / "report.json")
        print(f"rendered {number}: mae={entry['mae']:.6f} rmse={entry['rmse']:.6f}", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest="command", required=True)
    prepare = commands.add_parser("prepare")
    prepare.add_argument("--assets", type=Path, required=True)
    prepare.add_argument("--corpus", type=Path, required=True)
    prepare.add_argument("--output", type=Path, required=True)
    prepare.add_argument("--articulation", type=float, default=1.0)
    prepare.add_argument("--bert-model", type=Path)
    prepare.set_defaults(handler=command_prepare)
    quantize = commands.add_parser("quantize")
    quantize.add_argument("--model", type=Path, required=True)
    quantize.add_argument("--calibration", type=Path, required=True)
    quantize.add_argument("--output", type=Path, required=True)
    quantize.set_defaults(handler=command_quantize)
    compare = commands.add_parser("compare")
    compare.add_argument("--baseline", type=Path, required=True)
    compare.add_argument("--candidate", type=Path, required=True)
    compare.add_argument("--calibration", type=Path, required=True)
    compare.add_argument("--output", type=Path, required=True)
    compare.add_argument("--limit", type=int, default=3)
    compare.set_defaults(handler=command_compare)
    compare_feeds = commands.add_parser("compare-feeds")
    compare_feeds.add_argument("--model", type=Path, required=True)
    compare_feeds.add_argument("--baseline-calibration", type=Path, required=True)
    compare_feeds.add_argument("--candidate-calibration", type=Path, required=True)
    compare_feeds.add_argument("--output", type=Path, required=True)
    compare_feeds.add_argument("--limit", type=int, default=0)
    compare_feeds.add_argument("--resume", action="store_true")
    compare_feeds.set_defaults(handler=command_compare_feeds)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    args.handler(args)
    return 0


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