"""Отдельный wall-clock замер Витерби на мягких значениях одного сигнала."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from pathlib import Path
import sys
import time

import numpy as np

from protocol import punctured_fec
from protocol.video003_modem import PHY_HEADER_SYMBOLS, PREAMBLE_SYMBOLS, receive_frame
from tools.video004_benchmark import (
    PACKETS, RATE, RESULT_ROOT, SYMBOL_RATE, TRANSIENT_SYMBOLS,
    load_fixture,
)


OUTPUT = RESULT_ROOT / "viterbi_benchmark.json"


def run(repetitions: int = 3) -> dict:
    if repetitions < 1:
        raise ValueError("число повторов должно быть положительным")
    from protocol.video004_viterbi import decode_soft as jit_decode
    received, manifest = load_fixture()
    cursor = 0
    soft_words = []
    for _ in range(PACKETS):
        result = receive_frame(received[cursor:], "QPSK", RATE, SYMBOL_RATE,
                               fec_decoder=jit_decode)
        if not result.crc_pass or result.coded_soft is None or result.packet is None:
            raise AssertionError("контрольный поток не принят")
        soft_words.append(result.coded_soft)
        body_count = math.ceil(punctured_fec.coded_bit_count(len(result.packet) * 8, RATE) / 2)
        cursor += result.start_symbol + PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS + body_count + TRANSIENT_SYMBOLS
    original = [punctured_fec.decode_soft(word, 1024 * 8, RATE) for word in soft_words]
    optimized = [jit_decode(word, 1024 * 8, RATE) for word in soft_words]
    if not all(np.array_equal(a, b) for a, b in zip(original, optimized)):
        raise AssertionError("Витерби изменил биты")
    times = {}
    for name, decoder in (("before", punctured_fec.decode_soft), ("after", jit_decode)):
        samples = []
        for _ in range(repetitions):
            started = time.perf_counter()
            outputs = [decoder(word, 1024 * 8, RATE) for word in soft_words]
            samples.append(time.perf_counter() - started)
            if not all(np.array_equal(a, b) for a, b in zip(original, outputs)):
                raise AssertionError("повтор Витерби изменил биты")
        samples.sort()
        times[name] = {"wall_times_s": samples, "median_s": samples[len(samples) // 2],
                       "realtime_factor": samples[len(samples) // 2] / manifest["signal_duration_s"]}
    digest = hashlib.sha256(b"".join(word.tobytes() for word in soft_words)).hexdigest()
    return {"signal_sha256": manifest["signal_sha256"], "soft_input_sha256": digest,
            "packet_count": PACKETS, "signal_duration_s": manifest["signal_duration_s"],
            "before": times["before"], "after": times["after"],
            "speedup": times["before"]["median_s"] / times["after"]["median_s"],
            "bits_equal": True, "rf_tx": 0}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--numba-packages", type=Path)
    parser.add_argument("--repetitions", type=int, default=3)
    args = parser.parse_args()
    if args.numba_packages is not None:
        if not args.numba_packages.is_dir():
            parser.error("каталог Numba отсутствует")
        sys.path.insert(0, str(args.numba_packages.resolve()))
    result = run(args.repetitions)
    OUTPUT.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"Viterbi: before={result['before']['median_s']:.3f}s "
          f"after={result['after']['median_s']:.3f}s speedup={result['speedup']:.1f}x")


if __name__ == "__main__":
    main()