// Varshamov-Tenengolts codes
// Source: HeytingLean.CodingTheory.VT.Basic
// Theorem preserved: vtCode_mem_iff

use crate::hamming_basic::Bit;

pub fn vt_moment(x: &[Bit]) -> usize {
    let mut moment = 0usize;
    let mut one_count = 0usize;
    for &b in x {
        if b == 1 {
            one_count += 1;
            moment += one_count;
        }
    }
    moment
}

pub fn vt_code_member(x: &[Bit], n: usize, a: usize, m: usize) -> bool {
    x.len() == n && n + 1 <= m && vt_moment(x) % m == a % m
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_vt_moment_empty() {
        assert_eq!(vt_moment(&[]), 0);
    }

    #[test]
    fn test_vt_moment_all_ones() {
        // [1,1,1] -> moments: 1, 2, 3 -> sum = 6
        assert_eq!(vt_moment(&[1, 1, 1]), 6);
    }

    #[test]
    fn test_vt_code_membership() {
        // VT(3, 0, 4): length 3, moment ≡ 0 mod 4, 4 ≥ 4
        // [0,0,0] has moment 0 ≡ 0 mod 4
        assert!(vt_code_member(&[0, 0, 0], 3, 0, 4));
    }
}