import math
import mmh3
from bitarray import bitarray


class BloomFilter:
    def __init__(self, expected_num: int, false_positive_rate: float):
        """
        初始化布隆过滤器
        :param expected_num: 预期插入的元素数量（n）
        :param false_positive_rate: 可接受的误判率（p），如0.01表示1%
        """
        self.expected_num = expected_num  # 预期元素数n
        self.false_positive_rate = false_positive_rate  # 误判率p

        # 计算位数组长度m：m = - (n * ln(p)) / (ln(2)^2)
        self.bit_size = int(- (expected_num * math.log(false_positive_rate)) / (math.log(2) ** 2))
        # 计算哈希函数个数k：k = (m / n) * ln(2)
        self.hash_func_num = int((self.bit_size / expected_num) * math.log(2))

        # 初始化位数组（全为0）
        self.bit_array = bitarray(self.bit_size)
        self.bit_array.setall(0)

        # 记录实际插入的元素数量（用于后续可能的动态调整，可选）
        self.actual_num = 0

    def add(self, item) -> None:
        """插入元素到布隆过滤器"""
        if self.actual_num >= self.expected_num:
            print(f"警告：实际元素数已超过预期值{n}，误判率可能升高！")

        # 对元素进行k次哈希，得到k个索引
        for seed in range(self.hash_func_num):
            # 用mmh3哈希（支持种子，不同种子=不同哈希函数），取绝对值后对bit_size取模
            hash_value = mmh3.hash(str(item), seed)
            index = abs(hash_value) % self.bit_size
            self.bit_array[index] = 1  # 标记为1

        self.actual_num += 1

    def contains(self, item) -> bool:
        """查询元素是否可能存在（True=可能存在，False=一定不存在）"""
        for seed in range(self.hash_func_num):
            hash_value = mmh3.hash(str(item), seed)
            index = abs(hash_value) % self.bit_size
            if self.bit_array[index] == 0:
                return False  # 有一位为0，一定不存在
        return True  # 所有位为1，可能存在


# 测试代码
if __name__ == "__main__":
    # 初始化：预期插入1000个元素，误判率控制在1%
    bloom = BloomFilter(expected_num=1000, false_positive_rate=0.01)

    # 插入1000个元素
    for i in range(1000):
        bloom.add(i)

    # 测试已插入的元素（应全部返回True）
    error = 0
    for i in range(1000):
        if not bloom.contains(i):
            error += 1
    print(f"已插入元素的误判率（假阴性）：{error/1000:.2%}")  # 理论上应为0%

    # 测试未插入的元素（可能存在假阳性）
    false_positive = 0
    test_num = 10000  # 测试10000个未插入的元素
    for i in range(1000, 1000 + test_num):
        if bloom.contains(i):
            false_positive += 1
    print(f"未插入元素的实际误判率：{false_positive/test_num:.2%}")  # 应接近1%