1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

use crate::{
    core::prelude::*,
    errors::prelude::*,
    extensions::prelude::*,
    numeric::prelude::*,
};

/// ArrayTrait - Binary Array bits operations
pub trait ArrayBinaryBits where Self: Sized + Clone {

    /// Unpacks elements of a uint8 array into a binary-valued output array
    ///
    /// # Arguments
    ///
    /// * `axis` - the dimension over which bit-unpacking is done. if none, array is flattened
    /// * `count` - the number of elements to unpack along axis. if negative, array is trimmed
    /// * `bit_order` - the order of the returned bits. defaults to `Big`
    ///
    /// # Examples
    ///
    /// ```
    /// use arr_rs::prelude::*;
    ///
    /// let expected: Result<Array<u8>, _> = array!([[0, 0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 0, 1, 1, 1]]);
    /// let array: Result<Array<u8>, _> = array!([[2], [7], [23]]);
    /// assert_eq!(expected, array.unpack_bits(Some(1), None, None));
    /// ```
    fn unpack_bits(&self, axis: Option<isize>, count: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError>;

    /// Packs the elements of a binary-valued array into bits in a uint8 array
    ///
    /// # Arguments
    ///
    /// * `axis` - the dimension over which bit-packing is done. if none, array is flattened
    /// * `bit_order` - the order of the returned bits. defaults to `Big`
    ///
    /// # Examples
    ///
    /// ```
    /// use arr_rs::prelude::*;
    ///
    /// let expected: Result<Array<u8>, _> = array!([[2], [7], [23]]);
    /// let array: Result<Array<u8>, _> = array!([[0, 0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 0, 1, 1, 1]]);
    /// assert_eq!(expected, array.pack_bits(Some(1), None));
    /// ```
    fn pack_bits(&self, axis: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError>;
}

impl ArrayBinaryBits for Array<u8> {

    fn unpack_bits(&self, axis: Option<isize>, count: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError> {
        if self.is_empty()? { return Array::empty() }
        let bit_order = bit_order.unwrap_or(BitOrder::Big);
        match axis {
            None => {
                let result = self.ravel()?
                    .into_iter()
                    .flat_map(|a| {
                        let mut elems = (0..8).rev().map(move |idx| (a >> idx) & 1).collect::<Vec<u8>>();
                        if bit_order == BitOrder::Little { elems.reverse_ext() } else { elems }
                    })
                    .collect::<Array<u8>>();
                let count = count.unwrap_or(self.len()? as isize * 8);
                if count >= 0 { result.slice(0 .. count as usize) }
                else { result.slice(0 .. self.len()? - count as usize) }
            },
            Some(axis) => {
                let axis = self.normalize_axis(axis);
                self.apply_along_axis(axis, |arr| arr.unpack_bits(None, count, Some(bit_order)))
            }
        }
    }

    fn pack_bits(&self, axis: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError> {
        if self.is_empty()? { return Array::empty() }
        let bit_order = bit_order.unwrap_or(BitOrder::Big);
        match axis {
            None => {
                let mut elements = self.get_elements()?;
                if elements.len() % 8 != 0 { elements.extend_from_slice(&vec![0; 8 - elements.len() % 8]) }

                let parts = elements.len() / 8;
                let result = (0 .. parts).map(|p| {
                    let subarray = elements[p * 8 .. (p + 1) * 8].iter()
                        .map(|i| if i > &0 { "1" } else { "0" })
                        .collect::<Vec<&str>>()
                        .join("");
                    let subarray = if bit_order == BitOrder::Little { subarray.chars().rev().collect() } else { subarray };
                    u8::from_str_radix(&subarray, 2).unwrap()
                }).collect();
                Ok(result)
            },
            Some(axis) => {
                let axis = self.normalize_axis(axis);
                self.apply_along_axis(axis, |arr| arr.pack_bits(None, Some(bit_order)))
            },
        }
    }
}

impl ArrayBinaryBits for Result<Array<u8>, ArrayError> {

    fn unpack_bits(&self, axis: Option<isize>, count: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError> {
        self.clone()?.unpack_bits(axis, count, bit_order)
    }

    fn pack_bits(&self, axis: Option<isize>, bit_order: Option<BitOrder>) -> Result<Array<u8>, ArrayError> {
        self.clone()?.pack_bits(axis, bit_order)
    }
}