import tensorflow as tf
from maestro.ai import nn

class Model(object):

	class DataSet(object):
		def __init__(self, data, target):
			self.data = data
			self.target = target

	def __init__(self, name, size):
		self.name = f'{name}_' + 'x'.join([str(x) for x in size])
		self.size = size

	def init(self, x, y, prediction, accuracy, error, optimizer, training, globalStep):
		self.x = x
		self.y = y
		self.prediction = prediction
		self.accuracy = accuracy
		self.error = error
		self.optimizer = optimizer
		self.training = training
		self.globalStep = globalStep

	def train(self, data, target, steps=10000, rate=0.001):
		# Run train
		pass

	def predict(self, data):
		# Run predictions
		pass

	def _run(self):
		# Run session
		pass

class NoteModel(Model):
	def __init__(self, size):
		super().__init__('NoteModel', size)

	def initModel(self, learningRate, gradientCliping=True):
		with tf.device('/cpu:0'):
			with tf.name_scope('input'):
				x = tf.placeholder(tf.float32, [None, self.size[0]], name='x')
				y = tf.placeholder(tf.float32, [None, self.size[-1]], name='y')

		with tf.device('/gpu:0'):
			# Todo: Calculate dropdown...
			pred = nn.logistic_regression(x, self.size)

			with tf.name_scope('error'):
				cost = -tf.reduce_sum(y * tf.log(pred), reduction_indices=1)
				with tf.name_scope('total'):
					error = tf.reduce_mean(cost, name='error')
			tf.summary.scalar('error', error)

			with tf.name_scope('accuracy'):
				accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)), tf.float32), name='accuracy')
			tf.summary.scalar('accuracy', accuracy)

			with tf.name_scope('train'):
				globalStep = tf.Variable(0, trainable=False, name='globalStep')
				optimizer = tf.train.AdamOptimizer(learning_rate=learningRate)

				if gradientCliping:
					gvs = optimizer.compute_gradients(error)
					capped_gvs = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in gvs]
					training = optimizer.apply_gradients(capped_gvs, name='training', global_step=globalStep)
				else:
					training = optimizer.minimize(error, name='training', global_step=globalStep)

		self.init(x, y, pred, accuracy, error, optimizer, training, globalStep)
