use bitvec::prelude::*;
use labrador_ldpc::decoder::DecodeFrom;

use crate::{
    buffer::{Buffer, BufferOverflow},
    error::EncodeError,
};

pub(crate) fn interleave<const N: usize>(
    data: &[u8],
    out: &mut Buffer<N>,
) -> Result<(), EncodeError> {
    let bv = data.view_bits::<Msb0>();

    let initial_length = out.len();
    for _ in 0..data.len() {
        out.try_push(0).map_err(|_| EncodeError::CatsOverflow)?;
    }
    let bv_out = out[initial_length..].view_bits_mut::<Msb0>();

    let mut out_i = 0;
    for i in 0..32 {
        for j in (0..bv.len()).step_by(32) {
            if i + j >= bv.len() {
                continue;
            }

            bv_out.set(out_i, bv[i + j]);

            out_i += 1;
        }
    }

    Ok(())
}

pub(crate) fn uninterleave<const N: usize>(
    data: &[u8],
    out: &mut Buffer<N>,
) -> Result<(), BufferOverflow> {
    let bv = data.view_bits::<Msb0>();

    for _ in 0..data.len() {
        out.try_push(0)?;
    }
    let bv_out = out.view_bits_mut::<Msb0>();

    let mut out_i = 0;
    for i in 0..32 {
        for j in (0..bv.len()).step_by(32) {
            if i + j >= bv.len() {
                continue;
            }

            bv_out.set(i + j, bv[out_i]);

            out_i += 1;
        }
    }

    Ok(())
}

pub(crate) fn uninterleave_soft<const N: usize, T: DecodeFrom>(
    data: &[T],
    out: &mut Buffer<N, T>,
) -> Result<(), BufferOverflow> {
    for _ in 0..data.len() {
        out.try_push(T::zero())?;
    }

    let mut out_i = 0;
    for i in 0..32 {
        for j in (0..data.len()).step_by(32) {
            if i + j >= data.len() {
                continue;
            }

            out[i + j] = data[out_i];

            out_i += 1;
        }
    }

    Ok(())
}

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

    #[test]
    fn interleaver_works() {
        let mut data = [0x84, 0x73, 0x12, 0xA3, 0xFF, 0x00, 0xC2, 0x1B, 0x77];
        let orig = Buffer::new_full(&mut data);

        let mut interleaved = [0; 10];
        let mut interleaved = Buffer::new(&mut interleaved, 0);
        interleaved.push(b'H');

        interleave(&orig, &mut interleaved).unwrap();

        let expected = [b'H', 0xCD, 0xB5, 0xDB, 0x2A, 0x0A, 0x52, 0x0C, 0x89, 0x4F];
        assert_eq!(expected, interleaved[..]);

        let mut uninterleaved = [0; 10];
        let mut uninterleaved = Buffer::new(&mut uninterleaved, 0);
        uninterleave(&interleaved[1..], &mut uninterleaved).unwrap();

        assert_eq!(*orig, *uninterleaved);
    }

    #[test]
    fn hard_interleave_soft_uninterleave() {
        let mut data = [0x84, 0x73, 0x12, 0xA3, 0xFF, 0x00, 0xC2, 0x1B, 0x77];
        let orig = Buffer::new_full(&mut data);

        let mut interleaved = [0; 10];
        let mut interleaved = Buffer::new(&mut interleaved, 0);
        interleaved.push(b'H');

        interleave(&orig, &mut interleaved).unwrap();

        let mut soft_interleaved = [0.0; 10 * 8];
        for (i, b) in interleaved.iter().enumerate() {
            for j in 0..8 {
                soft_interleaved[8 * i + j] = f32::from_hard_bit(b & (1 << (7 - j)) > 0);
            }
        }

        let mut uninterleaved = [0.0; 10 * 8];
        let mut uninterleaved = Buffer::new(&mut uninterleaved, 0);
        uninterleave_soft(&soft_interleaved[8..], &mut uninterleaved).unwrap();

        assert_eq!(orig.len() * 8, uninterleaved.len());
        for (i, b) in orig.iter().enumerate() {
            for j in 0..8 {
                assert_eq!(uninterleaved[8 * i + j].hard_bit(), *b & (1 << (7 - j)) > 0);
            }
        }
    }
}