#[cfg(test)]
mod props;

mod dict;
mod pair;
mod repr;

mod size_of {
    use std::{mem, ops};
    type Run = ops::RangeInclusive<u16>;
    pub static U16: usize = mem::size_of::<u16>();
    pub static U32: usize = mem::size_of::<u32>();
    pub static U64: usize = mem::size_of::<u64>();
    pub static RUN: usize = mem::size_of::<Run>();
}

pub use self::dict::{PopCount, Rank, Select0, Select1};

/// Set is a set of `u32`.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Set {
    entries: Vec<Entry>,
}

/// Entry is a component of `Set`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Entry {
    slot: u16,
    bits: repr::Block,
}

impl Set {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn clear(&mut self) {
        *self = Self::default();
    }

    pub fn get(&self, x: u32) -> bool {
        let (slot, bit) = split(x);
        self.entries
            .binary_search_by_key(&slot, |e| e.slot)
            .map(|i| self.entries[i].bits.get(bit))
            .unwrap_or(false)
    }

    pub fn put(&mut self, x: u32, enabled: bool) -> bool {
        let (slot, bit) = split(x);
        if enabled {
            match self.entries.binary_search_by_key(&slot, |e| e.slot) {
                Ok(i) => self.entries[i].bits.put(bit, true),
                Err(i) => {
                    let bucket = {
                        let mut bits = repr::Block::default();
                        bits.put(bit, true);
                        Entry { slot, bits }
                    };
                    self.entries.insert(i, bucket);
                    false
                }
            }
        } else if let Ok(i) = self.entries.binary_search_by_key(&slot, |e| e.slot) {
            self.entries[i].bits.put(bit, false)
        } else {
            false
        }
    }

    pub fn insert(&mut self, x: u32) -> bool {
        self.put(x, true)
    }
    pub fn remove(&mut self, x: u32) -> bool {
        self.put(x, false)
    }

    pub fn shrink(&mut self) {
        self.entries.retain(|b| b.bits.count1() > 0);
        self.entries.shrink_to_fit();
    }
}

/// SetEntries is an iterator for bitwise ops.
pub struct SetEntries<'a>(Mapping<'a, fn(&'a Entry) -> CowEntry<'a>>);

type Mapping<'a, F> = ::std::iter::Map<::std::slice::Iter<'a, Entry>, F>;

type CowEntry<'a> = ::std::borrow::Cow<'a, Entry>;

pub struct And<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    pairs: pair::And<L, R>,
}
pub fn and<'a, L, R>(lhs: L, rhs: R) -> And<'a, L::IntoIter, R::IntoIter>
where
    L: IntoIterator<Item = CowEntry<'a>>,
    R: IntoIterator<Item = CowEntry<'a>>,
{
    And {
        pairs: pair::and(lhs, rhs),
    }
}

pub struct Or<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    pairs: pair::Or<L, R>,
}
pub fn or<'a, L, R>(lhs: L, rhs: R) -> Or<'a, L::IntoIter, R::IntoIter>
where
    L: IntoIterator<Item = CowEntry<'a>>,
    R: IntoIterator<Item = CowEntry<'a>>,
{
    Or {
        pairs: pair::or(lhs, rhs),
    }
}

pub struct AndNot<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    pairs: pair::AndNot<L, R>,
}
pub fn and_not<'a, L, R>(lhs: L, rhs: R) -> AndNot<'a, L::IntoIter, R::IntoIter>
where
    L: IntoIterator<Item = CowEntry<'a>>,
    R: IntoIterator<Item = CowEntry<'a>>,
{
    AndNot {
        pairs: pair::and_not(lhs, rhs),
    }
}

pub struct Xor<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    pairs: pair::Xor<L, R>,
}
pub fn xor<'a, L, R>(lhs: L, rhs: R) -> Xor<'a, L::IntoIter, R::IntoIter>
where
    L: IntoIterator<Item = CowEntry<'a>>,
    R: IntoIterator<Item = CowEntry<'a>>,
{
    Xor {
        pairs: pair::xor(lhs, rhs),
    }
}

fn split(u: u32) -> (u16, u16) {
    let q = u / (1 << 16);
    let r = u % (1 << 16);
    (q as u16, r as u16)
}

fn merge(n: u16, m: u16) -> u32 {
    (n as u32) * (1 << 16) + (m as u32)
}

impl Set {
    pub fn and<'a, T>(&'a self, that: T) -> And<'a, SetEntries<'a>, T::IntoIter>
    where
        T: IntoIterator<Item = CowEntry<'a>>,
    {
        and(self, that)
    }

    pub fn or<'a, T>(&'a self, that: T) -> Or<'a, SetEntries<'a>, T::IntoIter>
    where
        T: IntoIterator<Item = CowEntry<'a>>,
    {
        or(self, that)
    }

    pub fn and_not<'a, T>(&'a self, that: T) -> AndNot<'a, SetEntries<'a>, T::IntoIter>
    where
        T: IntoIterator<Item = CowEntry<'a>>,
    {
        and_not(self, that)
    }

    pub fn xor<'a, T>(&'a self, that: T) -> Xor<'a, SetEntries<'a>, T::IntoIter>
    where
        T: IntoIterator<Item = CowEntry<'a>>,
    {
        xor(self, that)
    }

    pub fn bits<'a>(&'a self) -> impl Iterator<Item = u32> + 'a {
        self.entries.iter().flat_map(|entry| {
            let slot = entry.slot;
            entry.bits.bits().map(move |half| merge(slot, half))
        })
    }
}

impl Entry {
    pub fn bits<'a>(&'a self) -> impl Iterator<Item = u32> + 'a {
        let slot = self.slot;
        self.bits.bits().map(move |half| merge(slot, half))
    }
}

impl ::std::cmp::PartialOrd for Entry {
    fn partial_cmp(&self, that: &Self) -> Option<::std::cmp::Ordering> {
        let s1 = self.slot;
        let s2 = that.slot;
        s1.partial_cmp(&s2)
    }
}
impl ::std::cmp::Ord for Entry {
    fn cmp(&self, that: &Self) -> ::std::cmp::Ordering {
        let s1 = self.slot;
        let s2 = that.slot;
        s1.cmp(&s2)
    }
}

impl<'a> Iterator for SetEntries<'a> {
    type Item = CowEntry<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next()
    }
}

impl<'a> IntoIterator for &'a Set {
    type Item = CowEntry<'a>;
    type IntoIter = SetEntries<'a>;
    fn into_iter(self) -> Self::IntoIter {
        SetEntries(self.entries.iter().map(|b| ::std::borrow::Cow::Borrowed(b)))
    }
}

impl<T: AsRef<[u32]>> From<T> for Set {
    fn from(v: T) -> Self {
        let mut bits = Self::new();
        for &bit in v.as_ref() {
            bits.insert(bit);
        }
        bits.shrink();
        bits
    }
}

impl<'a> ::std::iter::FromIterator<Entry> for Set {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = Entry>,
    {
        let entries = iter.into_iter().collect();
        Set { entries }
    }
}
impl<'a> ::std::iter::FromIterator<CowEntry<'a>> for Set {
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = CowEntry<'a>>,
    {
        let entries = iter.into_iter().map(|entry| entry.into_owned()).collect();
        Set { entries }
    }
}

macro_rules! impl_fns {
    ( $( $tyname:ident ),* ) => ($(
        impl<'a, L, R> $tyname<'a, L, R>
        where
            L: Iterator<Item = CowEntry<'a>>,
            R: Iterator<Item = CowEntry<'a>>,
        {
            pub fn and<T>(self, that: T) -> And<'a, Self, T::IntoIter>
            where
                T: IntoIterator<Item = CowEntry<'a>>,
            {
                and(self, that)
            }

            pub fn or<T>(self, that: T) -> Or<'a, Self, T::IntoIter>
            where
                T: IntoIterator<Item = CowEntry<'a>>,
            {
                or(self, that)
            }

            pub fn and_not<T>(self, that: T) -> AndNot<'a, Self, T::IntoIter>
            where
                T: IntoIterator<Item = CowEntry<'a>>,
            {
                and_not(self, that)
            }

            pub fn xor<T>(self, that: T) -> Xor<'a, Self, T::IntoIter>
            where
                T: IntoIterator<Item = CowEntry<'a>>,
            {
                xor(self, that)
            }
        }

        impl<'a, L, R> $tyname<'a, L, R>
        where
            L: Iterator<Item = CowEntry<'a>>,
            R: Iterator<Item = CowEntry<'a>>,
        {
            pub fn bits(self) -> impl Iterator<Item = u32> + 'a
            where
                L: 'a,
                R: 'a,
            {
                self.flat_map(move |entry| {
                    let slot = entry.slot;
                    entry
                        .into_owned()
                        .bits
                        .into_bits()
                        .map(move |half| merge(slot, half))
                })
            }
        }
    )*)
}
impl_fns!(And, Or, AndNot, Xor);

impl<'a, L, R> Iterator for And<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    type Item = CowEntry<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        self.pairs.next().and_then(|pair| match pair {
            (Some(mut lhs), Some(rhs)) => {
                lhs.to_mut().bits.and_assign(&rhs.bits);
                Some(lhs)
            }
            _ => None,
        })
    }
}

impl<'a, L, R> Iterator for Or<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    type Item = CowEntry<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        self.pairs.next().and_then(|pair| match pair {
            (Some(mut lhs), Some(rhs)) => {
                lhs.to_mut().bits.or_assign(&rhs.bits);
                Some(lhs)
            }
            (Some(lhs), None) => Some(lhs),
            (None, Some(rhs)) => Some(rhs),
            (None, None) => None,
        })
    }
}

impl<'a, L, R> Iterator for AndNot<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    type Item = CowEntry<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        self.pairs.next().and_then(|pair| match pair {
            (Some(mut lhs), Some(rhs)) => {
                lhs.to_mut().bits.andnot_assign(&rhs.bits);
                Some(lhs)
            }
            (Some(lhs), None) => Some(lhs),
            _ => None,
        })
    }
}

impl<'a, L, R> Iterator for Xor<'a, L, R>
where
    L: Iterator<Item = CowEntry<'a>>,
    R: Iterator<Item = CowEntry<'a>>,
{
    type Item = CowEntry<'a>;
    fn next(&mut self) -> Option<Self::Item> {
        self.pairs.next().and_then(|pair| match pair {
            (Some(mut lhs), Some(rhs)) => {
                lhs.to_mut().bits.xor_assign(&rhs.bits);
                Some(lhs)
            }
            (Some(lhs), None) => Some(lhs),
            (None, Some(rhs)) => Some(rhs),
            _ => None,
        })
    }
}

impl ::std::ops::Index<u32> for Set {
    type Output = bool;
    fn index(&self, i: u32) -> &Self::Output {
        static T: bool = true;
        static F: bool = false;
        if self.get(i) {
            &T
        } else {
            &F
        }
    }
}

impl PopCount<u64> for Set {
    const MAX_BOUND: u64 = 1 << 32;
    fn count1(&self) -> u64 {
        self.entries
            .iter()
            .map(|b| u64::from(b.bits.count1()))
            .sum()
    }
}

impl Rank<u32> for Set {
    fn rank1(&self, i: u32) -> u32 {
        let (hi, lo) = split(i);
        let mut rank = 0;
        for bucket in &self.entries {
            if bucket.slot > hi {
                break;
            } else if bucket.slot == hi {
                rank += u32::from(bucket.bits.rank1(lo));
                break;
            } else {
                rank += bucket.bits.count1();
            }
        }
        rank
    }
}

impl Select1<u32> for Set {
    fn select1(&self, mut c: u32) -> Option<u32> {
        if self.count1() <= u64::from(c) {
            return None;
        }
        for bucket in &self.entries {
            let w = bucket.bits.count1();
            if c >= w {
                c -= w;
            } else {
                let s = u32::from(bucket.bits.select1(c as u16).unwrap());
                let k = u32::from(bucket.slot) << 16;
                return Some(s + k);
            }
        }
        None
    }
}

impl Select0<u32> for Set {
    fn select0(&self, c: u32) -> Option<u32> {
        /// Find the smallest index i in range at which f(i) is true, assuming that
        /// f(i) == true implies f(i+1) == true.
        macro_rules! search {
            ($start: expr, $end: expr, $func: expr) => {{
                let mut i = $start;
                let mut j = $end;
                while i < j {
                    let h = i + (j - i) / 2;
                    if $func(h) {
                        j = h; // f(j) == true
                    } else {
                        i = h + 1; // f(i-1) == false
                    }
                }
                i
            }};
        }

        if self.count0() <= u64::from(c) {
            return None;
        }
        let fun = |i| self.rank0(i as u32) > c;
        let pos = search!(0u64, 1 << 32, fun);
        if pos < (1 << 32) {
            Some(pos as u32 - 1)
        } else {
            None
        }
    }
}

mod impl_io {
    use super::{size_of, Entry, PopCount, Set, repr::{self, Block, Compressed}};
    use {byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}, std::io};

    impl Set {
        // The cookie header spans either 32 bits or 64 bits.
        //
        // If the first 32 bits have the value `SERIAL_NO_RUN`,
        // then there is no `Run` slot in `Set`.
        // the next 32 bits are used to store the number of slots.
        // If the bitset is empty (i.e., it has no container),
        // then you should choose this cookie header.
        //
        // If the 16 least significant bits of the 32-bit cookie have the value `SERIAL_COOKIE`,
        // the 16 most significant bits of the 32-bit cookie are used to store
        // the number of slots minus 1.
        // That is, if you shift right by 16 the cookie and add 1, you get the number of slots.
        //
        // Then we store `RunIndex` following the initial 32 bits,
        // as a bitset to indicate whether each of the slots is a `Run` or not.
        //
        // The LSB of the first byte corresponds to the first stored slots and so forth.
        // Thus if follows that the least significant 16 bits of the first 32 bits
        // of a serialized bitsets should either have the value `SERIAL_NO_RUN`
        // or the value SERIAL_COOKIE. In other cases, we should abort the decoding.
        //
        // After scanning the cookie header, we know how many containers are present in the bitset.

        pub fn write_to<W: io::Write>(&self, mut write: W) -> io::Result<()> {
            let w = &mut write;

            let mut compressed = Vec::with_capacity(self.entries.len());
            let mut bitset = vec![0u8; (self.entries.len() + 7) / 8];
            let mut hasrun = false;

            for (i, b) in self.entries.iter().enumerate() {
                let c = b.bits.compressed();
                if c.is_run() {
                    hasrun = true;
                    bitset[i / 8] |= 1 << (i % 8);
                }
                compressed.push((b.slot, c));
            }

            let runidx = RunIndex { hasrun, bitset };

            let (size_of_cookie, size_of_runidx) = if runidx.is_empty() {
                (2 * size_of::U32, 0)
            } else {
                (2 * size_of::U16, runidx.bytes().len())
            };

            let size_of_header = 2 * size_of::U16 * self.entries.len();
            let sum_size_of = size_of_cookie + size_of_runidx + size_of_header;

            // serial cookie
            if runidx.is_empty() {
                w.write_u32::<LittleEndian>(SERIAL_NO_RUN)?;
                w.write_u32::<LittleEndian>(self.entries.len() as u32)?;
            } else {
                w.write_u16::<LittleEndian>(SERIAL_COOKIE)?;
                w.write_u16::<LittleEndian>((self.entries.len() - 1) as u16)?;
                w.write_all(runidx.bytes())?;
            };

            // header
            for &(slot, ref bits) in &compressed {
                w.write_u16::<LittleEndian>(slot)?;
                w.write_u16::<LittleEndian>((bits.count1() - 1) as u16)?;
            }

            if runidx.is_empty() || self.entries.len() >= NO_OFFSET_THRESHOLD as usize {
                // offset
                let mut offset = sum_size_of + 2 * size_of::U16 * self.entries.len();
                for &(_, ref bits) in &compressed {
                    w.write_u32::<LittleEndian>(offset as u32)?;
                    let pop = bits.count1();
                    match *bits {
                        Compressed::Arr(_) => {
                            assert!(pop as usize > repr::SEQ_MAX_LEN);
                            offset += (1 << 16) / 8;
                        }
                        Compressed::Seq(_) => {
                            assert!(pop as usize <= repr::SEQ_MAX_LEN);
                            offset += size_of::U16 * pop as usize;
                        }
                        Compressed::Run(ref run) => {
                            offset += size_of::U16;
                            offset += 2 * size_of::U16 * run.ranges.len();
                        }
                    }
                }
            }

            for &(_, ref bits) in &compressed {
                match *bits {
                    Compressed::Seq(ref seq) => seq.write_to(w)?,
                    Compressed::Arr(ref arr) => arr.write_to(w)?,
                    Compressed::Run(ref run) => run.write_to(w)?,
                }
            }

            Ok(())
        }

        pub fn read_from<R: io::Read>(mut read: R) -> io::Result<Self> {
            let r = &mut read;

            match r.read_u32::<LittleEndian>()? {
                cookie if cookie == SERIAL_NO_RUN => {
                    let slot_len = r.read_u32::<LittleEndian>()? as usize;
                    let header = read_header(r, slot_len)?;

                    discard_offset(r, slot_len)?;

                    let mut entries = Vec::with_capacity(slot_len);

                    for (slot, pop) in header {
                        let pop = pop as usize;
                        let bits = if pop > repr::SEQ_MAX_LEN {
                            let arr = repr::Arr::read_from(r)?;
                            Block::from(arr)
                        } else {
                            let seq = repr::Seq::read_from(r, pop)?;
                            Block::from(&seq)
                        };
                        entries.push(Entry { slot, bits });
                    }
                    Ok(Set { entries })
                }

                cookie if cookie & 0x_0000_FFFF == u32::from(SERIAL_COOKIE) => {
                    let slot_len = (cookie.wrapping_shr(16) + 1) as usize;
                    let bytes_len = (slot_len + 7) / 8;

                    let hasrun = true;
                    let bitset = {
                        let mut buf = vec![0; bytes_len];
                        r.read_exact(&mut buf)?;
                        buf
                    };
                    let runidx = RunIndex { hasrun, bitset };
                    let header = read_header(r, slot_len)?;

                    if runidx.is_empty() || slot_len >= NO_OFFSET_THRESHOLD as usize {
                        discard_offset(r, slot_len)?;
                    }

                    let mut entries = Vec::with_capacity(slot_len);

                    for (i, (slot, pop)) in header.into_iter().enumerate() {
                        let pop = pop as usize;

                        let bits = if runidx.bitset[i / 8] & (1 << (i % 8)) > 0 {
                            let run = repr::Run::read_from(r)?;
                            Block::from(&run)
                        } else if pop > repr::SEQ_MAX_LEN {
                            let arr = repr::Arr::read_from(r)?;
                            Block::from(arr)
                        } else {
                            let seq = repr::Seq::read_from(r, pop)?;
                            Block::from(&seq)
                        };
                        entries.push(Entry { slot, bits });
                    }
                    Ok(Set { entries })
                }

                x => Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("unexpected cookie value: {}", x),
                )),
            }
        }
    }

    const SERIAL_COOKIE: u16 = 12_347; // `Arr`, `Seq` and `Run`
    const SERIAL_NO_RUN: u32 = 12_346; // `Arr` and `Seq`
    const NO_OFFSET_THRESHOLD: u8 = 4;

    struct RunIndex {
        hasrun: bool,
        bitset: Vec<u8>,
    }

    impl RunIndex {
        fn is_empty(&self) -> bool {
            !self.hasrun
        }
        fn bytes(&self) -> &[u8] {
            &self.bitset
        }
    }

    fn read_header<R: io::Read>(r: &mut R, size: usize) -> io::Result<Vec<(u16, u32)>> {
        let mut vec = Vec::with_capacity(size);
        for _ in 0..size {
            let key = r.read_u16::<LittleEndian>()?;
            let pop = r.read_u16::<LittleEndian>()?;
            vec.push((key, u32::from(pop) + 1));
        }
        // vec is sorted?
        Ok(vec)
    }

    fn discard_offset<R: io::Read>(r: &mut R, size: usize) -> io::Result<()> {
        for _ in 0..size {
            let _offset = r.read_u32::<LittleEndian>()?;
        }
        Ok(())
    }
}