#coding:utf-8
# Bloom filters in Python
# Adam Langley <agl@imperialviolet.org>
# 给CountedBloom加了一个max_count 张沈鹏 <zsp007@gmail.com>
# Bloom-Filter算法简介
# http://www.googlechinablog.com/2007/07/bloom-filter.html
# http://zh.wikipedia.org/wiki/%E5%B8%83%E9%9A%86%E8%BF%87%E6%BB%A4%E5%99%A8
# 这个计算器可以帮你求最佳的参数
# http://www.cc.gatech.edu/~manolios/bloom-filters/calculator.html
# CountedBloom 的 buckets 参数对应于计算器的m,也就是"m denotes the number of bits in the Bloom filter"

import array
import struct

mixarray = array.array ('B', '\x00' * 256)
# The mixarray is based on RC4 and is used as diffusion in the hashing function

def mixarray_init (mixarray):
    for i in range (256):
        mixarray[i] = i
    k = 7
    for j in range (4):
        for i in range (256):
            s = mixarray[i]
            k = (k + s) % 256
            mixarray[i] = mixarray[k]
            mixarray[k] = s

mixarray_init (mixarray)

class Bloom (object):
    '''Bloom filters provide a fast and compact way of checking set membership. They do this by introducing a risk of a 
  false positive (but there are no false negatives).

  For more information see http://www.cs.wisc.edu/~cao/papers/summary-cache/node8.html'''
    def __init__ (self, bytes, hashes, data = None):
        '''@bytes is the size of the bloom filter in 8-bit bytes and @hashes is the number of hash functions to use. Consult the
    web page linked above for values to use. If in doubt, bytes = num_elements and hashes = 4'''
        self.hashes = hashes
        self.bytes = bytes

        if data == None:
            self.a = self._make_array (bytes)
        else:
            assert len (data) == bytes
            self.a = data

    def init_from_counted (self, cnt):
        '''Set the contents of this filter from the contents of the counted filter @cnt. You have to match sizes'''
        if self.bytes * 8 != (len (cnt.a) * 2):
            raise ValueError ('Filters are not the same size')
        for i in xrange (len (cnt.a)):
            b = cnt.a[i]
            b1 = (b & 0xf0) >> 4
            b2 = (b & 0x0f)
            if b1:
                self.a[(i * 2) // 8] |= self.bitmask[(i * 2) % 8]
            if b2:
                self.a[(i * 2 + 1) // 8] |= self.bitmask[(i * 2 + 1) % 8]

    def _make_array (self, size):
        a = array.array ('B')
        # stupidly, there's no good way that I can see of resizing an array without allocing a huge string to do so
        # thus I use this, slightly odd, method:

        blocklen = 256
        arrayblock = array.array ('B', '\x00' * blocklen)
        todo = size
        while (todo >= blocklen):
            a.extend (arrayblock)
            todo -= blocklen
        if todo:
            a.extend (array.array ('B', '\x00' * todo))
        # now a is of the right length
        return a

    def _hashfunc (self, n, val):
        '''Apply the nth hash function'''

        global mixarray

        b = [ord(x) for x in struct.pack ('I', val)]
        c = array.array ('B', [0, 0, 0, 0])
        for i in range (4):
            c[i] = mixarray[(b[i] + n) % 256]

        return struct.unpack ('I', c.tostring())[0]

    bitmask = [0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01]

    def insert (self, val):
        for i in range (self.hashes):
            n = self._hashfunc (i, val) % (self.bytes * 8)
            self.a[n // 8] |= self.bitmask[n % 8]

    def __contains__ (self, val):
        for i in range (self.hashes):
            n = self._hashfunc (i, val) % (self.bytes * 8)
            if not self.a[n // 8] & self.bitmask[n % 8]:
                return 0
        return 1
MAX_COUNT = 15
class CountedBloom (Bloom):
    '''Just like a Bloom filter, but provides counting (e.g. you can delete as well). This uses 4 bits per bucket, so is
  generally four times larger than the same non-counted bloom filter.'''

    def __init__ (self, buckets, hashes):
        '''Please note that @buckets must be even. Also note that with a Bloom object you give the number of *bytes* and each byte is 8 buckets. Here you're giving the number of buckets.'''
        assert buckets % 2 == 0
        self.hashes = hashes
        self.buckets = buckets

        self.a = self._make_array (buckets // 2)

    def insert (self, val):
        masks  = [(0x0f, 0xf0), (0xf0, 0x0f)]
        shifts = [4, 0           ]

        for i in range (self.hashes):
            n = self._hashfunc (i, val) % self.buckets
            byte         = n // 2
            bucket = n % 2
            (notmask, mask) = masks[bucket]
            shift        = shifts[bucket]
            bval         = ((self.a[byte] & mask) >> shift)
            if bval < MAX_COUNT: # we shouldn't increment it if it's at the maximum
                bval += 1

            self.a[byte] = (self.a[byte] & notmask) | (bval << shift)
    def __contains__ (self, val):
        masks        = [(0x0f, 0xf0), (0xf0, 0x0f)]
        shifts = [4, 0]

        for i in range (self.hashes):
            n = self._hashfunc (i, val) % self.buckets
            byte         = n // 2
            bucket = n % 2
            (notmask, mask) = masks[bucket]
            shift        = shifts[bucket]
            bval          = ((self.a[byte] & mask) >> shift)

            if bval == 0:
                return 0
        return 1

    def max_count(self, val):
        masks        = [(0x0f, 0xf0), (0xf0, 0x0f)]
        shifts = [4, 0]
        count_val = MAX_COUNT
        for i in range (self.hashes):
            n = self._hashfunc (i, val) % self.buckets
            byte         = n // 2
            bucket = n % 2
            (notmask, mask) = masks[bucket]
            shift        = shifts[bucket]
            bval          = ((self.a[byte] & mask) >> shift)

            if bval < MAX_COUNT:
                if bval == 0:
                    return 0
                else:
                    count_val = bval
        return count_val

    def __delitem__ (self, val):
        masks  = [(0x0f, 0xf0), (0xf0, 0x0f)]
        shifts = [4, 0]

        for i in range (self.hashes):
            n = self._hashfunc (i, val) % self.buckets
            byte         = n // 2
            bucket = n % 2
            (notmask, mask) = masks[bucket]
            shift        = shifts[bucket]
            bval          = ((self.a[byte] & mask) >> shift)

            if bval < MAX_COUNT: # we shouldn't decrement it if it's at the maximum
                bval -= 1

            self.a[byte] = (self.a[byte] & notmask) | (bval << shift)

__all__ = ['Bloom']

if __name__ == '__main__':
    print 'Testing bloom filter: there should be no assertion failures'
    a = Bloom (3, 4)

    a.insert (45)
    print a.a
    a.insert (17)
    print a.a
    a.insert (12)
    print a.a
    assert 45 in a

    assert 45 in a
    assert not 33 in a
    assert 45 in a
    assert 17 in a
    assert 12 in a

    c = 0
    for x in range (255):
        if x in a:
            c += 1
    print c
    print float(c)/255


    a = CountedBloom (24, 4)
    a.insert (45)
    print a.a
    a.insert (17)
    print a.a
    a.insert (12)
    a.insert (12)
    print "a.max_count(12)", a.max_count(12)
    a.insert ("张沈鹏")
    a.insert ("张沈鹏")
    a.insert ("张沈鹏")
    print "a.max_count(zsp)", a.max_count(12)
    print a.a
    assert 45 in a

    assert 45 in a
    assert not 33 in a
    assert 45 in a
    assert 17 in a
    assert 12 in a

    c = 0
    for x in range (255):
        if x in a:
            c += 1
    print c
    print float(c)/255

    del a[45]
    assert not 45 in a

    a2 = Bloom (3, 4)
    a2.init_from_counted (a)

    print a2.a

    assert 17 in a2
    assert 12 in a2
    assert not 45 in a