from .utils import write_memmap
from .utils import read_memmap
from .utils import load_obj

import os
from collections import defaultdict

class BaseHMF():
	
	def __init__(self, root_dirpath, memmap_map, verbose):
		"""If root_dirpath exists, load the existing memmap so we can get arrays.
		
		(root)

		attributes
		dirpath
		node_type = root
		nodes
		
			group_name1

				attributes
				dirpath
				node_type = group
				nodes

					array_name1

						attributes
						dirpath
						node_type = array
						dtype
						shape

			group_name2

			group_name3 ...

		Question: why not combine the node name and the dict following it into a single dict?
		- this will in fact make searching thru nodes at nodes key much harder since we can't 
		  query it using key name, but entire dictionaries will come up --> forcing us to 
		  query using a dictionary.
		
		"""
		
		if(memmap_map):

			if(verbose):

				print('Loading from existing HMF')

			self.memmap_map = memmap_map

		else:

			if(verbose):

				print('Creating new HMF')

			self.memmap_map = dict()
			self.memmap_map['attributes'] = dict()
			self.memmap_map['nodes'] = dict()
			self.memmap_map['dirpath'] = root_dirpath
			self.memmap_map['node_type'] = 'root'

			self.memmap_map['dataframe_colnames'] = defaultdict(dict)
			self.memmap_map['grouped'] = False

			self.memmap_map['grouped'] = dict()
			self.memmap_map['group_sizes'] = dict()
			self.memmap_map['group_names'] = dict()
			self.memmap_map['group_items'] = dict()

			self.memmap_map['multi_pdfs'] = False

		self.verbose = verbose
	
	def set_group(self, group_dirpath):
		"""
		Todo: validate group_dirpath
		"""
		
		self.update_memmap_map_group(group_dirpath)

		
	def update_memmap_map_group(self, group_dirpath):
		
		memmap_map_pos = self.memmap_map['nodes']
		# dict with node_names of group names / array names

		if(group_dirpath=='/'):
			return 

		# iterate through each level of group (horizontal traversal)
		for idx, node_name in enumerate(group_dirpath.split('/')):

			# skip / because root is not part of dict ( (root) )
			if idx==0:

				if(node_name!=''):
					raise ValueError('Path much start with the root directory "/"')

				continue	

			# if node_name is one of memmap_map_pos keys
			if(node_name in memmap_map_pos):

				if self.verbose:
					print('{} found in memmap'.format(node_name))

				if(not self._is_group_node(memmap_map_pos, node_name)):

					if self.verbose:
						print('{} is not a group node so skipping'.format(node_name))

					continue

				# take the current level dirpath
				group_dirpath_pos = memmap_map_pos[node_name]['dirpath']

				# move to the groups dict position of the current node_name
				memmap_map_pos = memmap_map_pos[node_name]['nodes']
				
			else:

				if self.verbose:
					print('{} not found in memmap'.format(node_name))
				
				# if node_name is not part of the key, create the node_name's map position
				memmap_map_pos[node_name] = dict()
				memmap_map_pos[node_name]['nodes'] = dict()  # create node names dict
				memmap_map_pos[node_name]['attributes'] = dict()
				memmap_map_pos[node_name]['node_type'] = 'group'

				# if it's a node_name right after / root position
				# will always visit - setting the group_dirpath_pos always
				if idx==1:

					# start the dirpath afresh
					memmap_map_pos[node_name]['dirpath'] = node_name
					group_dirpath_pos = memmap_map_pos[node_name]['dirpath']

				else:

					# if it's beyond the level right after root /
					# the dirpath is combination of all previous node_name dirpaths
					memmap_map_pos[node_name]['dirpath'] = self._assemble_dirpath(group_dirpath_pos, node_name)

				# update the map position before moving to the next level of node_name (if there is any)
				memmap_map_pos = memmap_map_pos[node_name]['nodes']
				
	def retrieve_memmap_map_pos_group(self, group_dirpath):
		"""
		group_pos: the dictionary position where nodes, dirpath, node_type can be queried
		i.e. right below the node names dict level


		check if it is indeed a group

		"""

		if(group_dirpath=='/'):
			return self.memmap_map

		memmap_map_pos = self.memmap_map['nodes']
		
		for idx, node_name in enumerate(group_dirpath.split('/')):
			
			if idx==0:
				continue
			
			if node_name in memmap_map_pos:
			
				memmap_map_group_pos = memmap_map_pos[node_name]
				memmap_map_pos = memmap_map_group_pos['nodes']

			else:
				raise 
				
		return memmap_map_group_pos

	def set_array(self, array_filepath, array):
		
		# validate array
		# validate dirpath
		
		self.update_memmap_map_array(array_filepath, array)
		
		self.write_array(array_filepath, array)
		
	def update_memmap_map_array(self, array_filepath, array):
		
		array_filepaths = array_filepath.split('/')
		
		group_dirpath = '/'.join(array_filepaths[0:-1])
		array_name = array_filepaths[-1]
		
		if len(array_filepaths) > 2:
			self.update_memmap_map_group(group_dirpath)
			
			memmap_map_group_pos = self.retrieve_memmap_map_pos_group(group_dirpath)
		else:
			memmap_map_group_pos = self.memmap_map

		# TODO: check if array_name is already in here
		
		memmap_map_group_pos['nodes'][array_name] = dict()
		memmap_map_group_pos['nodes'][array_name]['node_type'] = 'array'
		memmap_map_group_pos['nodes'][array_name]['attributes'] = dict()
		memmap_map_group_pos['nodes'][array_name]['nodes'] = dict()

		if(memmap_map_group_pos['node_type']!='root'):

			filepath = self._assemble_dirpath(memmap_map_group_pos['dirpath'], array_name)

		else:

			filepath = array_name

		memmap_map_group_pos['nodes'][array_name]['dirpath'] = filepath
		
		dtype = str(array.dtype)
		shape = array.shape
		
		memmap_map_group_pos['nodes'][array_name]['dtype'] = dtype
		memmap_map_group_pos['nodes'][array_name]['shape'] = shape

	def retrieve_memmap_map_pos_array(self, array_filepath):
		
		array_filepaths = array_filepath.split('/')
		
		group_dirpath = '/'.join(array_filepaths[0:-1])
		array_name = array_filepaths[-1]
		
		if len(array_filepaths) > 2:
			self.update_memmap_map_group(group_dirpath)
			
			memmap_map_group_pos = self.retrieve_memmap_map_pos_group(group_dirpath)
		else:
			memmap_map_group_pos = self.memmap_map
			
		return memmap_map_group_pos['nodes'][array_name]

	def write_array(self, array_filepath, array):
		
		memmap_map_array_pos = self.retrieve_memmap_map_pos_array(array_filepath)
		
		dtype = memmap_map_array_pos['dtype']
		shape = memmap_map_array_pos['shape']
		filepath = os.path.join(
			self.memmap_map['dirpath'],
			memmap_map_array_pos['dirpath'])
		
		write_memmap(filepath, dtype, shape, array)

	def get_array(self, array_filepath, idx=None):
		
		memmap_map_array_pos = self.retrieve_memmap_map_pos_array(array_filepath)
		
		dtype = memmap_map_array_pos['dtype']
		shape = memmap_map_array_pos['shape']
		filepath = os.path.join(
			self.memmap_map['dirpath'],
			memmap_map_array_pos['dirpath'])
		
		return read_memmap(filepath, dtype, shape, idx)

	def set_node_attr(self, attr_dirpath, key, value):

		self.update_memmap_map_group(attr_dirpath)
			
		memmap_map_group_pos = self.retrieve_memmap_map_pos_group(attr_dirpath)

		memmap_map_group_pos['attributes'][key] = value

	def get_node_attr(self, attr_dirpath, key):
		
		memmap_map_group_pos = self.retrieve_memmap_map_pos_group(attr_dirpath)

		return memmap_map_group_pos['attributes'][key]

	def _assemble_dirpath(self, *args):

		args = [str(elem) for elem in args]

		return '__'.join(args)

	def _is_group_node(self, memmap_map_pos, node_name):

		return memmap_map_pos[node_name]['node_type'] == 'group'

	def _is_array_node(self, memmap_map_pos, node_name):

		return memmap_map_pos[node_name]['node_type'] == 'array'	  
	