#!/usr/bin/env python3

import os
from termcolor import colored, cprint
import requests
import sys
from tqdm import tqdm
import json

import tensorflow as tf
import numpy as np
import model, sample, encoder
import gpt_2_simple as gpt2

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
tf.logging.set_verbosity(tf.logging.ERROR)

class GPT2Client(object):
	def __init__(self, model_name='117M', save_dir='models'):
		"""
		Attributes
		----------
		attr: model_name (string)
			- default: '117M'
			- desc: Downloads the '117M' GPT-2 model. Can be set to '345M' model 
		
		attr: save_dir (string)
			- default: 'models'
			- desc: Name of directory where the weights, checkpoints, and 
					hyper-parameters are downloaded and saved
		"""
		self.model_name = model_name
		self.save_dir = save_dir

	def download_model(self):
		""" Creates `models` directory and downloads model weights and checkpoints """

		if self.model_name not in ['117M', '345M']:
			raise AssertionError('Please choose from either 117M or 345M parameter models only. This library does support other model sizes.')
		else:
			subdir = os.path.join(self.save_dir, self.model_name)
			if not os.path.exists(subdir):
				os.makedirs(subdir)
			
			for filename in ['checkpoint', 'encoder.json', 'hparams.json', 'model.ckpt.data-00000-of-00001', 'model.ckpt.index', 'model.ckpt.meta', 'vocab.bpe']:
				r = requests.get('https://storage.googleapis.com/gpt-2/models/' + self.model_name + '/' + filename, stream=True)

				with open(os.path.join(subdir, filename), 'wb') as f:
					file_size = int(r.headers['content-length'])
					chunk_size = 1000
					with tqdm(ncols=100, desc='Downloading {}'.format(colored(filename, 'cyan', attrs=['bold'])), total=file_size, unit_scale=True) as pbar:
						for chunk in r.iter_content(chunk_size=chunk_size):
							f.write(chunk)
							pbar.update(chunk_size)

	def generate(self, interactive=False, n_samples=1, words=None, display=True, return_text=False):
		""" Returns generated text sample
		
		Parameters
		----------
		arg: interactive (bool)
			- default: False
			- desc: Toggles interactive mode which prompts user for input text

		arg: n_samples (int)
			- default: 0
			- desc: Number of samples to be generated by GPT-2 Model. If 0, it generates indefinitely
		
		arg: words (int)
			- default: None
			- desc: Number of words in generated sample

		arg: display (bool)
			- default: True
			- desc: Prints out text to console when set to True

		arg: return_true (bool)
			- default: False
			- desc: Returns generated text when set to True

		Returns:
			Generated string
		"""

		if not interactive:
			# Generate random samples from scratch
			print (colored('Generating sample...', 'yellow'))

			models_dir = models_dir = os.path.expanduser(os.path.expandvars(self.save_dir))
			enc = encoder.get_encoder(self.model_name, self.save_dir)
			hparams = model.default_hparams()

			with open(os.path.join(self.save_dir, self.model_name, 'hparams.json')) as f:
				data = json.load(f)
				hparams.override_from_dict(data)

			length = hparams.n_ctx

			with tf.Session(graph=tf.Graph()) as sess:
				np.random.seed(None)
				tf.set_random_seed(None)

				batch_size = 1
				temperature = 1
				top_k = 40

				output = sample.sample_sequence(
					hparams=hparams,
					length=length,
					start_token=enc.encoder['<|endoftext|>'],
					batch_size=batch_size,
					temperature=temperature, 
					top_k=top_k
				)

				saver = tf.train.Saver()
				ckpt = tf.train.latest_checkpoint(os.path.join(self.save_dir, self.model_name))
				saver.restore(sess, ckpt)

				generated = 0
				text = []
				while n_samples == 0 or generated < n_samples:
					out = sess.run(output)
					for i in range(batch_size):
						generated += batch_size
						text.append(enc.decode(out[i]))
						print (colored('---------------------SAMPLE---------------------\n', 'cyan'))

						if display:
							print (text)

						if return_text:
							return text

		else:
			# Generate random samples from prompt
			models_dir = models_dir = os.path.expanduser(os.path.expandvars(self.save_dir))
			enc = encoder.get_encoder(self.model_name, self.save_dir)
			hparams = model.default_hparams()

			with open(os.path.join(self.save_dir, self.model_name, 'hparams.json')) as f:
				data = json.load(f)
				hparams.override_from_dict(data)

			length = hparams.n_ctx

			with tf.Session(graph=tf.Graph()) as sess:
				batch_size = 1
				temperature = 1
				top_k = 40

				context = tf.placeholder(tf.int32, [batch_size, None])
				np.random.seed(None)
				tf.set_random_seed(None)

				output = sample.sample_sequence(
					hparams=hparams,
					length=length,
					start_token=enc.encoder['<|endoftext|>'],
					batch_size=batch_size,
					temperature=temperature, 
					top_k=top_k
				)

				saver = tf.train.Saver()
				ckpt = tf.train.latest_checkpoint(os.path.join(self.save_dir, self.model_name))
				saver.restore(sess, ckpt)

				generated = 0
				text = None

				for _ in range(n_samples):
					prompt = input(colored('Enter a prompt got GPT-2 >> ', 'cyan'))
					print ('{}: {}\n'.format(colored('Prompt', attrs=['bold']), colored(prompt, 'green')))
					print (colored('Generating sample...', 'yellow'))

					context_tokens = enc.encode(prompt)
					generated = 0
					text = []
					for _ in range(n_samples // batch_size):
						out = sess.run(output, feed_dict={
							context: [context_tokens for _ in range(batch_size)]
						})[:, len(context_tokens):]

						for i in range(batch_size):
							generated += 1
							text.append(enc.decode(out[i]))
							print (colored('---------------------SAMPLE---------------------\n', 'cyan'))

							if display:
								print (text)

							if return_text:
								return text

	def finetune(self, dataset, return_text=True):
		"""  Returns generated text sample

		Parameters
		----------
		arg: dataset (object)
			- desc: Custom dataset text file

		arg: return_text (bool)
			- default: True
			- desc: Toggles whether to return custom-generated text in an array after fine-tuning

		Returns:
			Generated string in an array
		"""
		sess = gpt2.start_tf_sess()
		gpt2.finetune(sess,
              dataset,
              model_name=self.model_name,
              steps=1000)   # steps is max number of training steps

		if return_text:
			text = gpt2.generate(sess, return_as_list=True)
			return text
		else:
			gpt2.generate(sess)