import os
import datetime
from tqdm import tqdm
import numpy as np

import tensorflow as tf
print('TF version', tf.__version__)

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), 'Physical GPUs,', len(logical_gpus), 'Logical GPUs')
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)


@tf.function
def do_stuff(wmat, tf_var):
    with tf.device("/gpu:0"):
        S = tf.constant(0.0)
        for i in tf.range(4):
            fi = tf.cast(i, dtype=tf.float32)
            A = tf.math.lgamma(tf.tanh(tf.matmul(wmat + fi, tf.transpose(wmat - fi, [0, 2, 1]))))
            S += tf.reduce_sum(A)
        error = tf.reduce_mean(tf_var)
        return error, S

exp_uuid = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

n_batches = 512


def gen():
    for i in range(n_batches):
        with tf.device("/cpu:0"): # Make sure it comes from CPU
            r = tf.ones((400,800))
        yield r

option_names = ['map().batch()', 'map_and_batch()']
for option in range(2):

    with tf.device("/cpu:0"):
        dataset = tf.data.Dataset.from_generator(gen, output_types=tf.float32)

        def my_identity(x):
            with tf.device("/cpu:0"):
                print("my_identity input:", x, x.device)
                y = tf.identity(x)
                print("my_identity output:", y, y.device)
                return y

        if option == 0:
            ## Option 0: map().batch()
            dataset = dataset.map(my_identity).batch(16)

        elif option == 1:
            ## Option 1: deprecated map_and_batch()
            dataset = dataset.apply(tf.data.experimental.map_and_batch(my_identity, 16))

    gpu_transform = tf.data.experimental.prefetch_to_device('/gpu:0', buffer_size=4)
    dataset = dataset.apply(gpu_transform)


    tf_var = tf.Variable(tf.zeros(3))
    adam = tf.keras.optimizers.Adam(1e-4)
    logpath = os.path.join('data', 'logs', 'pa_' + exp_uuid + '_' + option_names[option])

    tf.profiler.experimental.start(logpath)
    start = datetime.datetime.now()
    for b, wmat in tqdm(enumerate(dataset)):
        with tf.GradientTape() as tape:

            if b == 0:
                print('\n\n dataset element device', wmat.device)
                print('\n')

            # Do some calculations
            result = do_stuff(wmat, tf_var)

        grads = tape.gradient(result[0], [tf_var])
        adam.apply_gradients(zip(grads, [tf_var]))
    stop = datetime.datetime.now()
    tf.profiler.experimental.stop()

    print(f'\n\nOption {option_names[option]}\n===========================\n')
    print(logpath)
    print('Time lapsed=', stop - start)
    print("\n\n")