import numpy as np
import os


class ClosedRows(object):
	def __init__(self, max_len=None):
		self.CR = set()
		if max_len is None:
			max_len = float('inf')

		self.MaxLen = max_len


	def pop(self):
		return self.CR.pop()


	def get_rows(self):
		return self.CR


	def add(self, element):
		if len(self.CR) == self.MaxLen:
			raise RuntimeError("Maximum size exceeded")

		self.CR.add(element)


	def __contains__(self, element):
		if element in self.CR:
			return True
		return False


	def serialize(self):
		return list(self.get_rows())


	def deserialize(self, data):
		self.CR = set(data)


	def __len__(self):
		return len(self.CR)


	def extend(self, start, stop):
		self.CR |= frozenset(range(start, stop))
		if len(self.CR) >= self.MaxLen:
			raise RuntimeError("Maximum size exceeded")


	def flush(self, size=None):
		self.CR = set()



class PersistentClosedRows(ClosedRows):
	def __init__(self, path, size=None, max_len=None):
		super().__init__(max_len=max_len)
		self.DType = 'i1'
		self.Path = path
		if os.path.exists(self.Path):
			self.CRBit = np.memmap(self.Path, dtype=self.DType, mode='readwrite')
			for i in range(len(self.CRBit)):
				if self.CRBit[i] == 0:
					self.CR.add(i)
		else:
			if size is None:
				raise RuntimeError("The size should correspond to array size")
			self.CR.add(0)
			self.ones(size)


	def pop(self):
		element = super().pop()
		self.CRBit[element] = 1
		return element


	def add(self, element):
		super().add(element)
		self.CRBit[element] = 0


	def extend(self, start, stop):
		cr_ = np.zeros(self.CRBit.shape[0], dtype=self.DType)
		cr_[:] = self.CRBit[:]
		cr_.resize(stop, refcheck=False)
		self.CRBit = np.memmap(self.Path, dtype=self.DType, mode='w+', shape=cr_.shape)
		self.CRBit[:] = cr_[:]
		super().extend(start, stop)


	def flush(self, size):
		super().flush(size)
		self.ones(size)


	def ones(self, size):
		self.CRBit = np.memmap(self.Path, dtype=self.DType, mode='w+', shape=(size,))
		self.CRBit[:] = 1
