viterbi 0.2.0

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-05
//! Generic constraint-length `K ≤ 9` end-to-end tests (RK1–RK4) plus the
//! `M → m = k-1` tail-length transition guards.
//!
//! The crate default `M = 6` is the CCSDS K=7 tail length. A K=9 code has memory
//! `m = k-1 = 8`, so its zero-tail flush must be 8 stages (not 6) or the encoder
//! never returns to state 0 and traceback-from-0 breaks. These tests exercise the
//! K=9 (`S = 256`) generic path end-to-end and assert the tail length is derived
//! from the code's own `k` in the encoder, the decoder, and the puncturer.

use viterbi::puncture::{DePuncturer, Puncturer};
use viterbi::{
    CodeParams, CodeProfile, CodedBlock, ConfigError, HardHamming, ParamError, PunctureMatrix,
    PuncturedRate, ViterbiDecoder, ViterbiEncoder,
};

/// Unpack the first `cb.nbits` MSB-first coded bits into 0/1 hard samples (mirrors
/// the crate's internal packing so the on-air stream lines up with the encoder).
fn hard_samples(cb: &CodedBlock) -> Vec<u8> {
    (0..cb.nbits)
        .map(|i| (cb.bytes[i / 8] >> (7 - (i % 8))) & 1)
        .collect()
}

#[test]
fn k9_profile_resolves_to_valid_non_catastrophic_params() {
    let p = CodeProfile::K9R1_2.params();
    assert_eq!(p.k(), 9);
    assert_eq!(p.n(), 2);
    assert_eq!(p.generators(), &[0o561, 0o753]);
    // `params()` is infallible; re-validate via the checking constructor (RK3):
    // the `[0o561, 0o753]` K=9 code is non-catastrophic and `k ≤ 9`.
    assert!(CodeParams::new(9, vec![0o561, 0o753], vec![false, false]).is_ok());
}

#[test]
fn k9_decoder_builds_at_s256() {
    // S = 256 = 2^(9-1): the const-generic survivor buffer scales to K=9 (RK1).
    assert!(ViterbiDecoder::<256, HardHamming>::new(CodeProfile::K9R1_2.params(), 1_000).is_ok());
}

#[test]
fn k9_with_s64_is_state_count_mismatch() {
    // S=64 but K=9 needs 2^8 = 256 → StateCountMismatch, no panic (RK3).
    assert!(matches!(
        ViterbiDecoder::<64, HardHamming>::new(CodeProfile::K9R1_2.params(), 1_000),
        Err(ConfigError::Param(ParamError::StateCountMismatch {
            states: 64,
            k: 9
        }))
    ));
}

#[test]
fn k9_clean_round_trip_recovers_payload() {
    let enc = ViterbiEncoder::new(CodeProfile::K9R1_2.params()).unwrap();
    let mut dec =
        ViterbiDecoder::<256, HardHamming>::new(CodeProfile::K9R1_2.params(), 4096).unwrap();
    let raw = b"K=9 generic path works!";
    let coded = enc.encode(raw).unwrap();
    let out = dec.decode_block(&coded).unwrap();
    assert_eq!(&out.bytes, raw);
    assert_eq!(out.nbits, raw.len() * 8);
}

#[test]
fn k9_encode_terminates_in_state_zero() {
    // The `M → m = k-1` transition guard. A lingering `M = 6` tail on a K=9 code
    // (m = 8) leaves the encoder two stages short of state 0. Two independent
    // witnesses that the zero-tail is exactly `m = k-1 = 8` stages:
    //  (a) coded length = (nbits + 8) * 2 — an 8-stage zero flush clears the K=9
    //      shift register back to state 0;
    //  (b) after a CLEAN decode the terminal state-0 metric is 0 — the zero-cost
    //      correct path truly terminates at state 0.
    // Both fail if `M = 6` lingered in the encoder or decoder.
    let enc = ViterbiEncoder::new(CodeProfile::K9R1_2.params()).unwrap();
    let mut dec =
        ViterbiDecoder::<256, HardHamming>::new(CodeProfile::K9R1_2.params(), 4096).unwrap();
    let raw = b"tail";
    let nbits = raw.len() * 8;
    let coded = enc.encode(raw).unwrap();
    assert_eq!(
        coded.nbits,
        (nbits + 8) * 2,
        "K=9 zero-tail must be m = k-1 = 8 stages"
    );
    let out = dec.decode_block(&coded).unwrap();
    assert_eq!(&out.bytes, raw);
    assert_eq!(
        dec.last_terminal_metric(),
        0,
        "a clean K=9 decode must terminate in state 0 (metric 0)"
    );
}

#[test]
fn k9_recovers_from_a_correctable_error() {
    let enc = ViterbiEncoder::new(CodeProfile::K9R1_2.params()).unwrap();
    let mut dec =
        ViterbiDecoder::<256, HardHamming>::new(CodeProfile::K9R1_2.params(), 4096).unwrap();
    let raw = b"correct me please";
    let mut coded = enc.encode(raw).unwrap();
    // Flip a single coded bit — well within the K=9 (d_free = 12, t = 5) capability.
    coded.bytes[0] ^= 0b0100_0000;
    let out = dec.decode_block(&coded).unwrap();
    assert_eq!(&out.bytes, raw);
}

#[test]
fn k9_with_puncturing_round_trips_and_tail_is_m_times_n() {
    // The one test that catches an `M`-stray in `puncture.rs`'s `m·n` verbatim tail
    // for K ≠ 7: a K=9 code (m = 8, n = 2) with a rate-2/3 puncture matrix.
    let params = CodeProfile::K9R1_2.params();
    let enc = ViterbiEncoder::new(params.clone()).unwrap();
    let matrix = PunctureMatrix::ccsds_rate(2, PuncturedRate::R2_3).unwrap();
    let punc = Puncturer::new(matrix.clone(), &params).unwrap();
    let dp = DePuncturer::<HardHamming>::new(matrix.clone(), &params).unwrap();

    let raw = b"pk9"; // 3 bytes = 24 info bits, period-aligned (24 % 2 == 0)
    let nbits = raw.len() * 8;

    let coded = enc.encode(raw).unwrap();
    assert_eq!(coded.nbits, (nbits + 8) * 2, "encode tail must use m = 8");

    let punctured = punc.puncture(&coded, nbits).unwrap();

    // On-air length = kept body bits + `m·n = 8·2 = 16` verbatim tail samples.
    // A lingering `M = 6` would place 12 tail samples, shifting the body/tail split.
    let m = 8usize;
    let n = 2usize;
    let kept_body = (nbits / matrix.period()) * matrix.kept_per_period(); // 12 * 3 = 36
    assert_eq!(
        punctured.nbits,
        kept_body + m * n,
        "punctured tail must be m·n = 16 verbatim samples (K=9)"
    );
    assert_eq!(
        dp.expected_punctured_len(nbits),
        Some(kept_body + m * n),
        "expected_punctured_len must reflect m = 8, not M = 6"
    );

    // Full round-trip: de-puncture (reinsert erasures) then decode on the K=9 core.
    let on_air = hard_samples(&punctured);
    let full = dp.depuncture(&on_air, nbits).unwrap();
    let mut dec = ViterbiDecoder::<256, HardHamming>::new(params, 4096).unwrap();
    let out = dec.decode(&full, nbits).unwrap();
    assert_eq!(&out.bytes, raw);
}