import tensorflow as tf
import numpy as np
from .modules.embedder import Embedder
from .modules.recovery import Recovery
from .modules.generator import Generator
from .modules.supervisor import Supervisor
from .modules.discriminator import Discriminator
from .modules.model_utils import extract_time, random_generator, batch_generator, MinMaxScaler, save_dict_to_json
import logging, os , datetime, time
#logging.disable(logging.WARNING) 


class TimeGAN(tf.keras.Model):
    def __init__(self, args):
        super().__init__()
        self.embedder = Embedder(args)
        self.recovery = Recovery(args)
        self.generator = Generator(args)
        self.supervisor = Supervisor(args)
        self.discriminator = Discriminator(args)
        self.mse = tf.keras.losses.MeanSquaredError()

    # E0_solver
    def recovery_forward(self, X, optimizer):
        # initial_hidden = self.embedder.initialize_hidden_state()
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            X_tilde = self.recovery(H, training=True)
            E_loss_T0 = self.mse(X, X_tilde)
            E_loss0 = 10 * tf.math.sqrt(E_loss_T0)
        
        var_list = self.embedder.trainable_weights + self.recovery.trainable_weights
        grads = tape.gradient(E_loss0, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return E_loss_T0

    # GS_solver
    def supervisor_forward(self, X, Z, optimizer):
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            #H_hat = self.generator(Z, training=True)
            H_hat_supervise = self.supervisor(H, training=True)
            G_loss_S = self.mse(H[:, 1:, :], H_hat_supervise[:, :-1, :])
        
        var_list = self.supervisor.trainable_weights
        grads = tape.gradient(G_loss_S, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return G_loss_S

    # G_solver
    def generator_forward(self, X, Z, optimizer=None, gamma=1):
        with tf.GradientTape() as tape:
            H = self.embedder(X)
            E_hat = self.generator(Z, training=True)

            # Supervisor & Recovery
            H_hat_supervise = self.supervisor(H, training=True)
            H_hat = self.supervisor(E_hat, training=True)
            X_hat = self.recovery(H_hat)

            # Discriminator
            Y_fake = self.discriminator(H_hat)
            Y_real = self.discriminator(H)
            Y_fake_e = self.discriminator(E_hat)

            # Generator loss
            G_loss_U = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_fake), Y_fake, from_logits=True)
            )
            G_loss_U_e = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_fake_e), Y_fake_e, from_logits=True)
            )
            G_loss_S = tf.math.reduce_mean(
                tf.keras.losses.mean_squared_error(H[:, 1:, :], H_hat_supervise[:, :-1, :])
            )
            # Difference in "variance" between X_hat and X
            G_loss_V1 = tf.math.reduce_mean(tf.math.abs(tf.math.sqrt(tf.nn.moments(X_hat, [0])[1] + 1e-6)
                                            - tf.math.sqrt(tf.nn.moments(X, [0])[1] + 1e-6)))
            # Difference in "mean" between X_hat and X
            G_loss_V2 = tf.math.reduce_mean(tf.math.abs((tf.nn.moments(X_hat, [0])[0])
                                            - (tf.nn.moments(X, [0])[0])))
            G_loss_V = G_loss_V1 + G_loss_V2
            #G_loss_V = tf.math.add(G_loss_V1, G_loss_V2)
            ## Sum of all G_losses
            G_loss = G_loss_U + gamma * G_loss_U_e + 100 * tf.math.sqrt(G_loss_S) + 100 * G_loss_V

        GS_var_list = self.generator.trainable_weights + self.supervisor.trainable_weights
        GS_grads = tape.gradient(G_loss, GS_var_list)
        optimizer.apply_gradients(zip(GS_grads, GS_var_list))
        
        return G_loss_U, G_loss_S, G_loss_V

    # D_solver
    def discriminator_forward(self, X, Z, optimizer=None, gamma=1, train_D=False):
        if train_D:
            with tf.GradientTape() as tape:
                H = self.embedder(X)
                E_hat = self.generator(Z)

                # Supervisor & Recovery
                H_hat_supervise = self.supervisor(H)
                H_hat = self.supervisor(E_hat)
                X_hat = self.recovery(H_hat)

                # Discriminator forward
                Y_fake = self.discriminator(H_hat)
                Y_real = self.discriminator(H)
                Y_fake_e = self.discriminator(E_hat)

                # Discriminator loss
                D_loss_real = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.ones_like(Y_real), Y_real, from_logits=True)
                )
                D_loss_fake = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake), Y_fake, from_logits=True)
                )
                D_loss_fake_e = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake_e), Y_fake_e, from_logits=True)
                )
                D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e

            D_var_list = self.discriminator.trainable_weights
            D_grads = tape.gradient(D_loss, D_var_list)
            optimizer.apply_gradients(zip(D_grads, D_var_list))

        else:
            # Checking if D_loss > 0.15
            H = self.embedder(X)
            E_hat = self.generator(Z)

            # Supervisor & Recovery
            H_hat_supervise = self.supervisor(H)
            H_hat = self.supervisor(E_hat)
            X_hat = self.recovery(H_hat)

            # Discriminator forward
            Y_fake = self.discriminator(H_hat)
            Y_real = self.discriminator(H)
            Y_fake_e = self.discriminator(E_hat)

            # Discriminator loss
            D_loss_real = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_real), Y_real, from_logits=True)
            )
            D_loss_fake = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake), Y_fake, from_logits=True)
            )
            D_loss_fake_e = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake_e), Y_fake_e, from_logits=True)
            )
            D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e
        

        return D_loss

    # E_solver
    def embedding_forward_joint(self, X, optimizer, eta):
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            X_tilde = self.recovery(H, training=True)
            E_loss_T0 = tf.math.reduce_mean(
                tf.keras.losses.mean_squared_error(X, X_tilde)
            )
            E_loss0 = 10 * tf.math.sqrt(E_loss_T0)

            H_hat_supervise = self.supervisor(H)
            G_loss_S = tf.math.reduce_mean(
                self.mse(H[:, 1:, :], H_hat_supervise[:, :-1, :])
            )
            E_loss = E_loss0 + eta * G_loss_S
        
        var_list = self.embedder.trainable_weights + self.recovery.trainable_weights
        grads = tape.gradient(E_loss, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return E_loss

    # Inference
    def generate(self, Z, ori_data_num, ori_time, max_val, min_val):
        """
        Args:
            Z: input random noises
            ori_data_num: the first dimension of ori_data.shape
            ori_time: timesteps of original data
            max_val: the maximum value of MinMaxScaler(ori_data)
            min_val: the minimum value of MinMaxScaler(ori_data)
        Return:
            generated_data: synthetic time-series data
        """
        E_hat = self.generator(Z)
        H_hat = self.supervisor(E_hat)
        generated_data_curr = self.recovery(H_hat)
        generated_data_curr = generated_data_curr.numpy()
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)
        
        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val
        
        return generated_data

    # Inference by autoencoder
    def ae_generate(self, ori_data, ori_data_num, ori_time, max_val, min_val):
        H = self.embedder(ori_data)
        generated_data_curr = self.recovery(H)
        generated_data_curr = generated_data_curr.numpy()
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)

        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val

        return generated_data

    # Direct inference (testing)
    def generator_inference(self, z_dim, ori_data, model_dir):
        """This function is currently not used.
        Args:
            Z: input random noises
            ori_data: the original dataset (for information extraction)
            trained_model: trained Generator
        Return:
            generated_data: synthetic time-series data
        """
        no, seq_len, dim = np.asarray(ori_data).shape
        ori_time, max_seq_len = extract_time(ori_data)
        # Normalization
        _, min_val, max_val = MinMaxScaler(ori_data)


        if z_dim == -1:  # choose z_dim for the dimension of noises
            z_dim = dim
        Z = random_generator(no, z_dim, ori_time, max_seq_len)

        # Load models
        self.recovery.load_weights(model_dir)
        self.supervisor.load_weights(model_dir)
        self.generator.load_weights(model_dir)

        E_hat = self.generator(Z)
        H_hat = self.supervisor(E_hat)
        generated_data_curr = self.recovery(H_hat)
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)
        
        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val
        
        return generated_data


def train_timegan(ori_data, dynamic_time, args):
    no, seq_len, dim = np.asarray(ori_data).shape

    if args.z_dim == -1:  # choose z_dim for the dimension of noises
        args.z_dim = dim

    ori_time, max_seq_len = extract_time(ori_data)
    # Normalization
    ori_data, min_val, max_val = MinMaxScaler(ori_data)

    model = TimeGAN(args)

    # Set up optimizers
    if args.optimizer == 'adam':
        optimizer = tf.keras.optimizers.Adam(epsilon=args.epsilon)

    print('Set up Tensorboard')
    current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    train_log_dir = os.path.join('tensorboard', current_time + '-' + args.exp_name)
    train_summary_writer = tf.summary.create_file_writer(train_log_dir)
    # save args to tensorboard folder
    save_dict_to_json(args.__dict__, os.path.join(train_log_dir, 'params.json'))

    # 1. Embedding network training
    print('Start Embedding Network Training')
    start = time.time()
    for itt in range(args.embedding_iterations):
    #for itt in range(1):
        X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size, use_tf_data=False)
        X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
        step_e_loss = model.recovery_forward(X_mb, optimizer)
        if itt % 100 == 0:
            print('step: '+ str(itt) + '/' + str(args.embedding_iterations) +
                    ', e_loss: ' + str(np.round(np.sqrt(step_e_loss),4)))
            # Write to Tensorboard
            with train_summary_writer.as_default():
                tf.summary.scalar('Embedding_loss', np.round(np.sqrt(step_e_loss),4), step=itt)

    print('Finish Embedding Network Training')
    end = time.time()
    print('Train embedding time elapsed: {} sec'.format(end-start))

    # 2. Training only with supervised loss
    if args.gen_type == 'autoencoder':
        print('Skip Training with Supervised Loss')
    else:
        print('Start Training with Supervised Loss Only')
        start = time.time()
        for itt in range(args.supervised_iterations):
        #for itt in range(1):
            X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
            Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)

            X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
            Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

            step_g_loss_s = model.supervisor_forward(X_mb, Z_mb, optimizer)
            if itt % 100 == 0:
                print('step: '+ str(itt)  + '/' + str(args.supervised_iterations) +', s_loss: '
                            + str(np.round(np.sqrt(step_g_loss_s),4)))
                # Write to Tensorboard
                with train_summary_writer.as_default():
                    tf.summary.scalar('Supervised_loss', np.round(np.sqrt(step_g_loss_s),4), step=itt)
        
        print('Finish Training with Supervised Loss Only')
        end = time.time()
        print('Train Supervisor time elapsed: {} sec'.format(end-start))

    # 3. Joint Training
    if args.gen_type == 'autoencoder':
        print('Skip Joint Training')
    else:
        print('Start Joint Training')
        start = time.time()
        for itt in range(args.joint_iterations):
        #for itt in range(1):
            # Generator training (two times as discriminator training)
            for g_more in range(2):
                X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
                Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)
                
                X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
                Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

                step_g_loss_u, step_g_loss_s, step_g_loss_v = model.generator_forward(X_mb,
                                                                                    Z_mb,
                                                                                    optimizer)
                step_e_loss_t0 = model.embedding_forward_joint(X_mb, optimizer, args.eta)

            # Discriminator training
            X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
            Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)

            X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
            Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

            check_d_loss = model.discriminator_forward(X_mb, Z_mb, train_D=False)
            if (check_d_loss > 0.15):
                step_d_loss = model.discriminator_forward(X_mb, Z_mb, optimizer, train_D=True)
            else:
                step_d_loss = check_d_loss

            if itt % 100 == 0:
                print('step: '+ str(itt) + '/' + str(args.joint_iterations) +
                    ', d_loss: ' + str(np.round(step_d_loss, 4)) +
                    ', g_loss_u: ' + str(np.round(step_g_loss_u, 4)) +
                    ', g_loss_s: ' + str(np.round(np.sqrt(step_g_loss_s), 4)) +
                    ', g_loss_v: ' + str(np.round(step_g_loss_v, 4)) +
                    ', e_loss_t0: ' + str(np.round(np.sqrt(step_e_loss_t0), 4)))
                # Write to Tensorboard
                with train_summary_writer.as_default():
                    tf.summary.scalar('Joint/Discriminator',
                                    np.round(step_d_loss, 4), step=itt)
                    tf.summary.scalar('Joint/Generator',
                                    np.round(step_g_loss_u, 4), step=itt)
                    tf.summary.scalar('Joint/Supervisor',
                                    np.round(step_g_loss_s, 4), step=itt)
                    tf.summary.scalar('Joint/Moments',
                                    np.round(step_g_loss_v, 4), step=itt)
                    tf.summary.scalar('Joint/Embedding',
                                    np.round(step_e_loss_t0, 4), step=itt)
        print('Finish Joint Training')
        end = time.time()
        print('Train jointly time elapsed: {} sec'.format(end-start))


    ## Synthetic data generation
    if args.gen_type == 'autoencoder':
        input_ori = tf.convert_to_tensor(ori_data, dtype=tf.float32)
        generated_data = model.ae_generate(input_ori, no, ori_time, max_val, min_val)
    else:    
        Z_mb = random_generator(no, args.z_dim, ori_time, args.max_seq_len)
        Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)
        generated_data = model.generate(Z_mb, no, ori_time, max_val, min_val)

    return generated_data, train_log_dir

import tensorflow as tf
import numpy as np
from .modules.embedder import Embedder
from .modules.recovery import Recovery
from .modules.generator import Generator
from .modules.supervisor import Supervisor
from .modules.discriminator import Discriminator
from .modules.model_utils import extract_time, random_generator, batch_generator, MinMaxScaler, save_dict_to_json
import logging, os , datetime, time
#logging.disable(logging.WARNING) 


class TimeGAN(tf.keras.Model):
    def __init__(self, args):
        super().__init__()
        self.embedder = Embedder(args)
        self.recovery = Recovery(args)
        self.generator = Generator(args)
        self.supervisor = Supervisor(args)
        self.discriminator = Discriminator(args)
        self.mse = tf.keras.losses.MeanSquaredError()

    # E0_solver
    def recovery_forward(self, X, optimizer):
        # initial_hidden = self.embedder.initialize_hidden_state()
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            X_tilde = self.recovery(H, training=True)
            E_loss_T0 = self.mse(X, X_tilde)
            E_loss0 = 10 * tf.math.sqrt(E_loss_T0)
        
        var_list = self.embedder.trainable_weights + self.recovery.trainable_weights
        grads = tape.gradient(E_loss0, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return E_loss_T0

    # GS_solver
    def supervisor_forward(self, X, Z, optimizer):
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            #H_hat = self.generator(Z, training=True)
            H_hat_supervise = self.supervisor(H, training=True)
            G_loss_S = self.mse(H[:, 1:, :], H_hat_supervise[:, :-1, :])
        
        var_list = self.supervisor.trainable_weights
        grads = tape.gradient(G_loss_S, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return G_loss_S

    # G_solver
    def generator_forward(self, X, Z, optimizer=None, gamma=1):
        with tf.GradientTape() as tape:
            H = self.embedder(X)
            E_hat = self.generator(Z, training=True)

            # Supervisor & Recovery
            H_hat_supervise = self.supervisor(H, training=True)
            H_hat = self.supervisor(E_hat, training=True)
            X_hat = self.recovery(H_hat)

            # Discriminator
            Y_fake = self.discriminator(H_hat)
            Y_real = self.discriminator(H)
            Y_fake_e = self.discriminator(E_hat)

            # Generator loss
            G_loss_U = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_fake), Y_fake, from_logits=True)
            )
            G_loss_U_e = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_fake_e), Y_fake_e, from_logits=True)
            )
            G_loss_S = tf.math.reduce_mean(
                tf.keras.losses.mean_squared_error(H[:, 1:, :], H_hat_supervise[:, :-1, :])
            )
            # Difference in "variance" between X_hat and X
            G_loss_V1 = tf.math.reduce_mean(tf.math.abs(tf.math.sqrt(tf.nn.moments(X_hat, [0])[1] + 1e-6)
                                            - tf.math.sqrt(tf.nn.moments(X, [0])[1] + 1e-6)))
            # Difference in "mean" between X_hat and X
            G_loss_V2 = tf.math.reduce_mean(tf.math.abs((tf.nn.moments(X_hat, [0])[0])
                                            - (tf.nn.moments(X, [0])[0])))
            G_loss_V = G_loss_V1 + G_loss_V2
            #G_loss_V = tf.math.add(G_loss_V1, G_loss_V2)
            ## Sum of all G_losses
            G_loss = G_loss_U + gamma * G_loss_U_e + 100 * tf.math.sqrt(G_loss_S) + 100 * G_loss_V

        GS_var_list = self.generator.trainable_weights + self.supervisor.trainable_weights
        GS_grads = tape.gradient(G_loss, GS_var_list)
        optimizer.apply_gradients(zip(GS_grads, GS_var_list))
        
        return G_loss_U, G_loss_S, G_loss_V

    # D_solver
    def discriminator_forward(self, X, Z, optimizer=None, gamma=1, train_D=False):
        if train_D:
            with tf.GradientTape() as tape:
                H = self.embedder(X)
                E_hat = self.generator(Z)

                # Supervisor & Recovery
                H_hat_supervise = self.supervisor(H)
                H_hat = self.supervisor(E_hat)
                X_hat = self.recovery(H_hat)

                # Discriminator forward
                Y_fake = self.discriminator(H_hat)
                Y_real = self.discriminator(H)
                Y_fake_e = self.discriminator(E_hat)

                # Discriminator loss
                D_loss_real = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.ones_like(Y_real), Y_real, from_logits=True)
                )
                D_loss_fake = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake), Y_fake, from_logits=True)
                )
                D_loss_fake_e = tf.math.reduce_mean(
                    tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake_e), Y_fake_e, from_logits=True)
                )
                D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e

            D_var_list = self.discriminator.trainable_weights
            D_grads = tape.gradient(D_loss, D_var_list)
            optimizer.apply_gradients(zip(D_grads, D_var_list))

        else:
            # Checking if D_loss > 0.15
            H = self.embedder(X)
            E_hat = self.generator(Z)

            # Supervisor & Recovery
            H_hat_supervise = self.supervisor(H)
            H_hat = self.supervisor(E_hat)
            X_hat = self.recovery(H_hat)

            # Discriminator forward
            Y_fake = self.discriminator(H_hat)
            Y_real = self.discriminator(H)
            Y_fake_e = self.discriminator(E_hat)

            # Discriminator loss
            D_loss_real = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.ones_like(Y_real), Y_real, from_logits=True)
            )
            D_loss_fake = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake), Y_fake, from_logits=True)
            )
            D_loss_fake_e = tf.math.reduce_mean(
                tf.keras.losses.binary_crossentropy(tf.zeros_like(Y_fake_e), Y_fake_e, from_logits=True)
            )
            D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e
        

        return D_loss

    # E_solver
    def embedding_forward_joint(self, X, optimizer, eta):
        with tf.GradientTape() as tape:
            H = self.embedder(X, training=True)
            X_tilde = self.recovery(H, training=True)
            E_loss_T0 = tf.math.reduce_mean(
                tf.keras.losses.mean_squared_error(X, X_tilde)
            )
            E_loss0 = 10 * tf.math.sqrt(E_loss_T0)

            H_hat_supervise = self.supervisor(H)
            G_loss_S = tf.math.reduce_mean(
                self.mse(H[:, 1:, :], H_hat_supervise[:, :-1, :])
            )
            E_loss = E_loss0 + eta * G_loss_S
        
        var_list = self.embedder.trainable_weights + self.recovery.trainable_weights
        grads = tape.gradient(E_loss, var_list)
        optimizer.apply_gradients(zip(grads, var_list))

        return E_loss

    # Inference
    def generate(self, Z, ori_data_num, ori_time, max_val, min_val):
        """
        Args:
            Z: input random noises
            ori_data_num: the first dimension of ori_data.shape
            ori_time: timesteps of original data
            max_val: the maximum value of MinMaxScaler(ori_data)
            min_val: the minimum value of MinMaxScaler(ori_data)
        Return:
            generated_data: synthetic time-series data
        """
        E_hat = self.generator(Z)
        H_hat = self.supervisor(E_hat)
        generated_data_curr = self.recovery(H_hat)
        generated_data_curr = generated_data_curr.numpy()
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)
        
        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val
        
        return generated_data

    # Inference by autoencoder
    def ae_generate(self, ori_data, ori_data_num, ori_time, max_val, min_val):
        H = self.embedder(ori_data)
        generated_data_curr = self.recovery(H)
        generated_data_curr = generated_data_curr.numpy()
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)

        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val

        return generated_data

    # Direct inference (testing)
    def generator_inference(self, z_dim, ori_data, model_dir):
        """This function is currently not used.
        Args:
            Z: input random noises
            ori_data: the original dataset (for information extraction)
            trained_model: trained Generator
        Return:
            generated_data: synthetic time-series data
        """
        no, seq_len, dim = np.asarray(ori_data).shape
        ori_time, max_seq_len = extract_time(ori_data)
        # Normalization
        _, min_val, max_val = MinMaxScaler(ori_data)


        if z_dim == -1:  # choose z_dim for the dimension of noises
            z_dim = dim
        Z = random_generator(no, z_dim, ori_time, max_seq_len)

        # Load models
        self.recovery.load_weights(model_dir)
        self.supervisor.load_weights(model_dir)
        self.generator.load_weights(model_dir)

        E_hat = self.generator(Z)
        H_hat = self.supervisor(E_hat)
        generated_data_curr = self.recovery(H_hat)
        generated_data = list()

        for i in range(ori_data_num):
            temp = generated_data_curr[i, :ori_time[i], :]
            generated_data.append(temp)
        
        # Renormalization
        generated_data = generated_data * max_val
        generated_data = generated_data + min_val
        
        return generated_data


def train_timegan(ori_data, dynamic_time, args):
    no, seq_len, dim = np.asarray(ori_data).shape

    if args.z_dim == -1:  # choose z_dim for the dimension of noises
        args.z_dim = dim

    ori_time, max_seq_len = extract_time(ori_data)
    # Normalization
    ori_data, min_val, max_val = MinMaxScaler(ori_data)

    model = TimeGAN(args)

    # Set up optimizers
    if args.optimizer == 'adam':
        optimizer = tf.keras.optimizers.Adam(epsilon=args.epsilon)

    print('Set up Tensorboard')
    current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    train_log_dir = os.path.join('tensorboard', current_time + '-' + args.exp_name)
    train_summary_writer = tf.summary.create_file_writer(train_log_dir)
    # save args to tensorboard folder
    save_dict_to_json(args.__dict__, os.path.join(train_log_dir, 'params.json'))

    # 1. Embedding network training
    print('Start Embedding Network Training')
    start = time.time()
    for itt in range(args.embedding_iterations):
    #for itt in range(1):
        X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size, use_tf_data=False)
        X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
        step_e_loss = model.recovery_forward(X_mb, optimizer)
        if itt % 100 == 0:
            print('step: '+ str(itt) + '/' + str(args.embedding_iterations) +
                    ', e_loss: ' + str(np.round(np.sqrt(step_e_loss),4)))
            # Write to Tensorboard
            with train_summary_writer.as_default():
                tf.summary.scalar('Embedding_loss', np.round(np.sqrt(step_e_loss),4), step=itt)

    print('Finish Embedding Network Training')
    end = time.time()
    print('Train embedding time elapsed: {} sec'.format(end-start))

    # 2. Training only with supervised loss
    if args.gen_type == 'autoencoder':
        print('Skip Training with Supervised Loss')
    else:
        print('Start Training with Supervised Loss Only')
        start = time.time()
        for itt in range(args.supervised_iterations):
        #for itt in range(1):
            X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
            Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)

            X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
            Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

            step_g_loss_s = model.supervisor_forward(X_mb, Z_mb, optimizer)
            if itt % 100 == 0:
                print('step: '+ str(itt)  + '/' + str(args.supervised_iterations) +', s_loss: '
                            + str(np.round(np.sqrt(step_g_loss_s),4)))
                # Write to Tensorboard
                with train_summary_writer.as_default():
                    tf.summary.scalar('Supervised_loss', np.round(np.sqrt(step_g_loss_s),4), step=itt)
        
        print('Finish Training with Supervised Loss Only')
        end = time.time()
        print('Train Supervisor time elapsed: {} sec'.format(end-start))

    # 3. Joint Training
    if args.gen_type == 'autoencoder':
        print('Skip Joint Training')
    else:
        print('Start Joint Training')
        start = time.time()
        for itt in range(args.joint_iterations):
        #for itt in range(1):
            # Generator training (two times as discriminator training)
            for g_more in range(2):
                X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
                Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)
                
                X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
                Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

                step_g_loss_u, step_g_loss_s, step_g_loss_v = model.generator_forward(X_mb,
                                                                                    Z_mb,
                                                                                    optimizer)
                step_e_loss_t0 = model.embedding_forward_joint(X_mb, optimizer, args.eta)

            # Discriminator training
            X_mb, T_mb = batch_generator(ori_data, ori_time, args.batch_size)
            Z_mb = random_generator(args.batch_size, args.z_dim, T_mb, args.max_seq_len)

            X_mb = tf.convert_to_tensor(X_mb, dtype=tf.float32)
            Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)

            check_d_loss = model.discriminator_forward(X_mb, Z_mb, train_D=False)
            if (check_d_loss > 0.15):
                step_d_loss = model.discriminator_forward(X_mb, Z_mb, optimizer, train_D=True)
            else:
                step_d_loss = check_d_loss

            if itt % 100 == 0:
                print('step: '+ str(itt) + '/' + str(args.joint_iterations) +
                    ', d_loss: ' + str(np.round(step_d_loss, 4)) +
                    ', g_loss_u: ' + str(np.round(step_g_loss_u, 4)) +
                    ', g_loss_s: ' + str(np.round(np.sqrt(step_g_loss_s), 4)) +
                    ', g_loss_v: ' + str(np.round(step_g_loss_v, 4)) +
                    ', e_loss_t0: ' + str(np.round(np.sqrt(step_e_loss_t0), 4)))
                # Write to Tensorboard
                with train_summary_writer.as_default():
                    tf.summary.scalar('Joint/Discriminator',
                                    np.round(step_d_loss, 4), step=itt)
                    tf.summary.scalar('Joint/Generator',
                                    np.round(step_g_loss_u, 4), step=itt)
                    tf.summary.scalar('Joint/Supervisor',
                                    np.round(step_g_loss_s, 4), step=itt)
                    tf.summary.scalar('Joint/Moments',
                                    np.round(step_g_loss_v, 4), step=itt)
                    tf.summary.scalar('Joint/Embedding',
                                    np.round(step_e_loss_t0, 4), step=itt)
        print('Finish Joint Training')
        end = time.time()
        print('Train jointly time elapsed: {} sec'.format(end-start))


    ## Synthetic data generation
    if args.gen_type == 'autoencoder':
        input_ori = tf.convert_to_tensor(ori_data, dtype=tf.float32)
        generated_data = model.ae_generate(input_ori, no, ori_time, max_val, min_val)
    else:    
        Z_mb = random_generator(no, args.z_dim, ori_time, args.max_seq_len)
        Z_mb = tf.convert_to_tensor(Z_mb, dtype=tf.float32)
        generated_data = model.generate(Z_mb, no, ori_time, max_val, min_val)

    return generated_data, train_log_dir

#logging.disable(logging.WARNING) # E0_solver # initial_hidden = self.embedder.initialize_hidden_state() # GS_solver #H_hat = self.generator(Z, training=True) # G_solver # Supervisor & Recovery # Discriminator # Generator loss # Difference in "variance" between X_hat and X # Difference in "mean" between X_hat and X #G_loss_V = tf.math.add(G_loss_V1, G_loss_V2) ## Sum of all G_losses # D_solver # Supervisor & Recovery # Discriminator forward # Discriminator loss # Checking if D_loss > 0.15 # Supervisor & Recovery # Discriminator forward # Discriminator loss # E_solver # Inference Args:             Z: input random noises             ori_data_num: the first dimension of ori_data.shape             ori_time: timesteps of original data             max_val: the maximum value of MinMaxScaler(ori_data)             min_val: the minimum value of MinMaxScaler(ori_data)         Return:             generated_data: synthetic time-series data # Renormalization # Inference by autoencoder # Renormalization # Direct inference (testing) This function is currently not used.         Args:             Z: input random noises             ori_data: the original dataset (for information extraction)             trained_model: trained Generator         Return:             generated_data: synthetic time-series data # Normalization # choose z_dim for the dimension of noises # Load models # Renormalization # choose z_dim for the dimension of noises # Normalization # Set up optimizers # save args to tensorboard folder # 1. Embedding network training #for itt in range(1): # Write to Tensorboard # 2. Training only with supervised loss #for itt in range(1): # Write to Tensorboard # 3. Joint Training #for itt in range(1): # Generator training (two times as discriminator training) # Discriminator training # Write to Tensorboard ## Synthetic data generation