
import tqdm
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

import utils.filesys as fs
import utils.utils as ut
import utils.tfw as tfw
import utils.printer as pr
import utils.datasets as ds



""" A combination of modules that interacts with environment """
class Model:
	def __init__(self, name, **kwargs):
		self.name = name
		self._make_model(**kwargs)
		self._set_description()
		self._make_var_dict()
		self._make_session()

	def __del__(self):
		self.sess.close()

	def _set_description(self, **kwargs):
		desc = 'name: {0}'.format(self.name)
		for module in self.modules:
			desc += '\n/{0}: {1}'.format(module.name, module.get_description())
		self._make_description(desc)

	def _make_description(self, desc):
		self.description = desc
		self.desc = tf.summary.text(
			'description', tf.convert_to_tensor(desc))

	def _make_model(self, **kwargs):
		self.modules = []

	def _make_var_dict(self):
		self.var_dict = {}
		for module in self.modules:
			self.var_dict.update(module.var_dict)

	def _make_session(self):
		self.model_saver = tf.train.Saver(self.var_dict)
		self.initializer = tf.global_variables_initializer()
		self.sess = tf.Session()
		self.reset()
		self.print_num_params()

	def print_num_params(self):
		for module in self.modules:
			pr.log('# of params in {0}: {1}'.format(
				module.name, module.get_num_params()))

	def reset(self):
		self.sess.run(self.initializer)

	def save(self, path):
		fs.check_dir(path, make=True)
		self.model_saver.save(self.sess, path)

	def load(self, path):
		fs.check_dir(path, fatal=True)
		self.model_saver.restore(self.sess, path)






#---------------------------------------------------------------------------#
#--------------------------------- System ----------------------------------#
#---------------------------------------------------------------------------#

class System:
	def __init__(self, dataset):
		self.dataset = dataset
	
	def train_x(self, train_func, batch_count=100, batch_size=20, procs=[], **kwargs):
		gen = ds.batch_generator(self.dataset['train'], batch_count, batch_size)
		for batch in gen:
			result = train_func(batch['x'], **kwargs)
			for proc in procs:
				proc.process(batch, result)
	
	def train_xy(self, train_func, batch_count=100, batch_size=20, procs=[], **kwargs):
		gen = ds.batch_generator(self.dataset['train'], batch_count, batch_size)
		for batch in gen:
			result = train_func(batch['x'], batch['y'], **kwargs)
			for proc in procs:
				proc.process(batch, result)
	
	def train_batch(self, train_func, batch_count=100, batch_size=20, procs=[], **kwargs):
		gen = ds.batch_generator(self.dataset['train'], batch_count, batch_size)
		for batch in gen:
			result = train_func(batch, batch_count, **kwargs)
			for proc in procs:
				proc.process(batch, result)

	def test_x(self, test_func, batch_size=100):
		gen = ds.epoch_generator(self.dataset['test'], batch_size)
		results = []
		for batch in gen:
			results.append(test_func(batch['x']))
		return np.mean(results)

	def test_xy(self, test_func, batch_size=100):
		gen = ds.epoch_generator(self.dataset['test'], batch_size)
		results = []
		for batch in gen:
			results.append(test_func(batch['x'], batch['y']))
		return np.mean(results)




#---------------------------------------------------------------------------#
#------------------------------- Processors --------------------------------#
#---------------------------------------------------------------------------#


class ResultProcessor:
	def process(self, batch, result):
		pass

class Logger(ResultProcessor):
	def __init__(self, model, log_dir='./log/', scalar_step=1, image_step=1000):
		self.log_dir = log_dir + str(ut.generate_id()) + '/'
		fs.check_dir(self.log_dir, make=True)
		self.summary_saver = tf.summary.FileWriter(self.log_dir, model.sess.graph)
		self.log(model.sess.run(model.desc))
		self.scalar_step = scalar_step
		self.image_step = image_step

	def process(self, batch, result):
		gs = result['global_step']
		s = result.get('summary', None)
		si = result.get('summary-image', None)
		if gs % self.scalar_step == 0:
			self.log(s, gs)
		if gs % self.image_step == 0:
			self.log(si, gs)

	def log(self, summary=None, global_step=0):
		if summary is not None:
			self.summary_saver.add_summary(summary, global_step)

class Reporter(ResultProcessor):
	def __init__(self, steps=100, kwords=[], log_dir=None):
		self.steps = steps
		self.kwords = kwords
		self.log_dir = log_dir

	def process(self, batch, result):
		step = batch['step']
		if step % self.steps == 0:
			report = '[step {0}]'.format(step)
			if 'global_step' in result:
				report += '[gstep {0}]'.format(result['global_step'])
			for word in self.kwords:
				report += '[{0} {1:.4f}]'.format(word, result[word])
			tqdm.tqdm.write(report)
			self.log2file(report)

	def log2file(self, msg):
		if self.log_dir is not None:
			with open(self.log_dir + "report.txt", "a") as file:
				file.write(msg + "\n")
				file.close()





class Presenter(ResultProcessor):
	def __init__(self, name, sess, func, fig_num=1, steps=100, logger=None):
		self.name = name
		self.sess = sess
		self.func = func
		self.fig_num = fig_num
		self.steps = steps
		self.logger = logger
		plt.ion()
		if logger is not None:
			image = func()
			self.image_ph = tf.placeholder(shape=image.shape, dtype=tf.float32)
			self.image_r = tfw.compact(self.image_ph)
			self.summary = tf.summary.image(name, self.image_r)

	def __del__(self):
		plt.ioff()
		plt.show()

	def process(self, batch, result):
		gstep = result['global_step']
		if gstep % self.steps == 0:
			images = self.func()
			# Summarizing the image
			if self.logger is not None:
				self.logger.log(
					self.sess.run(self.summary, {self.image_ph: images}), gstep)
			# Displaying the image
			plt.figure(self.fig_num)
			ut.show2(images)
			plt.pause(1e-5)

