/// Returns the Walsh spectrum of an n-input boolean function packed as a truth-table code.
pub fn walsh_spectrum(code: u128, n: usize) -> Vec<i64> {
    let size = 1usize << n;
    let mut t: Vec<i64> = (0..size)
        .map(|x| if (code >> x) & 1 == 1 { -1 } else { 1 })
        .collect();
    let mut len = 1;
    while len < size {
        let mut i = 0;
        while i < size {
            for j in i..i + len {
                let a = t[j];
                let b = t[j + len];
                t[j] = a + b;
                t[j + len] = a - b;
            }
            i += len << 1;
        }
        len <<= 1;
    }
    t
}

/// Returns how far the packed function sits from every affine function, zero when it is one.
///
/// ```
/// assert_eq!(mrlynum::boolean::nonlinearity(0b0110, 2), 0);
/// assert_eq!(mrlynum::boolean::nonlinearity(0b1000, 2), 1);
/// ```
pub fn nonlinearity(code: u128, n: usize) -> i64 {
    if n == 0 {
        return 0;
    }
    let max_w = walsh_spectrum(code, n)
        .into_iter()
        .map(|w| w.abs())
        .max()
        .unwrap_or(0);
    (1i64 << (n - 1)) - max_w / 2
}

/// Reports whether the packed function outputs one on exactly half of its inputs.
pub fn is_balanced(code: u128, n: usize) -> bool {
    let size = 1u32 << n;
    (code & ((1u128 << size) - 1)).count_ones() == size / 2
}

/// Returns the mean chance that flipping one input bit flips the output, 0.5 at full avalanche.
pub fn sac(code: u128, n: usize) -> f64 {
    if n == 0 {
        return 0.0;
    }
    let size = 1usize << n;
    let mut total = 0.0;
    for bit in 0..n {
        let mut flips = 0usize;
        for x in 0..size {
            let f = (code >> x) & 1;
            let g = (code >> (x ^ (1 << bit))) & 1;
            if f != g {
                flips += 1;
            }
        }
        total += flips as f64 / size as f64;
    }
    total / n as f64
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn linear_function_has_zero_nonlinearity() {
        let n = 3;
        let size = 1usize << n;
        let mut code: u128 = 0;
        for x in 0..size {
            if x & 1 == 1 {
                code |= 1 << x;
            }
        }
        assert_eq!(nonlinearity(code, n), 0);
    }
    #[test]
    fn n3_max_nonlinearity_is_two() {
        let n = 3;
        let max = (0..(1u128 << (1 << n)))
            .map(|c| nonlinearity(c, n))
            .max()
            .unwrap();
        assert_eq!(max, 2);
    }
    #[test]
    fn n4_reaches_bent_level_six() {
        let n = 4;
        let max = (0..(1u128 << (1 << n)))
            .map(|c| nonlinearity(c, n))
            .max()
            .unwrap();
        assert_eq!(max, 6);
    }
    #[test]
    fn balance_detects_half_ones() {
        assert!(is_balanced(0b0011, 2));
        assert!(!is_balanced(0b0111, 2));
    }
    #[test]
    fn carpet_corners_expand_on_four_walsh_characters() {
        let code = 7;
        for x in 0..2u32 {
            for y in 0..2u32 {
                let (sx, sy) = (1 - 2 * (x as i64), 1 - 2 * (y as i64));
                assert_eq!(
                    ((code >> (2 * x + y)) & 1) as i64,
                    (3 + sx + sy - sx * sy) / 4
                );
            }
        }
        let spectrum = walsh_spectrum(code, 2);
        assert_eq!(spectrum, vec![-2, -2, -2, 2]);
        let quarters: Vec<i64> = spectrum
            .iter()
            .enumerate()
            .map(|(s, &w)| ((if s == 0 { 4 } else { 0 }) - w) / 2)
            .collect();
        assert_eq!(quarters, vec![3, 1, 1, -1]);
    }
    #[test]
    fn sac_of_constant_is_zero() {
        assert_eq!(sac(0, 3), 0.0);
        let n = 3;
        let size = 1usize << n;
        let mut code: u128 = 0;
        for x in 0..size {
            if x & 1 == 1 {
                code |= 1 << x;
            }
        }
        assert!((sac(code, n) - 1.0 / 3.0).abs() < 1e-12);
    }
}