import tensorflow as tf
print(tf.__version__)
from debug import clr


mnist = tf.keras.datasets.mnist
(x_train, y_train), (_, _) = mnist.load_data()
x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_train = x_train / 255.0
y_train = y_train.astype('int32')

buffer_size = 5000
batch_size = 100
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(buffer_size)
ds = ds.batch(batch_size)

class My_Model(tf.keras.Model):
  def __init__(self):
    super(My_Model, self).__init__()
    self.flatten = tf.keras.layers.Flatten()
    self.fc1 = tf.keras.layers.Dense(10)

  def call(self, x, training=True):
    x = self.flatten(x)
    x = self.fc1(x)
    return x

def loss(logits, labels):
  return tf.reduce_mean(
      tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels))

step_counter = tf.Variable(0,  trainable = False, dtype=tf.int64)

learning_rate = clr.cyclic_learning_rate(step_counter,
                              learning_rate=0.01,
                              max_lr=0.1,
                              step_size=100,
                              gamma=0.99994,
                              mode='triangular',
                              name=None)
    

opt = tf.optimizers.Adam(learning_rate=learning_rate())

model = My_Model()

    
epochs = 2
for epoch in range(2):
  for batch, (images, labels) in enumerate(ds):
    with tf.GradientTape() as tape:
      logits = model(images, training=True)
      loss_value = loss(logits, labels)
      # print('loss', loss_value.numpy())
    grads = tape.gradient(loss_value, model.variables)
    opt.apply_gradients(zip(grads, model.trainable_variables))
    print('LR {}'.format(opt.learning_rate.numpy()))
    step_counter.assign_add(1)
    print('Step Counter  {}'.format(step_counter.numpy()))