#![cfg_attr(feature="clippy", feature(plugin))]
#![cfg_attr(feature="clippy", plugin(clippy))]

extern crate rand;
extern crate siphasher;

use rand::Rand;
use std::borrow::Borrow;
use std::cmp::max;
use std::hash::{Hash, Hasher};

use siphasher::sip::SipHasher13;
type FastHasher = SipHasher13;

use std::marker::PhantomData;
use std::mem;

macro_rules! cms_define {
    ($CountMinSketch:ident, $Counter:ty) => (

pub struct $CountMinSketch<K> {
    counters: Vec<Vec<$Counter>>,
    offsets: Vec<usize>,
    hashers: [FastHasher; 2],
    mask: usize,
    k_num: usize,
    reset_idx: usize,
    phantom_k: PhantomData<K>,
}

impl<K> $CountMinSketch<K>
    where K: Hash
{
    pub fn new(capacity: usize, probability: f64, tolerance: f64) -> Result<Self, &'static str> {
        let width = Self::optimal_width(capacity, tolerance);
        let k_num = Self::optimal_k_num(probability);
        let counters: Vec<Vec<$Counter>> = vec![vec![0; width]; k_num];
        let offsets = vec![0; k_num];
        let hashers = [Self::sip_new(), Self::sip_new()];
        let cms = $CountMinSketch {
            counters: counters,
            offsets: offsets,
            hashers: hashers,
            mask: Self::mask(width),
            k_num: k_num,
            reset_idx: 0,
            phantom_k: PhantomData,
        };
        Ok(cms)
    }

    pub fn add<Q: ?Sized>(&mut self, key: &Q, value: $Counter) where Q: Hash, K: Borrow<Q> {
        let mut hashes = [0u64, 0u64];
        let lowest = (0..self.k_num)
            .map(|k_i| {
                let offset = self.offset(&mut hashes, key, k_i);
                self.offsets[k_i] = offset;
                self.counters[k_i][offset]
            })
            .min()
            .unwrap();
        for k_i in 0..self.k_num {
            let offset = self.offsets[k_i];
            if self.counters[k_i][offset] == lowest {
                self.counters[k_i][offset] = self.counters[k_i][offset].saturating_add(value);
            }
        }
    }

    pub fn increment<Q: ?Sized>(&mut self, key: &Q) where Q: Hash, K: Borrow<Q> {
        self.add(key, 1)
    }

    pub fn estimate<Q: ?Sized>(&self, key: &Q) -> $Counter where Q: Hash, K: Borrow<Q> {
        let mut hashes = [0u64, 0u64];
        (0..self.k_num)
            .map(|k_i| {
                let offset = self.offset(&mut hashes, key, k_i);
                self.counters[k_i][offset]
            })
            .min()
            .unwrap()
    }

    pub fn estimate_memory(capacity: usize,
                           probability: f64,
                           tolerance: f64)
                           -> Result<usize, &'static str> {
        let width = Self::optimal_width(capacity, tolerance);
        let k_num = Self::optimal_k_num(probability);
        Ok(width * mem::size_of::<$Counter>() * k_num)
    }

    pub fn clear(&mut self) {
        for k_i in 0..self.k_num {
            for counter in &mut self.counters[k_i] {
                *counter = 0
            }
        }
        self.reset_idx = 0;
        self.hashers = [Self::sip_new(), Self::sip_new()];
    }

    pub fn reset(&mut self) {
        for k_i in 0..self.k_num {
            for counter in &mut self.counters[k_i] {
                *counter /= 2;
            }
        }
        self.reset_idx = 0;
    }

    pub fn reset_next(&mut self) -> Option<usize> {
        let idx = self.reset_idx;
        for k_i in 0..self.k_num {
            self.counters[k_i][idx] /= 2
        }
        let next = idx.wrapping_add(1) & self.mask;
        self.reset_idx = next;
        if next != 0 { Some(next) } else { None }
    }

    fn optimal_width(capacity: usize, tolerance: f64) -> usize {
        let e = tolerance / (capacity as f64);
        let width = (2.0 / e).round() as usize;
        max(2, width).checked_next_power_of_two().expect("Width would be way too large")
    }

    fn mask(width: usize) -> usize {
        assert!(width > 1);
        assert!(width & (width - 1) == 0);
        width - 1
    }

    fn optimal_k_num(probability: f64) -> usize {
        max(1, ((1.0 - probability).ln() / 0.5f64.ln()) as usize)
    }

    fn sip_new() -> FastHasher {
        let mut rng = rand::thread_rng();
        FastHasher::new_with_keys(Rand::rand(&mut rng), Rand::rand(&mut rng))
    }

    fn offset<Q: ?Sized>(&self, hashes: &mut [u64; 2], key: &Q, k_i: usize)
    -> usize where Q: Hash, K: Borrow<Q> {
        if k_i < 2 {
            let sip = &mut self.hashers[k_i as usize].clone();
            key.hash(sip);
            let hash = sip.finish();
            hashes[k_i as usize] = hash;
            hash as usize & self.mask
        } else {
            hashes[0]
                .wrapping_add((k_i as u64).wrapping_mul(hashes[1]) %
                              0xffffffffffffffc5) as usize & self.mask
        }
    }
}

)} // macro_rules! cms_define

cms_define!(CountMinSketch8, u8);
cms_define!(CountMinSketch16, u16);
cms_define!(CountMinSketch32, u32);
cms_define!(CountMinSketch64, u64);

#[cfg(test)]
mod tests {
    #[test]
    fn test_overflow() {
        use CountMinSketch8;

        let mut cms = CountMinSketch8::<&str>::new(100, 0.95, 10.0).unwrap();
        for _ in 0..300 {
            cms.increment("key");
        }
        assert!(cms.estimate("key") == u8::max_value());
    }

    #[test]
    fn test_increment() {
        use CountMinSketch16;

        let mut cms = CountMinSketch16::<&str>::new(100, 0.95, 10.0).unwrap();
        for _ in 0..300 {
            cms.increment("key");
        }
        assert!(cms.estimate("key") == 300);
    }

    #[test]
    fn test_increment_multi() {
        use CountMinSketch64;

        let mut cms = CountMinSketch64::<u64>::new(100, 0.99, 2.0).unwrap();
        for i in 0..1_000_000 {
            cms.increment(&(i % 100));
        }
        for key in 0..100 {
            assert!(cms.estimate(&key) >= 9_000);
        }
        cms.reset();
        for key in 0..100 {
            assert!(cms.estimate(&key) < 11_000);
        }
    }
}