import os
import tensorflow as tf


class Net(tf.keras.Model):
	"""A simple linear model."""
	def __init__(self):
		super(Net, self).__init__()
		self.l1 = tf.keras.layers.Dense(5)

	def call(self, x):
		return self.l1(x)


def toy_dataset():
	inputs = tf.range(10.)[:, None]
	labels = inputs * 5. + tf.range(5.)[None, :]
	dataset = tf.data.Dataset.from_tensor_slices(
		dict(x=inputs, y=labels)).repeat(1).batch(2)
	dataset = dataset.shuffle(buffer_size=4)
	dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
	return dataset

def train_step(net, x, y, optimizer):
	"""Trains `net` on `example` using `optimizer`."""
	with tf.GradientTape() as tape:
		output = net(x)
		loss = tf.reduce_mean(tf.abs(output - y))
	variables = net.trainable_variables
	gradients = tape.gradient(loss, variables)
	optimizer.apply_gradients(zip(gradients, variables))
	return loss


def generate_tf2_checkpoints():

	opt = tf.keras.optimizers.Adam(0.1)
	net = Net()
	dataset = toy_dataset()
	ckpt = tf.train.Checkpoint(optimizer=opt, net=net)
	manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)


	ckpt.restore(manager.latest_checkpoint)
	if manager.latest_checkpoint:
		step = opt.iterations.numpy() + 1
		print("Restored from {}".format(manager.latest_checkpoint))
	else:
		step = 0
		print("Initializing from scratch.")

	for _ in range(10):
		for element in dataset:
			loss = train_step(net, element['x'], element['y'], opt)
			step += 1
			if step % 10 == 0:
				save_path = manager.save()
				print("Saved checkpoint for step {}: {}".format(step, save_path))
				print("loss {:1.2f}".format(loss.numpy()))





if __name__ == '__main__':

	os.environ["CUDA_VISIBLE_DEVICES"]=""
	gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
	for gpu in gpus:
		tf.config.experimental.set_memory_growth(gpu, True)

	generate_tf2_checkpoints()