import numpy as np
import os
import collections


class Index(object):
	def __init__(self):
		self.N2IMap = collections.OrderedDict()
		self.I2NMap = collections.OrderedDict()


	def pop_index(self, index):
		if index in self.I2NMap:
			row_name = self.I2NMap.pop(index)
			del self.N2IMap[row_name]
			return True
		return False


	def get_row_index(self, row_name):
		return self.N2IMap.get(row_name)


	def get_row_name(self, row_index):
		return self.I2NMap.get(row_index)


	def add_row(self, name, index):
		self.N2IMap[name] = index
		self.I2NMap[index] = name


	def flush(self, indexes):
		n2imap = collections.OrderedDict()
		i2nmap = collections.OrderedDict()
		i = 0
		saved_indexes = []
		for row_name, row_index in self.N2IMap.items():
			if row_index not in indexes:
				n2imap[row_name] = i
				i2nmap[i] = row_name
				saved_indexes.append(row_index)
				i += 1

		self.N2IMap = n2imap
		self.I2NMap = i2nmap
		return saved_indexes


	def serialize(self):
		return {
			"N2IMap": self.N2IMap,
			"I2NMap": self.I2NMap,
		}


	def deserialize(self, data):
		self.N2IMap = data["N2IMap"]
		self.I2NMap = data["I2NMap"]


	def extend(self, size):
		pass


	def __contains__(self, row_name):
		if row_name in self.N2IMap:
			return True

		return False


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



class PersistentIndex(Index):
	def __init__(self, path, size):
		super().__init__()
		self.DType = 'U30'
		self.Path = path

		if os.path.exists(self.Path):
			self.Map = np.memmap(self.Path, dtype=self.DType, mode='readwrite')
			for i in range(len(self.Map)):
				value = self.Map[i]
				if value != '':
					self.I2NMap[i] = value
					self.N2IMap[value] = i
		else:
			if size is None:
				raise RuntimeError("The size should correspond to array size")

			self.Map = np.memmap(self.Path, dtype=self.DType, mode='w+', shape=(size,))


	def pop_index(self, index):
		if super().pop_index(index):
			self.Map[index] = ''
			return True
		return False


	def add_row(self, name, index):
		super().add_row(name, index)
		self.Map[index] = name


	def extend(self, size):
		map_ = np.zeros(self.Map.shape[0], dtype=self.DType)
		map_[:] = self.Map[:]
		map_.resize(size, refcheck=False)
		self.Map = np.memmap(self.Path, dtype=self.DType, mode='w+', shape=map_.shape)
		self.Map[:] = map_[:]


	def flush(self, closed_indexes):
		saved_indexes = super().flush(closed_indexes)
		self.Map = self.Map.take(saved_indexes, axis=0)
		map_ = np.memmap(self.Path, dtype=self.DType, mode='w+', shape=self.Map.shape)
		map_[:] = self.Map[:]
		self.Map = map_
		return saved_indexes
