from .EvaluationInput import EvaluationInput
from .get_type import get_type

from slytherin.collections import remove_list_duplicates
from slytherin import get_size
from slytherin.hash import hash_object
from chronometry import Timer
from disk import Path
from joblib import delayed
from chronometry.progress import ProgressBar
from pandas import DataFrame, Series

import dill
import pickle


class Memory:
	def __init__(
			self, key, pensieve, function, label=None, precursors=None, safe=True, metadata=False, materialize=True,
			_update=True, _stale=True
	):
		"""
		:param str key: unique name/identifier of the memory
		:param Pensieve pensieve: the pensieve this memory belongs to
		:param callable function: a function to be called on precursor memories
		:param list[Memory] or NoneType precursors: precursor memories to this memory
		:param bool safe: when True only a copy of the content is returned to avoid mutating it from outside
		:param dict metadata: an optional dictionary that carries meta data about this memory
		:param bool materialize: when False, memory runs the function when it needs the content, rather than keeping it
		:param bool _update: if True the precursors will be updated
		:param bool _stale:
		"""
		# make precursors unique
		precursors = precursors or []
		self._key = key
		self._label = label
		self._pensieve = pensieve
		self._content = None
		self._materialize = materialize
		self._safe = safe
		if self.pensieve is not None:
			if self.key not in self.pensieve._successor_keys:
				self.pensieve._successor_keys[self.key] = []
			if self.key not in self.pensieve._precursor_keys:
				self.pensieve._precursor_keys[self.key] = []
		self._frozen = False
		self._deep_freezed = False
		self._stale = _stale
		self._function = function
		self._metadata = metadata or {}

		self._total_time = None

		self._size = None
		self._precursors_hash = None
		self._content_type = None
		self._content_access_count = 0

		if _update:
			self.update(precursors, function)

	__PARAMS__ = [
		'key', 'label', 'materialize', 'safe', 'frozen', 'deep_freezed', 'stale', 'metadata', 'total_time', 'size',
		'precursors_hash', 'content_type', 'content_access_count'
	]

	@property
	def num_threads(self):
		return self.pensieve._num_threads

	def clean_copy(self, include_function=False, stale=False, update=False):

		result = self.__class__(
			key=self.key, pensieve=None,
			function=self._function if include_function else None,
			precursors=None,
			safe=self._safe, metadata=self._metadata.copy(),
			materialize=self._materialize, _update=update, _stale=stale
		)
		return result

	def partial_copy(self, include_function=False, stale=False, update=False, include_precursor_hash=True):
		result = self.clean_copy(include_function=include_function, stale=stale, update=update)
		result._content = self._content
		result._content = self._content
		result._frozen = self._frozen
		result._total_time = self._total_time
		result._size = self._size
		result._precursors_hash = self._precursors_hash if include_precursor_hash else None
		result._content_type = self._content_type
		result._content_access_count = self._content_access_count
		return result

	def __hashkey__(self):
		return self.__class__.__name__, self.parameters, self.precursor_keys

	@property
	def parameters(self):
		return {param: getattr(self, f'_{param}') for param in self.__PARAMS__}

	def __getstate__(self):
		"""
		:rtype: dict
		"""
		stale = self._stale
		try:
			function_dump = dill.dumps(obj=self._function, )
		except RecursionError as e:
			print(f'\nrecursion error during the dill.dumping of the "{self.key}" memory \n\n')
			raise e

		state = {
			'parameters': self.parameters,
			'function': function_dump
		}

		if not stale:
			try:
				state['serialized'] = pickle.dumps(obj=self._content, protocol=pickle.HIGHEST_PROTOCOL)
				state['serialized_by'] = 'pickle'
			except:
				try:
					state['serialized'] = dill.dumps(obj=self._content, protocol=dill.HIGHEST_PROTOCOL)
					state['serialized_by'] = 'dill'
				except:
					state['serialized_by'] = None
		else:
			state['serialized_by'] = None

		return state

	def __setstate__(self, state):
		"""
		:type state: dict
		"""
		parameters = state['parameters']
		for name, value in parameters.items():
			setattr(self, f'_{name}', value)
		if self._stale:
			self._content = None
			self._precursors_hash = None
		else:
			if state['serialized_by'] == 'dill':
				self._content = dill.loads(str=state['serialized'])

			elif state['serialized_by'] == 'pickle':
				self._content = pickle.loads(state['serialized'])

			else:
				self._stale = True
				self._content = None
				self._precursors_hash = None
		self._function = dill.loads(str=state['function'])
		self._pensieve = None

	def save(self, path):
		"""
		:type path: str or Path
		"""
		path = Path(string=path)
		path.make_dir()
		parameters = {param: getattr(self, f'_{param}') for param in self.__PARAMS__}
		try:
			(path + 'content.pensieve').save(obj=self._content)
		except:
			parameters['stale'] = True
		(path + 'parameters.pensieve').save(obj=parameters)
		(path + 'function.pensieve').save(obj=self._function, method='dill')

	@classmethod
	def load(cls, path, pensieve):
		path = Path(string=path)
		parameters = (path + 'parameters.pensieve').load()
		function = (path + 'function.pensieve').load(method='dill')
		memory = cls(
			pensieve=pensieve, function=function, precursors=None,
			key=parameters['key'], _update=False
		)
		for name, value in parameters.items():
			setattr(memory, f'_{name}', value)
		memory._function = function
		try:
			memory._content = (path + 'content.pensieve').load()
		except:
			memory._content = None
			memory._stale = True
		return memory

	def get_summary(self):
		result = {
			'key': self.key,
			'content_type': self._content_type,
			'materialized': self._materialize,
			'frozen': self._frozen,
			'evaluation_time': self.evaluation_time,
			'total_time': self.total_time,
			'size': self.size,
			'precursors_hash': self._precursors_hash,
			'precursors': self.precursor_keys
		}
		for key, value in self._metadata:
			new_key = f'metadata_{key}' if key in result else key
			result[new_key] = value
		return result

	@property
	def size(self):
		if self._size is None:
			result = 0
			result += get_size(self._key, exclude_objects=[self._pensieve])
			result += get_size(self._content, exclude_objects=[self._pensieve])
			result += get_size(self._precursors_hash, exclude_objects=[self._pensieve])
			result += get_size(self._content_type, exclude_objects=[self._pensieve])
			result += get_size(self._safe, exclude_objects=[self._pensieve])
			result += get_size(self._frozen, exclude_objects=[self._pensieve])
			result += get_size(self._stale, exclude_objects=[self._pensieve])
			result += get_size(self._function, exclude_objects=[self._pensieve])
			result += get_size(self._metadata, exclude_objects=[self._pensieve])
			result += get_size(result)
			self._size = result
		return self._size

	@property
	def evaluation_time(self):
		if self.key in self.pensieve.function_durations.measurements:
			return self.pensieve.function_durations.measurements[self.key].mean_duration
		else:
			return None

	@property
	def total_time(self):
		if self._total_time is None:
			if len(self.precursor_keys) == 0:
				self._total_time = self.evaluation_time
			else:
				self._total_time = self.evaluation_time
				for precursor in self.precursors:
					self._total_time += precursor.total_time
		return self._total_time

	@property
	def speed(self):
		evaluation_time = self.evaluation_time
		if evaluation_time is not None:
			return self.size / evaluation_time
		else:
			return None

	def display(self, p=None):
		print(self._get_dependency_summary())
		if p is None:
			try:
				from IPython.core.display import display
				display(self.content)
			except ImportError:
				print(self.content)
		else:
			p.pretty(self.content)

	def _get_dependency_summary(self):
		precursor_keys = self.precursor_keys
		successor_keys = self.successor_keys
		if len(precursor_keys) > 0:
			precursors = 'precursors: ' + ', '.join(precursor_keys)
		else:
			precursors = 'no precursors'

		if len(successor_keys) > 0:
			successors = 'successors: ' + ', '.join(successor_keys)
		else:
			successors = 'no successors'
		return f'{precursors}\n{successors}'

	def __sizeof__(self):
		return self.size

	def __eq__(self, other):
		if not isinstance(other, self.__class__):
			return False
		if self.key != other.key:
			return False

		all_keys = set(self.precursor_keys).union(other.precursor_keys)
		if len(all_keys) > len(self.precursor_keys) or len(all_keys) > len(other.precursor_keys):
			return False

		if self.content != other.content:
			return False

		return True

	def __hash__(self):
		return hash(self.key)

	def __repr__(self):
		return f'Memory:{self.key}'

	@property
	def is_frozen(self):
		return self._frozen

	@property
	def is_stale(self):
		return self._stale

	def freeze(self, forever=False):
		self._frozen = True
		self._deep_freezed = forever
		if forever:
			self._function = None

	def deep_freeze(self):
		self.freeze(forever=True)

	def unfreeze(self):
		if not self._deep_freezed:
			self._frozen = False
			if self._stale:
				self.mark_stale()
		else:
			print(f'{self.key} is deep-freezed and cannot be thawed!')

	@property
	def pensieve(self):
		"""
		:rtype: .Pensieve.Pensieve
		"""
		return self._pensieve

	@property
	def key(self):
		return self._key

	@property
	def label(self):
		output = self._label or self.key.replace('__', '\n').replace('_', ' ')

		if self._deep_freezed:
			frozen_label = 'deep-freezed'
		else:
			frozen_label = 'frozen'

		if self.is_stale and self.is_frozen:
			output += f'\n( stale & {frozen_label} )'
		elif self.is_stale and not self.is_frozen:
			output += '\n( stale )'
		elif not self.is_stale and self.is_frozen:
			output += f'\n( {frozen_label} )'
		else:
			output += f'\n{self._content_type}'

		return output

	@property
	def precursor_keys(self):
		"""
		:type: list[str]
		"""
		return list(self.pensieve.get_precursor_keys(memory=self))

	@property
	def successors(self):
		return self.pensieve.get_successors(memory=self)

	@property
	def has_precursors(self):
		return len(self.precursor_keys) > 0

	@property
	def successor_keys(self):
		"""
		:type: list[str]
		"""
		return list(self.pensieve.get_successor_keys(memory=self))

	@property
	def precursors(self):
		"""
		:rtype: list[Memory]
		"""
		return self.pensieve.get_precursors(memory=self)

	@property
	def has_successors(self):
		return len(self.successor_keys) > 0

	def erase_successor(self, successor):
		"""
		:param Memory or str successor: the successor memory or its key that should be removed
		"""
		if isinstance(successor, str):
			self.pensieve._successor_keys[self.key].remove(successor)
		else:
			self.pensieve._successor_keys[self.key].remove(successor.key)

	# ************************* COMPUTATION **********************************

	def update(self, precursors, function, label=None, metadata=None, materialize=None):
		"""
		:type precursors: list[Memory]
		:type function: callable
		:type metadata: NoneType or dict
		:type materialize: bool or NoneType
		"""
		# make precursors unique:
		if self.is_frozen:
			raise MemoryError(f'{self.key} is frozen. You cannot change a frozen memory!')

		precursors = precursors or []
		precursors = remove_list_duplicates(precursors)

		precursor_keys = [p.key for p in precursors]

		removed_precursor_keys = [key for key in self.precursor_keys if key not in precursor_keys]
		new_precursor_keys = [key for key in precursor_keys if key not in self.precursor_keys]

		self.pensieve._precursor_keys[self.key] = precursor_keys
		for precursor_key in removed_precursor_keys:
			self.pensieve._successor_keys[precursor_key].remove(self.key)
		for precursor_key in new_precursor_keys:
			self.pensieve._successor_keys[precursor_key].append(self.key)

		self._function = function
		self.mark_stale()

		if metadata is not None:
			self._metadata = metadata
		if materialize is not None:
			self._materialize = materialize

		if label is not None:
			self._label = label

	def evaluate(self):
		_ = self.content

	@property
	def stale_precursors(self):
		"""
		:rtype: list[Memory]
		"""
		return [precursor for precursor in self.precursors if precursor.is_stale]

	def _get_stale_dependencies(self, unique=False):
		"""
		:rtype: list[Memory]
		"""
		dependencies = []
		for precursor in self.stale_precursors:
			dependencies.append(precursor)
			dependencies += precursor._get_stale_dependencies(unique=False)

		if unique:
			unique_dependencies = []
			for dependency in dependencies:
				if dependency not in unique_dependencies:
					unique_dependencies.append(dependency)
			return unique_dependencies
		else:
			return dependencies

	@property
	def stale_dependencies(self):
		"""
		:rtype: list[Memory]
		"""
		return self._get_stale_dependencies(unique=True)

	def get_update_schedule(self):
		"""
		:rtype: list[list[Memory]]
		"""
		jobs = self.stale_dependencies
		jobs.reverse()

		added_to_schedule = []
		schedule = []
		while len(jobs) > 0:
			job_round = []
			for job in jobs:
				if len(job.stale_precursors) == 0:
					job_round.append(job)
				elif all([dependency in added_to_schedule for dependency in job.stale_dependencies]):
					job_round.append(job)

			for job in job_round:
				jobs.remove(job)
			added_to_schedule += job_round
			schedule.append(job_round)
		return schedule

	@property
	def type_significance(self):
		if isinstance(self._content, (DataFrame, Series, list, dict, set)):
			return 2
		else:
			return 1

	@property
	def content(self):
		if not self._materialize:
			self.set_content(content=None, precursors_hash=None)
			content, precursors_hash = self.get_content_and_hash()

		elif self.is_frozen or not self.is_stale:
			content = self._content

		else:
			content, precursors_hash = self.get_content_and_hash()
			self.set_content(content=content, precursors_hash=precursors_hash)
			content = self._content

		return content

	def set_content(self, content, precursors_hash):
		if self.is_frozen:
			raise MemoryError(f'{self.key} is frozen. You cannot change a frozen memory!')
		self._content = content
		self._stale = False
		self._precursors_hash = precursors_hash

	def mark_stale(self):
		if self._materialize:
			self._stale = True
		self._size = None
		for successor in self.successors:
			successor.mark_stale()

	def get_content_and_hash(self):
		if self.num_threads == 1:
			precursor_keys_to_contents = {p.key: p.content for p in self.precursors}
		else:
			def get_content(p):
				return p.content

			precursors = self.precursors

			schedule = self.get_update_schedule()

			progress_bar = ProgressBar(
				total=sum([len(schedule_round) for schedule_round in schedule]),
				echo=self.pensieve._echo
			)

			progress_amount = 0
			for schedule_round in schedule:
				progress_bar.show(amount=progress_amount, text=f'updating {len(schedule_round)} memories')
				self.pensieve.processor(delayed(get_content)(job) for job in schedule_round)
				progress_amount += len(schedule_round)
			if progress_amount > 0:
				progress_bar.show(amount=progress_amount, text=f'{self.key} updated!')

			contents = self.pensieve.processor(delayed(get_content)(p) for p in precursors)
			keys = [precursor.key for precursor in precursors]
			precursor_keys_to_contents = {key: content for key, content in zip(keys, contents)}

		if len(self.precursor_keys) == 0:
			new_hash = hash_object(self._function)
			if new_hash == self._precursors_hash and self._materialize:
				new_content = self._content

			else:
				timer = Timer(start_now=True, unit='timedelta')
				new_content = self._function()
				timer.stop()
				self.pensieve.function_durations.add_measurement(name=self.key, timer=timer)

		elif len(self.precursor_keys) == 1:
			precursor_content = list(precursor_keys_to_contents.values())[0]
			new_hash = hash_object((self._function, precursor_content))
			if new_hash == self._precursors_hash and self._materialize:
				new_content = self._content

			else:
				timer = Timer(start_now=True, unit='timedelta')
				new_content = self._function(precursor_content)
				timer.stop()
				self.pensieve.function_durations.add_measurement(name=self.key, timer=timer)

		else:
			inputs = EvaluationInput(inputs=precursor_keys_to_contents)
			new_hash = hash_object((self._function, inputs))
			if new_hash == self._precursors_hash and self._materialize:
				new_content = self._content

			else:
				timer = Timer(start_now=True, unit='timedelta')
				new_content = self._function(inputs.originals)

				timer.stop()
				self.pensieve.function_durations.add_measurement(name=self.key, timer=timer)

		self._content_type = get_type(new_content)

		self._content_access_count += 1
		return new_content, new_hash

	@property
	def graphviz_edges_str(self):
		if not self.has_precursors:
			return self.label if not self.has_successors else None
		else:
			edges = [
				f'{precursor.label} -> {self.label}'
				for precursor in self.precursors
			]
			return '\n'.join(edges)

	def __graph_node__(self):
		"""
		:rtype: dict
		"""
		return {
			'label': self.label,
			'value': None,
			'metadata': self._metadata
		}
