import tensorflow as tf
import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.utils import shuffle
from tqdm import tqdm_notebook as tqdm

tf.compat.v1.enable_eager_execution()


class AbstractIrmGame:
    """ Abstract class for IRM games. """

    def __init__(self, models, optimizers, extra_grad, variable_phi, n_epochs, batch_size, termination_acc, warm_start):
        self.models = models  # List of models for all the environments
        self.optimizers = optimizers  # List of optimizers for all the environments
        self.extra_grad = extra_grad
        self.variable_phi = variable_phi

        self.n_epochs = n_epochs  # Number of epochs
        self.batch_size = batch_size  # Batch size for each gradient update
        self.termination_acc = termination_acc  # Threshold on accuracy below which we terminating
        self.warm_start = warm_start  # minimum number of steps we have to train before terminating

        self.n_env = len(models) if not variable_phi else len(models) - 1

        self.keras_criterion = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
        self.torch_criterion = torch.nn.CrossEntropyLoss()

        self.grads = []
        self.grad_norms = [[] for _ in range(self.n_env)]
        self.losses = [[] for _ in range(self.n_env)]
        self.train_accs = []
        self.test_accs = []
        self.env_train_accs = [[] for _ in range(self.n_env)]

    @staticmethod
    def to_array(data):
        raise NotImplementedError

    @staticmethod
    def to_tensor(data):
        raise NotImplementedError

    @staticmethod
    def zeros(shape):
        raise NotImplementedError

    def predict(self, x, shape, keep_grad_idx, as_array):
        raise NotImplementedError

    def loss(self, x, y, i_env):
        raise NotImplementedError

    def zero_grad(self):
        raise NotImplementedError

    def update_optimizer(self, i_env):
        raise NotImplementedError

    def evaluate(self, x, y):
        accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

        y_, env_preds = self.predict(x=x, shape=(y.shape[0], 2), keep_grad_idx=None, as_array=True)

        accuracy.update_state(y_true=y, y_pred=y_)
        acc = accuracy.result().numpy()

        env_accs = []
        for i_env in range(self.n_env):
            accuracy.reset_states()
            accuracy.update_state(y_true=y, y_pred=env_preds[i_env])
            env_accs.append(accuracy.result().numpy())

        return acc, env_accs

    def concatenate_train_data(self, data_tuple):
        x = data_tuple[0][0]  # Combined data from environments
        for i in range(1, self.n_env):
            x_c = data_tuple[i][0]
            x = np.concatenate((x, x_c), axis=0)

        y = data_tuple[0][1]  # Combined labels from environments
        for i in range(1, self.n_env):
            y_c = data_tuple[i][1]
            y = np.concatenate((y, y_c), axis=0)

        return x, y

    def fit(self, data_tuple_train, data_tuple_test, env_wise=False):
        x_train_all, y_train_all = self.concatenate_train_data(data_tuple_train)
        x_test_all, y_test_all = data_tuple_test[0], data_tuple_test[1]

        flag = False
        n_examples = data_tuple_train[0][0].shape[0]
        steps = 0

        for i_epoch in range(self.n_epochs):
            print("Epoch %i/%i..." % (i_epoch + 1, self.n_epochs))

            epoch_data = []
            for env in range(self.n_env):
                x_env = data_tuple_train[env][0]
                y_env = data_tuple_train[env][1]
                epoch_data.append(shuffle(x_env, y_env))

            count = 0
            for offset in tqdm(range(0, n_examples, self.batch_size)):
                end = offset + self.batch_size
                x_batches = []  # list to store batches for each environment
                y_batches = []  # list to store batches of labels for each environment
                self.grads = []  # list to store gradients
                countp = count % self.n_env  # countp decides the index of the model which trains in the current step

                self.zero_grad()
                for i_env in range(self.n_env):
                    x_batches.append(epoch_data[i_env][0][offset:end, :])
                    y_batches.append(epoch_data[i_env][1][offset:end, :])

                    grad, loss_value = self.loss(i_env=i_env,
                                                 x=x_batches[i_env],
                                                 y=y_batches[i_env])

                    self.grads.append(grad)
                    self.losses[i_env].append(loss_value)

                ###
                # Old
                # Update the environment whose turn it is to learn
                # self.update_optimizer(i_env=countp)
                ##
                # New
                for i_env in range(self.n_env):
                    self.update_optimizer(i_env)
                ###

                # Compute training accuracy
                train_acc, env_train_accs = self.evaluate(x=x_train_all, y=y_train_all)
                self.train_accs.append(train_acc)
                for i_env in range(self.n_env):
                    self.env_train_accs[i_env].append(env_train_accs[i_env])

                # Compute test accuracy
                test_acc, env_test_accs = self.evaluate(x=x_test_all, y=y_test_all)
                self.test_accs.append(test_acc)

                # for i_env in range(self.n_env):
                #     self.grad_norms[i_env].append(
                #         tf.linalg.global_norm(self.grads[i_env])
                #     )

                if steps >= self.warm_start and train_acc < self.termination_acc:
                    # Terminate after warm start and train acc touches threshold we dont want it to fall below
                    # flag = True
                    # print("Early termination.")
                    # break
                    pass

                count = count + 1
                steps = steps + 1

            self.plot(env_wise)

            if flag:
                break

        # print train and test accuracy
        print("Training accuracy: %.4f" % self.train_accs[-1])
        print("Testing accuracy: %.4f" % self.test_accs[-1])

    def plot(self, env_wise):
        fig, ax1 = plt.subplots(figsize=(10, 6))
        ax1.set_xlabel("Training steps")
        ax1.set_ylabel("Accuracy")
        ax1.plot(self.train_accs, label="train acc")
        ax1.plot(self.test_accs, label="test acc")

        if env_wise:
            for i_env in range(self.n_env):
                ax1.plot(self.env_train_accs[i_env], label="train acc - env %i" % (i_env + 1))

        plt.legend()
        plt.show()

        fig, ax1 = plt.subplots(figsize=(10, 6))
        ax1.set_xlabel("Training steps")
        ax1.set_ylabel("Loss")

        if env_wise:
            for i_env in range(self.n_env):
                ax1.plot(self.losses[i_env], label="loss - env %i" % (i_env + 1))

        plt.legend()
        plt.show()

        fig, ax1 = plt.subplots(figsize=(10, 6))
        ax1.set_xlabel("Training steps")
        ax1.set_ylabel("Gradient norms")

        if env_wise:
            for i_env in range(self.n_env):
                ax1.plot(self.grad_norms[i_env], label="gradient norm - env %i" % (i_env + 1))

        plt.legend()
        plt.show()

    def mean_plot(self, env_wise):
        n = len(self.train_accs)
        train_accs = [np.mean([self.train_accs[2*i], self.train_accs[2*i+1]]) for i in range(n//2)]
        test_accs = [np.mean([self.test_accs[2*i], self.test_accs[2*i+1]]) for i in range(n//2)]
        env_train_accs = [[np.mean([self.env_train_accs[i_env][2*i], self.env_train_accs[i_env][2*i + 1]])
                           for i in range(n // 2)] for i_env in range(self.n_env)]

        plt.figure(figsize=(10, 6))
        plt.xlabel("Training steps pairs")
        plt.ylabel("Mean accuracy")

        plt.plot(train_accs, label="mean train acc")
        plt.plot(test_accs, label="mean test acc")

        if env_wise:
            for i_env in range(self.n_env):
                plt.plot(env_train_accs[i_env], label="train - env %i" % (i_env + 1))

        plt.legend()
        plt.show()


class TensorflowIrmGame(AbstractIrmGame):
    @staticmethod
    def to_array(data):
        return data

    @staticmethod
    def to_tensor(data):
        return data

    @staticmethod
    def zeros(shape):
        return tf.zeros(shape, dtype=tf.float32)

    def predict(self, x, shape, keep_grad_idx, as_array):
        x = self.to_tensor(x)
        y = self.zeros(shape)
        env_preds = []

        for i_env in range(self.n_env):
            env_pred = self.models[i_env](x)
            y = y + (1./self.n_env) * env_pred
            env_preds.append(env_pred)

        if as_array:
            y = self.to_array(y)
            for i_env in range(self.n_env):
                env_preds[i_env] = self.to_array(env_preds[i_env])

        return y, env_preds

    def loss(self, x, y, i_env):
        with tf.GradientTape() as tape:
            y_, _ = self.predict(x=x, shape=(y.shape[0], 2), keep_grad_idx=None, as_array=False)
            loss_value = self.keras_criterion(y_true=y, y_pred=y_)

        return tape.gradient(loss_value, self.models[i_env].trainable_variables), loss_value

    def zero_grad(self):
        pass

    def update_optimizer(self, i_env):
        self.optimizers[i_env].apply_gradients(zip(self.grads[i_env], self.models[i_env].trainable_variables))


class PytorchIrmGame(AbstractIrmGame):
    @staticmethod
    def to_array(data):
        return data.data.numpy()

    @staticmethod
    def to_tensor(data):
        return torch.tensor(data, dtype=torch.float32)

    @staticmethod
    def zeros(shape):
        return torch.zeros(shape, dtype=torch.float32)

    def predict(self, x, shape, keep_grad_idx, as_array):
        x = self.to_tensor(x)
        y = self.zeros(shape)
        env_preds = []

        if not as_array:
            for i_env in range(self.n_env):
                if i_env == keep_grad_idx:
                    env_pred = self.models[i_env](x)
                else:
                    with torch.no_grad():
                        env_pred = self.models[i_env](x)

                y = y + (1. / self.n_env) * env_pred
                env_preds.append(env_pred.detach())

        else:
            with torch.no_grad():
                for i_env in range(self.n_env):
                    env_pred = self.models[i_env](x)

                    y = y + (1. / self.n_env) * env_pred
                    env_preds.append(env_pred)

        if as_array:
            y = self.to_array(y)
            for i_env in range(self.n_env):
                env_preds[i_env] = self.to_array(env_preds[i_env])

        return y, env_preds

    def loss(self, x, y, i_env):
        pred, _ = self.predict(x=x, shape=(y.shape[0], 2), keep_grad_idx=i_env, as_array=False)
        y = torch.tensor(y).squeeze()

        loss = self.torch_criterion(target=y, input=pred)
        loss.backward()

        return None, loss.item()

    def zero_grad(self):
        for i_env in range(self.n_env):
            self.optimizers[i_env].zero_grad()

    def update_optimizer(self, i_env):
        self.optimizers[i_env].step()


class fixed_irm_game_model:
    def __init__(self, model_list, learning_rate, num_epochs, batch_size, termination_acc, warm_start, beta_1=0.9):
        self.model_list = model_list  # list of models for all the environments
        self.num_epochs = num_epochs  # number of epochs
        self.batch_size = batch_size  # batch size for each gradient update
        self.termination_acc = termination_acc  # threshold on accuracy below which we terminating
        self.warm_start = warm_start  # minimum number of steps we have to train before terminating due to accuracy
        # falling below threshold
        self.learning_rate = learning_rate  # learning rate in adam
        self.beta_1 = beta_1

    def fit(self, data_tuple_list):
        n_e = len(data_tuple_list)  # number of environments
        # combine the data from the different environments x_in: combined data from environments, y_in: combined
        # labels from environments, e_in: combined environment indices from environments
        x_in = data_tuple_list[0][0]
        for i in range(1, n_e):
            x_c = data_tuple_list[i][0]
            x_in = np.concatenate((x_in, x_c), axis=0)
        y_in = data_tuple_list[0][1]
        for i in range(1, n_e):
            y_c = data_tuple_list[i][1]
            y_in = np.concatenate((y_in, y_c), axis=0)
        e_in = data_tuple_list[0][2]
        for i in range(1, n_e):
            e_c = data_tuple_list[i][2]
            e_in = np.concatenate((e_in, e_c), axis=0)

            # cross entropy loss

        def loss_comb(model_list, x, y):
            loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            n_e = len(model_list)
            y_ = tf.zeros_like(y, dtype=tf.float32)
            # predict the model output from the ensemble
            for i in range(n_e):
                model_i = model_list[i]
                y_ = y_ + 0.5 * model_i(x)

            return loss_object(y_true=y, y_pred=y_)

        # gradient of cross entropy loss for environment e
        def grad_comb(model_list, inputs, targets, e):
            with tf.GradientTape() as tape:
                loss_value = loss_comb(model_list, inputs, targets)
            return loss_value, tape.gradient(loss_value, model_list[e].trainable_variables)

        model_list = self.model_list
        learning_rate = self.learning_rate
        beta_1 = self.beta_1

        # initialize optimizer for all the environments and representation learner and store it in a list
        optimizer_list = []
        for e in range(n_e):
            optimizer_list.append(tf.keras.optimizers.Adam(learning_rate=learning_rate, beta_1=beta_1))

        ####### train
        train_accuracy_results_0 = []  # list to store training accuracy

        flag = False
        num_epochs = self.num_epochs
        batch_size = self.batch_size
        num_examples = data_tuple_list[0][0].shape[0]
        period = n_e
        termination_acc = self.termination_acc
        warm_start = self.warm_start
        steps = 0
        for epoch in range(num_epochs):
            print("Epoch: " + str(epoch))
            datat_list = []
            for e in range(n_e):
                x_e = data_tuple_list[e][0]
                y_e = data_tuple_list[e][1]
                datat_list.append(shuffle(x_e, y_e))
            count = 0
            for offset in tqdm(range(0, num_examples, batch_size)):
                end = offset + batch_size
                batch_x_list = []  # list to store batches for each environment
                batch_y_list = []  # list to store batches of labels for each environment
                loss_value_list = []  # list to store loss values
                grads_list = []  # list to store gradients
                countp = count % period  # countp decides the index of the model which trains in the current step
                for e in range(n_e):
                    batch_x_list.append(datat_list[e][0][offset:end, :])
                    batch_y_list.append(datat_list[e][1][offset:end, :])
                    loss_value, grads = grad_comb(model_list, batch_x_list[e], batch_y_list[e], e)
                    grads_list.append(grads)
                # update the environment whose turn it is to learn
                optimizer_list[countp].apply_gradients(zip(grads_list[countp], model_list[countp].trainable_variables))

                # computing training accuracy
                y_ = tf.zeros_like(y_in, dtype=tf.float32)
                for e in range(n_e):
                    y_ = y_ + model_list[e](x_in)
                epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
                acc_train = np.float(epoch_accuracy(y_in, y_))
                train_accuracy_results_0.append(acc_train)

                if steps >= warm_start and acc_train < termination_acc:  ## Terminate after warm start and train
                    # acc touches threshold we dont want it to fall below
                    flag = True
                    print("Early termination.")
                    break

                count = count + 1
                steps = steps + 1
                self.train_accuracy_results = train_accuracy_results_0

            if flag:
                break

        self.model_list = model_list

        self.x_in = x_in
        self.y_in = y_in

    def evaluate(self, data_tuple_test):
        ##### evaluations jmtd
        x_test = data_tuple_test[0]
        y_test = data_tuple_test[1]
        x_in = self.x_in
        y_in = self.y_in

        model_list = self.model_list
        n_e = len(model_list)
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

        ytr_ = tf.zeros_like(y_in, dtype=tf.float32)
        for e in range(n_e):
            ytr_ = ytr_ + model_list[e](x_in)
        train_acc = np.float(train_accuracy(y_in, ytr_))

        yts_ = tf.zeros_like(y_test, dtype=tf.float32)
        for e in range(n_e):
            yts_ = yts_ + model_list[e](x_test)

        test_acc = np.float(test_accuracy(y_test, yts_))

        self.train_acc = train_acc
        self.test_acc = test_acc


class no_oscillation_irm_game_model:
    def __init__(self, model_list, learning_rate, num_epochs, batch_size, termination_acc, warm_start):

        self.model_list = model_list  # list of models for all the environments
        self.num_epochs = num_epochs  # number of epochs
        self.batch_size = batch_size  # batch size for each gradient update
        self.termination_acc = termination_acc  # threshold on accuracy below which we terminating
        self.warm_start = warm_start  # minimum number of steps we have to train before terminating due to accuracy
        # falling below threshold
        self.learning_rate = learning_rate  # learning rate in adam

    def fit(self, data_tuple_list):
        n_e = len(data_tuple_list)  # number of environments
        # combine the data from the different environments x_in: combined data from environments, y_in: combined
        # labels from environments, e_in: combined environment indices from environments
        x_in = data_tuple_list[0][0]
        for i in range(1, n_e):
            x_c = data_tuple_list[i][0]
            x_in = np.concatenate((x_in, x_c), axis=0)
        y_in = data_tuple_list[0][1]
        for i in range(1, n_e):
            y_c = data_tuple_list[i][1]
            y_in = np.concatenate((y_in, y_c), axis=0)
        e_in = data_tuple_list[0][2]
        for i in range(1, n_e):
            e_c = data_tuple_list[i][2]
            e_in = np.concatenate((e_in, e_c), axis=0)

            # cross entropy loss

        def loss_comb(model_list, x, y):
            loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            n_e = len(model_list)
            y_ = tf.zeros_like(y, dtype=tf.float32)
            # predict the model output from the ensemble
            for i in range(n_e):
                model_i = model_list[i]
                y_ = y_ + 0.5 * model_i(x)

            return loss_object(y_true=y, y_pred=y_)

        # gradient of cross entropy loss for environment e
        def grad_comb(model_list, inputs, targets, e):
            with tf.GradientTape() as tape:
                loss_value = loss_comb(model_list, inputs, targets)
            return loss_value, tape.gradient(loss_value, model_list[e].trainable_variables)

        model_list = self.model_list
        learning_rate = self.learning_rate

        # initialize optimizer for all the environments and representation learner and store it in a list
        optimizer_list = []
        for e in range(n_e):
            optimizer_list.append(tf.keras.optimizers.Adam(learning_rate=learning_rate))

        ####### train

        train_accuracy_results_0 = []  # list to store training accuracy

        flag = 'false'
        num_epochs = self.num_epochs
        batch_size = self.batch_size
        num_examples = data_tuple_list[0][0].shape[0]
        period = n_e
        termination_acc = self.termination_acc
        warm_start = self.warm_start
        steps = 0
        for epoch in range(num_epochs):
            print("Epoch: " + str(epoch))
            datat_list = []
            for e in range(n_e):
                x_e = data_tuple_list[e][0]
                y_e = data_tuple_list[e][1]
                datat_list.append(shuffle(x_e, y_e))
            count = 0
            for offset in range(0, num_examples, batch_size):
                end = offset + batch_size
                batch_x_list = []  # list to store batches for each environment
                batch_y_list = []  # list to store batches of labels for each environment
                loss_value_list = []  # list to store loss values
                grads_list = []  # list to store gradients
                countp = count % period  # countp decides the index of the model which trains in the current step
                for e in range(n_e):
                    batch_x_list.append(datat_list[e][0][offset:end, :])
                    batch_y_list.append(datat_list[e][1][offset:end, :])
                    loss_value, grads = grad_comb(model_list, batch_x_list[e], batch_y_list[e], e)
                    grads_list.append(grads)
                # update the environment whose turn it is to learn
                optimizer_list[countp].apply_gradients(zip(grads_list[countp], model_list[countp].trainable_variables))

                # computing training accuracy
                y_ = tf.zeros_like(y_in, dtype=tf.float32)
                for e in range(n_e):
                    y_ = y_ + model_list[e](x_in)
                epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
                acc_train = np.float(epoch_accuracy(y_in, y_))
                train_accuracy_results_0.append(acc_train)

                if (
                        steps >= warm_start and acc_train < termination_acc):  ## Terminate after warm start and train
                    # acc touches threshold we dont want it to fall below
                    flag = 'true'
                    break

                count = count + 1
                steps = steps + 1
                self.train_accuracy_results = train_accuracy_results_0
            if (flag == 'true'):
                break
        self.model_list = model_list

        self.x_in = x_in
        self.y_in = y_in

    def evaluate(self, data_tuple_test):
        ##### evaluations jmtd
        x_test = data_tuple_test[0]
        y_test = data_tuple_test[1]
        x_in = self.x_in
        y_in = self.y_in

        model_list = self.model_list
        n_e = len(model_list)
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

        ytr_ = tf.zeros_like(y_in, dtype=tf.float32)
        for e in range(n_e):
            ytr_ = ytr_ + model_list[e](x_in)
        train_acc = np.float(train_accuracy(y_in, ytr_))

        yts_ = tf.zeros_like(y_test, dtype=tf.float32)
        for e in range(n_e):
            yts_ = yts_ + model_list[e](x_test)

        test_acc = np.float(test_accuracy(y_test, yts_))

        self.train_acc = train_acc
        self.test_acc = test_acc


class variable_irm_game_model:
    def __init__(self, model_list, learning_rate, num_epochs, batch_size, termination_acc, warm_start):
        self.model_list = model_list  # list of models for the environments and representation learner
        self.num_epochs = num_epochs  # number of epochs
        self.batch_size = batch_size  # batch size for each gradient update
        self.termination_acc = termination_acc  # threshold on accuracy below which we terminate
        self.warm_start = warm_start  # minimum number of steps before terminating
        self.learning_rate = learning_rate  # learning rate for Adam optimizer

    def fit(self, data_tuple_list):
        n_e = len(data_tuple_list)  # number of environments
        # combine the data from the different environments x_in: combined data (features) from environments, y_in:
        # combined labels from environments, e_in: combined environment indices from environments
        x_in = data_tuple_list[0][0]
        for i in range(1, n_e):
            x_c = data_tuple_list[i][0]
            x_in = np.concatenate((x_in, x_c), axis=0)
        y_in = data_tuple_list[0][1]
        for i in range(1, n_e):
            y_c = data_tuple_list[i][1]
            y_in = np.concatenate((y_in, y_c), axis=0)
        e_in = data_tuple_list[0][2]
        for i in range(1, n_e):
            e_c = data_tuple_list[i][2]
            e_in = np.concatenate((e_in, e_c), axis=0)

            # cross entropy loss

        def loss_comb(model_list, x, y):
            loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            n_e = len(model_list) - 1
            y_ = tf.zeros_like(y, dtype=tf.float32)
            # pass the data from the representation learner
            z = model_list[n_e](x)
            # pass the output from the representation learner into the environments and aggregate them 
            for i in range(n_e):
                model_i = model_list[i]
                y_ = y_ + 0.5 * model_i(z)

            return loss_object(y_true=y, y_pred=y_)

        # gradient of cross entropy loss for environment e
        def grad_comb(model_list, inputs, targets, e):
            with tf.GradientTape() as tape:
                loss_value = loss_comb(model_list, inputs, targets)
            return loss_value, tape.gradient(loss_value, model_list[e].trainable_variables)

        model_list = self.model_list
        learning_rate = self.learning_rate

        # initialize optimizer for all the environments and representation learner and store it in a list
        optimizer_list = []
        for e in range(n_e + 1):
            if (e <= n_e - 1):
                optimizer_list.append(tf.keras.optimizers.Adam(learning_rate=learning_rate))
            if (e == n_e):
                optimizer_list.append(tf.keras.optimizers.Adam(learning_rate=learning_rate * 0.1))

        ####### train

        train_accuracy_results_0 = []  # list to store training accuracy

        flag = 'false'
        num_epochs = self.num_epochs
        batch_size = self.batch_size
        num_examples = data_tuple_list[0][0].shape[0]
        period = n_e + 1  #
        termination_acc = self.termination_acc
        warm_start = self.warm_start
        steps = 0
        for epoch in range(num_epochs):
            print("Epoch: " + str(epoch))
            datat_list = []
            for e in range(n_e + 1):
                if (e <= n_e - 1):
                    x_e = data_tuple_list[e][0]
                    y_e = data_tuple_list[e][1]
                    datat_list.append(shuffle(x_e, y_e))
                if (e == n_e):
                    datat_list.append(shuffle(x_in, y_in))
            count = 0
            for offset in range(0, num_examples, batch_size):
                end = offset + batch_size
                batch_x_list = []  # list to store batches for each environment
                batch_y_list = []  # list to store batches of labels for each environment
                loss_value_list = []  # list to store loss values
                grads_list = []  # list to store gradients
                countp = period - 1 - (
                            count % period)  # countp decides the index of the model which trains in the current step
                for e in range(n_e + 1):
                    batch_x_list.append(datat_list[e][0][offset:end, :])
                    batch_y_list.append(datat_list[e][1][offset:end, :])
                    loss_value, grads = grad_comb(model_list, batch_x_list[e], batch_y_list[e], e)
                    grads_list.append(grads)

                # update either a representation learner or an environment model  
                optimizer_list[countp].apply_gradients(zip(grads_list[countp], model_list[countp].trainable_variables))

                # computing training accuracy
                y_ = tf.zeros_like(y_in, dtype=tf.float32)
                z_in = model_list[n_e](x_in)
                for e in range(n_e):
                    y_ = y_ + model_list[e](z_in)
                epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
                acc_train = np.float(epoch_accuracy(y_in, y_))
                train_accuracy_results_0.append(acc_train)

                if (steps >= warm_start and acc_train < termination_acc):  # Terminate after warm start and train
                    # accuracy touches threshold we dont want it to fall below
                    flag = 'true'
                    break

                count = count + 1
                steps = steps + 1
                self.train_accuracy_results = train_accuracy_results_0

            if (flag == 'true'):
                break

        self.model_list = model_list

        self.x_in = x_in
        self.y_in = y_in

    def evaluate(self, data_tuple_test):
        ##### evaluations jmtd
        x_test = data_tuple_test[0]
        y_test = data_tuple_test[1]
        x_in = self.x_in
        y_in = self.y_in

        model_list = self.model_list
        n_e = len(model_list) - 1
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

        # compute training accuracy
        ytr_ = tf.zeros_like(y_in, dtype=tf.float32)
        z_in = model_list[n_e](x_in)
        for e in range(n_e):
            ytr_ = ytr_ + model_list[e](z_in)
        train_acc = np.float(train_accuracy(y_in, ytr_))

        # compute testing accuracy
        z_test = model_list[n_e](x_test)
        yts_ = tf.zeros_like(y_test, dtype=tf.float32)
        for e in range(n_e):
            yts_ = yts_ + model_list[e](z_test)

        test_acc = np.float(test_accuracy(y_test, yts_))

        self.train_acc = train_acc
        self.test_acc = test_acc


class standard_erm_model:
    def __init__(self, model, num_epochs, batch_size, learning_rate):

        self.model = model
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate

    def fit(self, data_tuple_list):
        learning_rate = self.learning_rate
        num_epochs = self.num_epochs
        n_e = len(data_tuple_list)
        x_in = data_tuple_list[0][0]
        for i in range(1, n_e):
            x_c = data_tuple_list[i][0]
            x_in = np.concatenate((x_in, x_c), axis=0)
        y_in = data_tuple_list[0][1]
        for i in range(1, n_e):
            y_c = data_tuple_list[i][1]
            y_in = np.concatenate((y_in, y_c), axis=0)
        e_in = data_tuple_list[0][2]
        for i in range(1, n_e):
            e_c = data_tuple_list[i][2]
            e_in = np.concatenate((e_in, e_c), axis=0)

            ### fit the model
        model = self.model
        batch_size = self.batch_size

        model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
                      loss='sparse_categorical_crossentropy',
                      metrics=['accuracy'])

        model.fit(x_in, y_in, epochs=num_epochs, batch_size=batch_size)

        self.x_in = x_in
        self.y_in = y_in

    def evaluate(self, data_tuple_test):
        ##### evaluations jmtd
        x_test = data_tuple_test[0]
        y_test = data_tuple_test[1]
        x_in = self.x_in
        y_in = self.y_in

        model = self.model
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()

        ytr_ = model.predict(x_in)
        train_acc = np.float(train_accuracy(y_in, ytr_))

        yts_ = model.predict(x_test)

        test_acc = np.float(test_accuracy(y_test, yts_))

        self.train_acc = train_acc
        self.test_acc = test_acc


class irm_model:
    def __init__(self, model, learning_rate, batch_size, steps_max, steps_threshold, gamma_new):
        self.model = model  # initialized model passed
        self.learning_rate = learning_rate  # learning rate for Adam optimizer
        self.batch_size = batch_size  # batch size per gradient update
        self.steps_max = steps_max  # maximum number of gradient steps
        self.steps_threshold = steps_threshold  # threshold on the number of steps after which we use penalty gamma_new
        self.gamma_new = gamma_new  # penalty value; note penalty is set to 1 initially and gamma_new only kicks in
        # after steps exceeed steps_threshold

    def fit(self, data_tuple_list):

        n_e = len(data_tuple_list)  # number of environments
        # combine the data from the different environments, x_in: combined data (features) from different environments
        x_in = data_tuple_list[0][0]
        for i in range(1, n_e):
            x_c = data_tuple_list[i][0]
            x_in = np.concatenate((x_in, x_c), axis=0)
        y_in = data_tuple_list[0][1]
        for i in range(1, n_e):
            y_c = data_tuple_list[i][1]
            y_in = np.concatenate((y_in, y_c), axis=0)
        e_in = data_tuple_list[0][2]
        for i in range(1, n_e):
            e_c = data_tuple_list[i][2]
            e_in = np.concatenate((e_in, e_c), axis=0)

        self.x_in = x_in
        self.y_in = y_in

        # cross entropy (we do not use the cross entropy from keras because there are issues when computing gradient
        # of the gradient)
        def cross_entropy_manual(y, y_pred):
            y_p = tf.math.log(tf.nn.softmax(y_pred))
            n_p = np.float(tf.shape(y_p)[0])
            ind_0 = tf.where(y == 0)[:, 0]
            ind_1 = tf.where(y == 1)[:, 0]
            y_p0 = tf.gather(y_p, ind_0)[:, 0]
            y_p1 = tf.gather(y_p, ind_1)[:, 1]
            ent_0 = tf.reduce_sum(y_p0)
            ent_1 = tf.reduce_sum(y_p1)
            total = -(ent_0 + ent_1) / n_p
            return total

        # cross entropy loss for environment e
        def loss_n(model, x, e, y, w, k):
            index = np.where(e == k)
            y1_ = model(x[index[0]]) * w
            y1 = y[index[0]]

            return cross_entropy_manual(y1, y1_)

            # gradient of cross entropy loss w.r.t w for environment e

        def grad_norm_n(model, x, e, y, w, k):
            with tf.GradientTape() as g:
                g.watch(w)
                loss_value = loss_n(model, x, e, y, w, k)
            return g.gradient(loss_value, w) ** 2

        # total cross entropy loss across all environments    
        def loss_0(model, x, e, y, w):
            y_ = model(x)
            loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
            return loss_object(y_true=y, y_pred=y_)

        # sum of cross entropy loss and penalty 
        def loss_total(model, x, e, y, w, gamma, n_e):
            loss0 = loss_0(model, x, e, y, w)
            loss_penalty = 0.0
            for k in range(n_e):
                loss_penalty += gamma * grad_norm_n(model, x, e, y, w, k)

            return (loss0 + loss_penalty) * (1 / gamma)

            # gradient of sum of cross entropy loss and penalty w.r.t model parameters

        def grad_total_n(model, x, e, y, w, gamma, n_e):
            with tf.GradientTape() as tape:
                loss_value = loss_total(model, x, e, y, w, gamma, n_e)
            return loss_value, tape.gradient(loss_value, model.trainable_variables)

        model = self.model
        learning_rate = self.learning_rate
        optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

        ## train 
        train_loss_results = []
        train_accuracy_results = []
        flag = 'false'
        batch_size = self.batch_size
        num_examples = x_in.shape[0]
        gamma = 1.0
        w = tf.constant(1.0)
        steps = 0
        steps_max = self.steps_max
        steps_threshold = self.steps_threshold
        gamma_new = self.gamma_new
        while (steps <= steps_max):
            (xt, yt, et) = shuffle(x_in, y_in, e_in)
            epoch_loss_avg = tf.keras.metrics.Mean()
            epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
            count = 0
            if (steps >= steps_threshold):
                gamma = gamma_new
            for offset in range(0, num_examples, batch_size):
                end = offset + batch_size
                batch_x, batch_y, batch_e = xt[offset:end, :], yt[offset:end, :], et[offset:end, :]
                loss_values, grads = grad_total_n(model, batch_x, batch_e, batch_y, w, gamma, n_e)
                optimizer.apply_gradients(zip(grads, model.trainable_variables))
                epoch_loss_avg(loss_values)
                epoch_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
                acc_train = np.float(epoch_accuracy(y_in, model(x_in)))
                train_loss_results.append(epoch_loss_avg.result())
                train_accuracy_results.append(epoch_accuracy.result())
                count = count + 1
                steps = steps + 1

    def evaluate(self, data_tuple_test):
        x_test = data_tuple_test[0]
        y_test = data_tuple_test[1]
        x_in = self.x_in
        y_in = self.y_in
        train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
        model = self.model
        ytr_ = model.predict(x_in)
        train_acc = np.float(train_accuracy(y_in, ytr_))

        yts_ = model.predict(x_test)
        test_acc = np.float(test_accuracy(y_test, yts_))

        self.train_acc = train_acc
        self.test_acc = test_acc

matplotlib.pyplot.legend tensorflow.zeros torch.zeros tensorflow.reduce_sum tensorflow.compat.v1.enable_eager_execution matplotlib.pyplot.plot numpy.concatenate numpy.mean torch.no_grad tensorflow.where numpy.where torch.nn.CrossEntropyLoss torch.tensor tensorflow.gather tensorflow.keras.metrics.Mean matplotlib.pyplot.figure tensorflow.shape tensorflow.zeros_like matplotlib.pyplot.show matplotlib.pyplot.ylabel tensorflow.GradientTape tensorflow.nn.softmax tensorflow.constant tensorflow.keras.metrics.SparseCategoricalAccuracy tensorflow.keras.losses.SparseCategoricalCrossentropy sklearn.utils.shuffle matplotlib.pyplot.subplots tensorflow.keras.optimizers.Adam matplotlib.pyplot.xlabel
[(8, 'tensorflow.compat.v1.enable_eager_execution', 'tf.compat.v1.enable_eager_execution', ([], {}), True, 'import tensorflow as tf\n'), (27, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (28, 'torch.nn.CrossEntropyLoss', 'torch.nn.CrossEntropyLoss', ([], {}), False, 'import torch\n'), (62, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (172, 'matplotlib.pyplot.subplots', 'plt.subplots', ([], {'figsize': '(10, 6)'}), True, 'import matplotlib.pyplot as plt\n'), (182, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (183, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (185, 'matplotlib.pyplot.subplots', 'plt.subplots', ([], {'figsize': '(10, 6)'}), True, 'import matplotlib.pyplot as plt\n'), (193, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (194, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (196, 'matplotlib.pyplot.subplots', 'plt.subplots', ([], {'figsize': '(10, 6)'}), True, 'import matplotlib.pyplot as plt\n'), (204, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (205, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (214, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(10, 6)'}), True, 'import matplotlib.pyplot as plt\n'), (215, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Training steps pairs"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (216, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Mean accuracy"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (218, 'matplotlib.pyplot.plot', 'plt.plot', (['train_accs'], {'label': '"""mean train acc"""'}), True, 'import matplotlib.pyplot as plt\n'), (219, 'matplotlib.pyplot.plot', 'plt.plot', (['test_accs'], {'label': '"""mean test acc"""'}), True, 'import matplotlib.pyplot as plt\n'), (225, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (226, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (240, 'tensorflow.zeros', 'tf.zeros', (['shape'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (280, 'torch.tensor', 'torch.tensor', (['data'], {'dtype': 'torch.float32'}), False, 'import torch\n'), (284, 'torch.zeros', 'torch.zeros', (['shape'], {'dtype': 'torch.float32'}), False, 'import torch\n'), (459, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (460, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (462, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (467, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_test'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (600, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (601, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (603, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (608, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_test'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (751, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (752, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (755, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (763, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_test'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (819, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (820, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (917, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), True, 'import tensorflow as tf\n'), (926, 'tensorflow.constant', 'tf.constant', (['(1.0)'], {}), True, 'import tensorflow as tf\n'), (956, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (957, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (81, 'numpy.concatenate', 'np.concatenate', (['(x, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (86, 'numpy.concatenate', 'np.concatenate', (['(y, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (209, 'numpy.mean', 'np.mean', (['[self.train_accs[2 * i], self.train_accs[2 * i + 1]]'], {}), True, 'import numpy as np\n'), (210, 'numpy.mean', 'np.mean', (['[self.test_accs[2 * i], self.test_accs[2 * i + 1]]'], {}), True, 'import numpy as np\n'), (260, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (352, 'numpy.concatenate', 'np.concatenate', (['(x_in, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (356, 'numpy.concatenate', 'np.concatenate', (['(y_in, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (360, 'numpy.concatenate', 'np.concatenate', (['(e_in, e_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (365, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (367, 'tensorflow.zeros_like', 'tf.zeros_like', (['y'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (495, 'numpy.concatenate', 'np.concatenate', (['(x_in, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (499, 'numpy.concatenate', 'np.concatenate', (['(y_in, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (503, 'numpy.concatenate', 'np.concatenate', (['(e_in, e_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (508, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (510, 'tensorflow.zeros_like', 'tf.zeros_like', (['y'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (634, 'numpy.concatenate', 'np.concatenate', (['(x_in, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (638, 'numpy.concatenate', 'np.concatenate', (['(y_in, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (642, 'numpy.concatenate', 'np.concatenate', (['(e_in, e_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (647, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (649, 'tensorflow.zeros_like', 'tf.zeros_like', (['y'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (788, 'numpy.concatenate', 'np.concatenate', (['(x_in, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (792, 'numpy.concatenate', 'np.concatenate', (['(y_in, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (796, 'numpy.concatenate', 'np.concatenate', (['(e_in, e_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (850, 'numpy.concatenate', 'np.concatenate', (['(x_in, x_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (854, 'numpy.concatenate', 'np.concatenate', (['(y_in, y_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (858, 'numpy.concatenate', 'np.concatenate', (['(e_in, e_c)'], {'axis': '(0)'}), True, 'import numpy as np\n'), (872, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['y_p0'], {}), True, 'import tensorflow as tf\n'), (873, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['y_p1'], {}), True, 'import tensorflow as tf\n'), (879, 'numpy.where', 'np.where', (['(e == k)'], {}), True, 'import numpy as np\n'), (896, 'tensorflow.keras.losses.SparseCategoricalCrossentropy', 'tf.keras.losses.SparseCategoricalCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (932, 'sklearn.utils.shuffle', 'shuffle', (['x_in', 'y_in', 'e_in'], {}), False, 'from sklearn.utils import shuffle\n'), (933, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', ([], {}), True, 'import tensorflow as tf\n'), (934, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (211, 'numpy.mean', 'np.mean', (['[self.env_train_accs[i_env][2 * i], self.env_train_accs[i_env][2 * i + 1]]'], {}), True, 'import numpy as np\n'), (223, 'matplotlib.pyplot.plot', 'plt.plot', (['env_train_accs[i_env]'], {'label': "('train - env %i' % (i_env + 1))"}), True, 'import matplotlib.pyplot as plt\n'), (303, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (319, 'torch.tensor', 'torch.tensor', (['y'], {}), False, 'import torch\n'), (377, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (388, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate', 'beta_1': 'beta_1'}), True, 'import tensorflow as tf\n'), (425, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (428, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (520, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (530, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), True, 'import tensorflow as tf\n'), (568, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (571, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (661, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (717, 'tensorflow.zeros_like', 'tf.zeros_like', (['y_in'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (721, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (802, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), True, 'import tensorflow as tf\n'), (866, 'tensorflow.nn.softmax', 'tf.nn.softmax', (['y_pred'], {}), True, 'import tensorflow as tf\n'), (868, 'tensorflow.where', 'tf.where', (['(y == 0)'], {}), True, 'import tensorflow as tf\n'), (869, 'tensorflow.where', 'tf.where', (['(y == 1)'], {}), True, 'import tensorflow as tf\n'), (870, 'tensorflow.gather', 'tf.gather', (['y_p', 'ind_0'], {}), True, 'import tensorflow as tf\n'), (871, 'tensorflow.gather', 'tf.gather', (['y_p', 'ind_1'], {}), True, 'import tensorflow as tf\n'), (888, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (911, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (944, 'tensorflow.keras.metrics.SparseCategoricalAccuracy', 'tf.keras.metrics.SparseCategoricalAccuracy', ([], {}), True, 'import tensorflow as tf\n'), (105, 'sklearn.utils.shuffle', 'shuffle', (['x_env', 'y_env'], {}), False, 'from sklearn.utils import shuffle\n'), (407, 'sklearn.utils.shuffle', 'shuffle', (['x_e', 'y_e'], {}), False, 'from sklearn.utils import shuffle\n'), (550, 'sklearn.utils.shuffle', 'shuffle', (['x_e', 'y_e'], {}), False, 'from sklearn.utils import shuffle\n'), (672, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': 'learning_rate'}), True, 'import tensorflow as tf\n'), (674, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', ([], {'learning_rate': '(learning_rate * 0.1)'}), True, 'import tensorflow as tf\n'), (867, 'tensorflow.shape', 'tf.shape', (['y_p'], {}), True, 'import tensorflow as tf\n'), (296, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (695, 'sklearn.utils.shuffle', 'shuffle', (['x_e', 'y_e'], {}), False, 'from sklearn.utils import shuffle\n'), (697, 'sklearn.utils.shuffle', 'shuffle', (['x_in', 'y_in'], {}), False, 'from sklearn.utils import shuffle\n')]