use crate::common::*;

/// Converts an iterator over contiguous u32 chunks into an iterator over usize

/// e.g. input [0b111000, 0b11] gives output [3, 4, 5, 32, 33]

/// observe that the bits per chunk are ordered from least to most significant bits, yielding smaller to larger usizes.

/// works by draining the inner u32 chunk iterator one u32 at a time, then draining that chunk until its 0.

struct BitChunkIter<I: Iterator<Item = u32>> {

    chunk_iter: I,

    next_bit_index: usize,

    cached: u32,

}

impl<I: Iterator<Item = u32>> BitChunkIter<I> {

    fn new(chunk_iter: I) -> Self {

        // first chunk is always a dummy zero, as if chunk_iter yielded Some(0).

        // Consequences:

        // 1. our next_bit_index is always off by 32 (we correct for it in Self::next) (no additional overhead)

        // 2. we cache u32 and not Option<u32>, because chunk_iter.next() is only called in Self::next.

        Self { chunk_iter, next_bit_index: 0, cached: 0 }

    }

}

impl<I: Iterator<Item = u32>> Iterator for BitChunkIter<I> {

    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {

        let mut chunk = self.cached;

        // loop until either:

        // 1. there are no more Items to return, or

        // 2. chunk encodes 1+ Items, one of which we will return.

        while chunk == 0 {

            // chunk is still empty! get the next one...

            chunk = self.chunk_iter.next()?;

            // ... and jump self.next_bit_index to the next multiple of 32.

            self.next_bit_index = (self.next_bit_index + 32) & !(32 - 1);

        }

        // assert(chunk > 0);

        // Shift the contents of chunk until the least significant bit is 1.

        // ... being sure to increment next_bit_index accordingly.

        #[inline(always)]

        fn skip_n_zeroes(chunk: &mut u32, n: usize, next_bit_index: &mut usize) {

            if *chunk & ((1 << n) - 1) == 0 {

                // n least significant bits are zero. skip n bits.

                *next_bit_index += n;

                *chunk >>= n;

            }

        }

        skip_n_zeroes(&mut chunk, 16, &mut self.next_bit_index);

        skip_n_zeroes(&mut chunk, 08, &mut self.next_bit_index);

        skip_n_zeroes(&mut chunk, 04, &mut self.next_bit_index);

        skip_n_zeroes(&mut chunk, 02, &mut self.next_bit_index);

        skip_n_zeroes(&mut chunk, 01, &mut self.next_bit_index);

        // least significant bit of chunk is 1.

        // assert(chunk & 1 == 1)

        // prepare our state for the next time Self::next is called.

        // Overwrite self.cached such that its shifted state is retained,

        // and jump over the bit whose index we are about to return.

        self.next_bit_index += 1;

        self.cached = chunk >> 1;

        // returned index is 32 smaller than self.next_bit_index because we use an

        // off-by-32 encoding to avoid having to cache an Option<u32>.

        Some(self.next_bit_index - 1 - 32)

    }

}

/*  --properties-->

     ___ ___ ___ ___

    |___|___|___|___|

  | |___|___|___|___|

  | |___|___|___|___|

  | |___|___|___|___|

  |

  V

 entity chunks (groups of 32)

*/

#[derive(Debug, Copy, Clone, Eq, PartialEq)]

struct Pair {

    entity: u32,

    property: u32,

}

impl From<[u32; 2]> for Pair {

    fn from([entity, property]: [u32; 2]) -> Self {

        Pair { entity, property }

    }

}

struct BitMatrix {

    bounds: Pair,

    buffer: *mut u32,

}

impl Drop for BitMatrix {

    fn drop(&mut self) {

        let total_chunks = Self::row_chunks(self.bounds.property) as usize

            * Self::column_chunks(self.bounds.entity) as usize;

        let layout = Self::layout_for(total_chunks);

        unsafe {

            // ?

            std::alloc::dealloc(self.buffer as *mut u8, layout);

        }

    }

}

impl Debug for BitMatrix {

    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {

        let row_chunks = Self::row_chunks(self.bounds.property) as usize;

        let column_chunks = Self::column_chunks(self.bounds.entity) as usize;

        for property in 0..row_chunks {

            for entity_chunk in 0..column_chunks {

                write!(f, "|")?;

                let mut chunk = unsafe { *self.buffer.add(row_chunks * entity_chunk + property) };

                let end =

                    if entity_chunk + 1 == column_chunks { self.bounds.entity % 32 } else { 32 };

                for _ in 0..end {

                    let c = match chunk & 1 {

                        0 => '0',

                        _ => '1',

                    };

                    write!(f, "{}", c)?;

                    chunk >>= 1;

                }

            }

            write!(f, "|\n")?;

        }

        Ok(())

    }

}

impl BitMatrix {

    #[inline]

    fn ceiling_to_mul_32(value: u32) -> u32 {

        (value + 31) & !31

    }

    #[inline]

    fn row_of(entity: u32) -> u32 {

        entity / 32

    }

    #[inline]

    fn row_chunks(property_bound: u32) -> u32 {

        property_bound

    }

    #[inline]

    fn column_chunks(entity_bound: u32) -> u32 {

        Self::ceiling_to_mul_32(entity_bound) / 32

    }

    #[inline]

    fn offsets_unchecked(&self, at: Pair) -> [usize; 2] {

        let o_in = at.entity as usize % 32;

        let row = Self::row_of(at.entity);

        let row_chunks = self.bounds.property;

        let o_of = row as usize * row_chunks as usize + at.property as usize;

        [o_of, o_in]

    }

    // returns a u32 which has bits 000...000111...111

    // for the last JAGGED chunk given the column size

    // if the last chunk is not jagged (when entity_bound % 32 == 0)

    // None is returned,

    // otherwise Some(x) is returned such that x & chunk would mask out

    // the bits NOT in 0..entity_bound

    fn last_row_chunk_mask(entity_bound: u32) -> Option<u32> {

        let zero_prefix_len = entity_bound % 32;

        if zero_prefix_len == 0 {

            None

        } else {

            Some(!0u32 >> zero_prefix_len)

        }

    }

    fn assert_within_bounds(&self, at: Pair) {

        assert!(at.entity < self.bounds.entity);

        assert!(at.property < self.bounds.property);

    }

    /////////

    fn reshape(&mut self, dims: [usize; 2]) {

        todo!()

    }

    fn new(bounds: Pair) -> Self {

        let total_chunks = Self::row_chunks(bounds.property) as usize

            * Self::column_chunks(bounds.entity) as usize;

        let layout = Self::layout_for(total_chunks);

        let buffer;

        unsafe {

            buffer = std::alloc::alloc(layout) as *mut u32;

            buffer.write_bytes(0u8, total_chunks);

        };

        Self { buffer, bounds }

    }

    fn set(&mut self, at: Pair) {

        self.assert_within_bounds(at);

        let [o_of, o_in] = self.offsets_unchecked(at);

        unsafe { *self.buffer.add(o_of) |= 1 << o_in };

    }

    fn unset(&mut self, at: Pair) {

        self.assert_within_bounds(at);

        let [o_of, o_in] = self.offsets_unchecked(at);

        unsafe { *self.buffer.add(o_of) &= !(1 << o_in) };

    }

    fn test(&self, at: Pair) -> bool {

        self.assert_within_bounds(at);

        let [o_of, o_in] = self.offsets_unchecked(at);

        unsafe { (*self.buffer.add(o_of) & 1 << o_in) != 0 }

    }

    fn iter<'a, 'b>(

        &'a self,

        buf: &'b mut Vec<u32>,

        mut fold_fn: impl FnMut(&'b [u32]) -> u32,

    ) -> impl Iterator<Item = usize> + 'b {

        let buf_start = buf.len();

        let row_chunks = Self::row_chunks(self.bounds.property) as usize;

        let column_chunks = Self::column_chunks(self.bounds.entity);

        let mut ptr = self.buffer;

        for _row in 0..column_chunks {

            let slice;

            unsafe {

                slice = std::slice::from_raw_parts(ptr, row_chunks);

                ptr = ptr.add(row_chunks);

            }

            buf.push(fold_fn(slice));

        }

        if let Some(mask) = Self::last_row_chunk_mask(self.bounds.entity) {

            *buf.iter_mut().last().unwrap() &= mask;

        }

        BitChunkIter::new(buf.drain(buf_start..))

    }

    fn layout_for(total_chunks: usize) -> std::alloc::Layout {

        unsafe {

            // this layout is ALWAYS valid:

            // 1. size is always nonzero

            // 2. size is always a multiple of 4 and 4-aligned

            std::alloc::Layout::from_size_align_unchecked(4 * total_chunks.max(1), 4)

        }

    }

}

#[test]

fn matrix_test() {

    let mut m = BitMatrix::new(Pair { entity: 50, property: 3 });

    m.set([2, 0].into());

    m.set([40, 1].into());

    m.set([40, 2].into());

    m.set([40, 0].into());

    println!("{:?}", &m);

    let mut buf = vec![];

    for index in m.iter(&mut buf, move |slice| slice[0] ^ slice[1] ^ slice[2]) {

        println!("index {}", index);

    }

}

// trait RwMatrixBits {

//     fn set(&mut self, at: [usize;2]);

//     fn unset(&mut self, at: [usize;2]);

//     fn set_entire_row(&mut self, row: usize);

//     fn unset_entire_row(&mut self, row: usize);

// }

// struct MatrixRefW<'a> {

//     _inner: usize,

// }

// impl<'a> MatrixRefW<'a> {

// }

// EXPERIMENTAL:

// mod predicate; // TODO later

// mod ecs;

mod bits;