# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Training to forget nuisance variables."""

import collections
import os

from absl import app
from absl import flags
from absl import logging

import numpy as np
import pandas as pd
from scipy import sparse
from six.moves import range
from six.moves import zip
import six.moves.cPickle as pickle
import tensorflow.compat.v1 as tf
from tensorflow.compat.v1 import gfile

from correct_batch_effects_wdn import io_utils
from correct_batch_effects_wdn import transform
from tensorflow.python.ops import gen_linalg_ops  # pylint: disable=g-direct-tensorflow-import

INPUT_NAME = "inputs"
OUTPUT_NAME = "outputs"
DISCRIMINATOR_NAME = "discriminator"
CLASSIFIER_NAME = "classifier"
UNIT_GAUSS_NAME = "unit_gauss"
CRITIC_LAYER_NAME = "critic_layer"
POSSIBLE_LOSSES = (DISCRIMINATOR_NAME, CLASSIFIER_NAME, UNIT_GAUSS_NAME)
INPUT_KEY_INDEX = 0
REGISTRY_NAME = "params_registry_python_3"
WASSERSTEIN_NETWORK = "WassersteinNetwork"
WASSERSTEIN_2_NETWORK = "Wasserstein2Network"
MEAN_ONLY_NETWORK = "MeanOnlyNetwork"
WASSERSTEIN_SQRD_NETWORK = "WassersteinSqrdNetwork"
WASSERSTEIN_CUBED_NETWORK = "WassersteinCubedNetwork"
POSSIBLE_NETWORKS = [
    WASSERSTEIN_NETWORK, WASSERSTEIN_2_NETWORK, WASSERSTEIN_SQRD_NETWORK,
    WASSERSTEIN_CUBED_NETWORK, MEAN_ONLY_NETWORK
]

FLAGS = flags.FLAGS

flags.DEFINE_string("input_df", None, "Path to the embedding dataframe.")
flags.DEFINE_string("save_dir", None, "location of file to save.")
flags.DEFINE_integer("num_steps_pretrain", None, "Number of steps to pretrain.")
flags.DEFINE_integer("num_steps", None, "Number of steps (after pretrain).")
flags.DEFINE_integer("disc_steps_per_training_step", None, "Number critic steps"
                     "to use per main training step.")
flags.DEFINE_enum("network_type", "WassersteinNetwork", POSSIBLE_NETWORKS,
                  "Network to use. Can be WassersteinNetwork.")
flags.DEFINE_integer("batch_n", 10, "Number of points to use per minibatch"
                     "for each loss.")
flags.DEFINE_float("learning_rate", 1e-4, "Initial learning rate to use.")
flags.DEFINE_float("epsilon", 0.01, "Regularization for covariance.")
flags.DEFINE_integer("feature_dim", 192, "Number of feature dimensions.")
flags.DEFINE_integer("checkpoint_interval", 4000, "Frequency to save to file.")
flags.DEFINE_spaceseplist("target_levels", "compound",
                          "dataframe target levels.")
flags.DEFINE_spaceseplist("nuisance_levels", "batch",
                          "dataframe nuisance levels.")
flags.DEFINE_integer(
    "layer_width", 2, "Width of network to use for"
    "approximating the Wasserstein distance.")
flags.DEFINE_integer(
    "num_layers", 2, "Number of layers to use for"
    "approximating the Wasserstein distance.")
flags.DEFINE_string(
    "reg_dir", None, "Directory to registry file, or None to"
    "save in save_dir.")
flags.DEFINE_float("lambda_mean", 0., "Penalty for the mean term of the affine"
                   "transformation.")
flags.DEFINE_float("lambda_cov", 0., "Penalty for the cov term of the affine"
                   "transformation.")
flags.DEFINE_integer("seed", 42, "Seed to use for numpy.")
flags.DEFINE_integer("tf_seed", 42, "Seed to use for tensorflow.")
flags.DEFINE_float(
    "cov_fix", 0.001, "Multiple of identity to add if using"
    "Wasserstein-2 distance.")

################################################################################
##### Functions and classes for storing and retrieving data
################################################################################


def get_dense_arr(matrix):
  """Convert a sparse matrix to numpy array.

  Args:
    matrix (matrix, sparse matrix, or ndarray): input

  Returns:
    dense numpy array.
  """
  if sparse.issparse(matrix):
    return matrix.toarray()
  else:
    return np.array(matrix)


class DataShuffler(object):
  """Object to hold and shuffle data.

  Adapted from
  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/learn/python/learn/datasets/mnist.py

  Attributes:
      inputs (ndarray): The inputs specified in __init__.
      outputs (ndarray): The outputs specified in __init__.
  """

  def __init__(self, inputs, outputs, random_state):
    """Inits DataShuffler given inputs, outputs, and a random state.

    Args:
      inputs (ndarray): 2-dimensional array containing the inputs, where each
        row is an individual input. The columns represent the different
        dimensions of an individual entry.
      outputs (ndarray): 2-dimensional array containing the outputs, where each
        row is an individual output.
      random_state (None or int): seed to feed to numpy.random.
    """
    assert inputs.shape[0] == outputs.shape[0]
    self.inputs = inputs
    self.outputs = outputs
    self._epochs_completed = -1
    self._num_examples = inputs.shape[0]
    self._random_state = random_state
    self._next_indices = []

  def next_batch(self, batch_size, shuffle=True):
    """Helper method for next_batch.

    Args:
      batch_size (int): Number of items to pick.
      shuffle (bool): whether to shuffle the data or not.

    Returns:
      A tuple of 2-dimensional ndarrays whose shape along the first axis is
        equal to batch_size. The rows in the first element correspond to inputs,
        and in the second element to the outputs.
    """
    indices = []
    while len(indices) < batch_size:
      if not self._next_indices:
        self._epochs_completed += 1
        self._next_indices = list(reversed(list(range(self._num_examples))))
        if shuffle:
          self._random_state.shuffle(self._next_indices)
      indices.append(self._next_indices.pop())
    return (get_dense_arr(self.inputs[indices]),
            get_dense_arr(self.outputs[indices]))


def _make_canonical_key(x):
  """Try to convert to a hashable type.

  First, if the input is a list of length 1, take its first component instead.
  Next, try to convert to a tuple if not hashable.

  Args:
    x (list, tuple, string, or None): input to convert to a key

  Returns:
    Hashable object
  """
  if isinstance(x, list) and len(x) == 1:
    x = x[0]
  if not isinstance(x, collections.Hashable):
    return tuple(x)
  else:
    return x


def split_df(df, columns_split):
  """Split a dataframe into two by column.

  Args:
    df (pandas dataframe): input dataframe to split.
    columns_split (int): Column at which to split the dataframes.

  Returns:
    df1, df2 (pandas dataframes): Split df into two dataframe based on column.
    The first has the first column_split-1 columns, and the second has columns
    columns_split onward.

  """
  return df.iloc[:, :columns_split], df.iloc[:, columns_split:]


def tuple_in_group_with_wildcards(needle, haystack):
  """Checks if needle is in haystack, allowing for wildcard components.

  Returns True if either needle or haystack is None, or needle is in haystack.
  Components in haystack are tuples. These tuples can have entries equal to None
  which serve as wildcard components.
  For example, if haystack = [(None, ...), ...], The first tuple in haystack
  has a wildcard for its first component.

  Args:
    needle (tuple): tuple to check if in group
    haystack (list): list of tuples of the same dimensions as tup.

  Returns:
    True if the tuple is in the group, False otherwise

  Raises:
    ValueError: Length of tup must match length of each tuple in group."
  """
  if needle is None or haystack is None:
    return True
  if any(len(needle) != len(it_needle) for it_needle in haystack):
    raise ValueError("Length of tup must match length of each tuple in group.")
  for it_needle in haystack:
    if all(needle_comp == it_needle_comp or it_needle_comp is None
           for needle_comp, it_needle_comp in zip(needle, it_needle)):
      return True
  return False


class DatasetHolder(object):
  """Object to hold datasets from which minibatches are sampled.

  Attributes:
    df_with_input_one_hot (pandas dataframe): Dataframe formatted as inputs
      including one-hot encoding for relevant categories.
    data_shufflers (instances of data_shuffler): The stored data shufflers.
    input_dim (int): Shape of original dataframe df used as an input.
    encoding_to_num (dict): Maps each encoding used to a unique integer.
  """

  def __init__(self,
               df,
               input_category_level=None,
               batch_input_info=None,
               input_output_type=np.array,
               random_state=None):
    """Inits DatasetHolder given a pandas dataframe and specifications.

    Args:
      df (pandas dataframe): all the input points, labeled by index.
      input_category_level (int, level name, or sequence of such, or None): Same
        as level used by pandas groupby method. Granularity level for labels
        added to inputs, unless no_input_labels=True. In that case the original
        inputs are provided
      batch_input_info (string or None): Formatting for the inputs
      input_output_type (class or module to store input and output information):
        Currently supports sparse.lil_matrix and np.array. np.array is the
        default setting.
      random_state (np.random.RandomState): Instance of RandomState used for
        shuffling data.

    Raises:
      ValueError: input_output_type not implemented. Use np.array or
        sparse.lil_matrix.
    """

    if input_output_type not in [np.array, sparse.lil_matrix]:
      raise ValueError("input_output_type not implemented. Use np.array or"
                       "sparse.lil_matrix.")
    self._input_output_type = input_output_type
    if input_output_type == np.array:
      self._zeros_maker = np.zeros
    else:
      self._zeros_maker = sparse.lil_matrix
    self._input_category_level = input_category_level
    self._batch_input_info = batch_input_info
    self.data_shufflers = collections.OrderedDict()
    self.input_dim = df.shape[1]
    self._num_input_categories = len(df.columns)
    self.df_with_input_one_hot = self.categorize_one_hot(
        df, category_level=self._input_category_level, drop_original=False)
    input_one_hot = self.df_with_input_one_hot.drop(
        self.df_with_input_one_hot.columns[list(range(self.input_dim))], axis=1)
    keys = [tuple(row) for row in input_one_hot.values]
    unique_keys = list(collections.OrderedDict.fromkeys(keys))
    self.encoding_to_num = {tuple(row): i for i, row in enumerate(unique_keys)}
    if random_state:
      self._random_state = random_state
    else:
      self._random_state = np.random.RandomState(seed=42)

  def get_random_state(self):
    """Get the random state.

    This is useful for testing.
    Returns:
      random_state (numpy random state).
    """
    return self._random_state

  def add_shufflers(self,
                    group_level=None,
                    label_level=None,
                    which_group_labels=None):
    """Add data_shuffler instances to the DatasetHolder.

    Args:
      group_level (int, level name, or sequence of such, or None): Same as level
        used by pandas groupby method, except when None. Indicates groups of
        inputs over which shufflers are created.
      label_level (int, level name, or sequence of such, or None): Same as level
        used by pandas groupby method, except when None. Indicates level to make
        distinct in outputs.
      which_group_labels (list or None): List to indicate which group_level
        values to use, or None to use all the possible ones.

    Returns:
      new_data_shufflers (dict): The newly added DataShuffler instances.
    """
    if group_level is None:
      groups = [[None, self.df_with_input_one_hot]]
    else:
      groups = self.df_with_input_one_hot.groupby(level=group_level)
    new_data_shufflers = {}
    for group_label, df_piece in groups:
      if not tuple_in_group_with_wildcards(group_label, which_group_labels):
        continue
      label_level_key = _make_canonical_key(label_level)
      data_shuffler_key = (group_label, label_level_key)
      data_shuffler = self.make_data_shuffler(
          df_piece, output_category_level=label_level)
      new_data_shufflers[data_shuffler_key] = data_shuffler
      self.data_shufflers.update(new_data_shufflers)
    return new_data_shufflers

  def make_data_shuffler(self,
                         df_with_input_one_hot,
                         output_category_level=None):
    """Generate a data shuffler from a given dataframe.

    Args:
      df_with_input_one_hot(pandas dataframe): Rows are known points.
      output_category_level (int, level name, or sequence of such, or None):
        Same as level used by pandas groupby method, except when None.

    Returns:
      a data_shuffler instance. The shuffler maps entries in the given dataframe
      to a one hot encoding of the labels at the provided granularity level.

    Raises:
      ValueError: Unknown value for batch_input_info.
    """
    input_df, input_one_hot = split_df(df_with_input_one_hot,
                                       self._num_input_categories)
    output_one_hot = self.categorize_one_hot(
        input_df, category_level=output_category_level)
    outputs = self._input_output_type(output_one_hot.values)
    if self._batch_input_info is not None:
      if self._batch_input_info == "multiplexed":
        inputs = self.encoding_to_multiplexed_encoding(input_df.values,
                                                       input_one_hot.values)
      elif self._batch_input_info == "one_hot":
        inputs = self._input_output_type(df_with_input_one_hot.values)
      else:
        raise ValueError("unknown value for batch_input_info.")
    else:
      inputs = self._input_output_type(input_df.values)
    return DataShuffler(inputs, outputs, self._random_state)

  def categorize_one_hot(self, df, category_level=None, drop_original=True):
    """Generate one-hot encoding from a given dataframe and selected categories.

    Args:
      df (pandas dataframe): input dataframe.
      category_level (int, level name, or sequence of such, or None): Same as
        level used by pandas groupby method, except when None. Used to indicate
        which indices to use for one-hot encoding.
      drop_original (bool): whether or not to drop the original table, leaving
        only the one-hot encoding vector.

    Returns:
      A tuple (df_one_hot, num_columns) where df_one_hot is the generated
        dataframe and num_columns is the number of columns of the original
        input dataframe df.
    """

    ## Ensure dataframe indices are strings
    index_names = df.index.names
    index = df.index.map(mapper=lambda x: tuple(str(v) for v in x))
    index.names = index_names
    df.index = index

    num_columns = len(df.columns)

    ## convert category_level from index to row values
    df_reset_index = df.reset_index(category_level)

    ## convert non-embedding values to one-hot encoding
    df_with_one_hot = pd.get_dummies(df_reset_index)

    ## restore indices
    df_with_one_hot = pd.DataFrame(data=df_with_one_hot.values, index=df.index)

    if drop_original:
      _, one_hot_df = split_df(df_with_one_hot, num_columns)
      return one_hot_df
    else:
      return df_with_one_hot

  def encoding_to_multiplexed_encoding(self, arr, encoding):
    """Generate a multiplexed encoding.

    For each entry, we have a row from arr and a row from encoding, specifying
    the original inputs the encoding value. In the multiplexed encoding
    scheme, the new dimension is num_bins * (num_dim + 1), where num_dim is the
    dimension of the original arr, and num_bins is the number of unique encoding
    values. Almost all values are set to zero, except for the values in bin i,
    where i is an index corresponding to the encoding value. Bin i is populated
    by the original entry from arr, followed by a 1.

    Args:
      arr (ndarray): 2 dimensional numpy array representing the original inputs.
        The rows are the individual entries, and the columns are the various
        coorindates.
      encoding (ndarray): 2 dimensional numpy array whose rows represent
        encodings.

    Returns:
      multiplexed_encoding (ndarray): 2 dimensional array as described above.

    Raises:
      ValueError: First dimension of arr and encoding must match.

    """
    if arr.shape[0] != encoding.shape[0]:
      raise ValueError("First dimension of arr and encoding must match.")
    num_points = arr.shape[0]
    num_dim = arr.shape[1]
    num_categories = len(self.encoding_to_num)
    ones = np.ones((num_points, 1))
    values_with_ones = np.hstack((arr, ones))
    ## TODO(tabakg): make this faster
    multiplexed_values = self._zeros_maker(
        (num_points, (num_dim + 1) * num_categories))
    for row_idx, (value, enc) in enumerate(zip(values_with_ones, encoding)):
      bin_num = self.encoding_to_num[tuple(enc)]
      multiplexed_values[row_idx, bin_num * (num_dim + 1):(bin_num + 1) *
                         (num_dim + 1)] = value
    return multiplexed_values


################################################################################
##### Code for Neural Network
################################################################################


def reverse_gradient(tensor):
  return -tensor + tf.stop_gradient(2 * tensor)


def make_tensor_dict(variable):
  """Returns a function that acts like a dictionary of tensors.

  Notice that the values provided to tf.case must be functions, which may have
  a particular scope. Since our input is a dictionary, we have to include its
  values in the scope of each of the functions provided to tf.case.

  Args:
    variable (dict): Input dictionary.

  Returns:
    f (function): A function mapping the keys from the given dictionary to its
    values in tensorflow.
  """

  ## TODO(tabakg): consider using tf.HashTable as a possible alternative.
  def variable_func(x):
    return lambda: variable[x]

  def dict_func(x):
    return tf.case(
        {tf.reduce_all(tf.equal(x, k)): variable_func(k) for k in variable})

  return dict_func


def make_wb(input_dim, output_dim, name):
  w = tf.Variable(
      np.eye(max(input_dim, output_dim))[:input_dim, :output_dim],
      name="w_" + name,
      dtype=tf.float32)
  b = tf.Variable(np.zeros(output_dim), name="b_" + name, dtype=tf.float32)
  return w, b


def make_discriminator_model(input_dim, activation, layer_width, num_layers):
  """Generates multi-layer model from the feature space to a scalar.

  Args:
    input_dim (int): Number of dimensions of the inputs.
    activation (function): activation layers to use.
    layer_width (int): Number of neurons per layer.
    num_layers (int): Number of layers not including the inputs.

  Returns:
    discriminator_model (function): Maps an input tensor to the output of the
       generated network.
  """
  w = {}
  b = {}
  w[0], b[0] = make_wb(input_dim, layer_width, name=CRITIC_LAYER_NAME + "0")
  for i in range(1, num_layers):
    w[i], b[i] = make_wb(
        layer_width, layer_width, name=CRITIC_LAYER_NAME + str(i))
  w[num_layers], b[num_layers] = make_wb(
      layer_width, 1, name=CRITIC_LAYER_NAME + str(num_layers))

  def discriminator_model(x0):
    x = {0: x0}
    for i in range(num_layers):
      x[i + 1] = activation(tf.matmul(x[i], w[i]) + b[i])
    return tf.matmul(x[num_layers], w[num_layers]) + b[num_layers]

  return discriminator_model


def wasserstein_distance(x_,
                         y_,
                         layer_width,
                         num_layers,
                         batch_n,
                         penalty_lambda=10.,
                         activation=tf.nn.softplus,
                         seed=None):
  """Estimator of the Wasserstein-1 distance between two distributions.

  This is based on the loss used in the following paper:

  Gulrajani, Ishaan, et al. "Improved training of wasserstein gans."
  arXiv preprint arXiv:1704.00028 (2017).

  One important difference between the following implementation and the paper
  above is that we only use the gradient constraint from above. This seems to
  work better for our problem in practice.

  Args:
    x_ (tf.Tensor): Empirical sample from first distribution. Its dimensions
      should be [batch_n, input_dim]
    y_ (tf.Tensor): Empirical sample from second distribution. Its dimensions
      should be [batch_n, input_dim].
    layer_width (int): Number of neurons to use for the discriminator model.
    num_layers (int): Number of layers to use for the discriminator model.
    batch_n (int): Number of elements per batch of x_ and y_.
    penalty_lambda (float): Penalty used to enforce the gradient condition.
      Specifically, the norm of discriminator_model should be no more than 1.
    activation (function): activation layers to use.
    seed (int): Used to randomly pick points where the gradient is evaluated.
      Using seed=None will seed differently each time the function is called.
      However, if using a global seed (tf.set_random_seed(seed)), this will
      reproduce the same results for each run.

  Returns:
    disc_loss (scalar tf.Tensor): The estimated Wasserstein-1 loss.
    gradient_penalty (scalalr tf.Tensor):: Value on the gradient penalty.
    discriminator_model (function): The function used to estimate the
      Wasserstein distance.
  """

  ## Get the number of elements and input size
  _, input_dim = x_.get_shape().as_list()

  ## make discriminator
  discriminator_model = make_discriminator_model(input_dim, activation,
                                                 layer_width, num_layers)

  ## random point for gradient penalty
  epsilon_rand = tf.random_uniform([batch_n, 1],
                                   minval=0,
                                   maxval=1,
                                   dtype=tf.float32,
                                   seed=seed)
  ## make this constant to test source of non-determinism.
  ## At least on desktop machine, this did not cause non-determinism.
  # epsilon_rand = tf.constant(0.5, shape=[batch_n, 1], dtype=tf.float32)
  x_hat = epsilon_rand * x_ + (1.0 - epsilon_rand) * y_

  ## gradient penalty
  (gradient,) = tf.gradients(discriminator_model(x_hat), [x_hat])
  gradient_penalty = penalty_lambda * tf.square(
      tf.maximum(0.0,
                 tf.norm(gradient, ord=2) - 1.0))

  ## calculate discriminator's loss
  disc_loss = (
      tf.reduce_mean(discriminator_model(y_) - discriminator_model(x_)) +
      gradient_penalty)

  return disc_loss, gradient_penalty, discriminator_model


def cov_tf(mean_x, x):
  """Compute the sample covariance matrix.

  Args:
    mean_x (tf.Tensor): Dimensions should be [1, input_dim]
    x (tf.Tensor): Dimensions should be [batch_n, input_dim]

  Returns:
    cov_xx (tf.Tensor): covariance of x.
  """
  mean_centered_x = x - mean_x
  n = tf.cast(tf.shape(x)[0], tf.float32)  ## number of points
  return tf.matmul(tf.transpose(mean_centered_x), mean_centered_x) / (n - 1)


def mat_sqrt_tf(mat):
  """Compute the square root of a non-negative definite matrix."""
  return gen_linalg_ops.matrix_square_root(mat)

  ### old version -- not sure if this would support gradients...
  # [e, v] = tf.self_adjoint_eig(mat)
  # return tf.matmul(tf.matmul(v, tf.diag(tf.sqrt(e))), tf.transpose(v))


def wasserstein_2_distance(x_, y_, mean_only=False):
  """Compute the wasserstein-2 distance squared between two distributions.

  This uses the closed form and assumes x_ and y_ are sampled from Gaussian
  distributions.

  Based on two_wasserstein_tf in
  research/biology/diagnose_a_well/analysis/distance.py

  Args:
    x_ (tf.Tensor): Empirical sample from first distribution. Its dimensions
      should be [batch_n, input_dim]
    y_ (tf.Tensor): Empirical sample from second distribution. Its dimensions
      should be [batch_n, input_dim].
    mean_only (bool): Restrict to use only the mean part.

  Returns:
    wasserstein-2 distance between x_ and y_.
  """
  mean_x_ = tf.reduce_mean(x_, axis=0, keep_dims=True)
  mean_y_ = tf.reduce_mean(y_, axis=0, keep_dims=True)

  if mean_only:
    return transform.sum_of_square(mean_x_ - mean_y_)

  ## If not using mean_only, compute the full W-2 distance metric:

  ## TESTING: adding identity to avoid error
  cov_x_ = cov_tf(mean_x_, x_) + FLAGS.cov_fix * tf.eye(FLAGS.feature_dim)
  cov_y_ = cov_tf(mean_y_, y_) + FLAGS.cov_fix * tf.eye(FLAGS.feature_dim)
  sqrt_cov_x_ = mat_sqrt_tf(cov_x_)

  prod = tf.matmul(tf.matmul(sqrt_cov_x_, cov_y_), sqrt_cov_x_)
  return transform.sum_of_square(mean_x_ -
                                 mean_y_) + tf.trace(cov_x_ + cov_y_ -
                                                     2.0 * mat_sqrt_tf(prod))


class Network(object):
  """Learns features that forget desired information.

  Since we may want different losses for different parts of the data, we use
  an instance of DatasetHolder containing instances of DataShuffler. They keys
  for the dictionary attributes are the same as the ones used for the
  DataShuffler instance, unless otherwise noted.

  There are several implementations, using ideas from the following paper to
  train a discriminator (critic), and worsen its performance during training:

  Ganin, Yaroslav, and Victor Lempitsky. "Unsupervised domain adaptation by
  backpropagation." International Conference on Machine Learning. 2015.

  This class has only the attributes common to all networks we will use for
  experiments. Inheriting networks will have additional components completing
  the network and implementaing train, predict, and evaluate.

  Atributes:
    holder (DatasetHolder): Holds the various DataShuffler instances to use.
    loss_specifiers (dict): Maps keys to loss types.
    losses (dict): Maps each loss type to a dictionary, which maps keys to
      the relevant tensorflow objects representing losses.
    input_dim_with_encoder (int): Total input size, including encoding.
    input_dim (int): Input size, not including encoding.
    encoding_dim (int): Input size of the encoding.
  """

  def __init__(self, holder, feature_dim, batch_n):
    """Inits Network with a given instance of DatasetHolder and parameters.

    Args:
      holder (instance of DatasetHolder): Provides data to train model
      feature_dim (int): Dimension of feature space. Currently supports the same
        dimension as the input space only. The reason is that for now we wish to
        initialize the transformation to the identity. However, this could be
        extended in the future.
      batch_n (int): Number of points to use from each DataShuffler during
        training. Currently this is the same for each DataShuffler.

    Raises:
      ValueError: Currently only supporting feature_dim == input_dim.
    """

    self.holder = holder
    self.losses = {loss: {} for loss in POSSIBLE_LOSSES}

    ## initialize useful dimensions and constants
    _, self._input_dim_with_encoder = holder.df_with_input_one_hot.shape
    self._input_dim = holder.input_dim
    self._encoding_dim = self._input_dim_with_encoder - self._input_dim
    self._feature_dim = feature_dim
    if self._input_dim != self._feature_dim:
      raise ValueError("Currently only supporting feature_dim == input_dim. "
                       "But we have feature_dim = %s and input_dim = %s" %
                       (feature_dim, holder.input_dim))
    self._batch_n = batch_n


def ensure_is_lst(z):
  """If a variable is not a list, return a list with only the variable.

  Args:
    z (object): input variable

  Returns:
    z if it is a list, otherwise [z]
  """
  if isinstance(z, list):
    return z
  else:
    return [z]


def get_filtered_key(input_key, indices):
  """Filter input keys to indices.

  Args:
    input_key (tuple): Contains input key information.
    indices: Which indices to select from the input key.

  Returns:
    filtered_key (tuple): The components of the input key designated by indices.
  """
  return tuple([input_key[i] for i in indices])


def sum_of_square(a):
  """Sum of squared elements of Tensor a."""
  return tf.reduce_sum(tf.square(a))


class WassersteinNetwork(Network):
  """Network with Wasserstein loss to forget batch effects.

  This network imposes only Wasserstein losses given a DatasetHolder
  containing instances of Datashuffler.

  The DataShuffler instances in the DatasetHolder are indexed by keys
  corresponding to input index levels for each DatasetHolder, as well as some
  specified output values. Here we only use the inputs (first component of each
  key, which we will call the input key).

  For example, this element might have form: (drug1, batch3).

  The arguments target_levels and nuisance_levels are used to specify the
  indices of the input key that will correspond to variables across which
  forgetting occurs, and variables that should be forgotten, respectively.
  That is, for each unique combination of target_levels, we consider the
  Wasserstein loss of each possible pairwise combination of nuisance_levels.

  For example, in the case when the input key has two elements (drug and batch
  in that order), specifying target_levels=[0] and nuisance_levels=[1] would
  construct the pairwise Wasserstein loss between all possible batches for each
  fixed individual drug.

  For this network, impose the condition that all DataShuffler instances have
  keys with matching input dimensionality. This ensures they all contain the
  necessary information to match distributions across some label categories
  but not others.

  The transformation from the inputs to the features is currently stored using
  the dictionaries w and b. This should be modified if a nonlinear function is
  used instead.

  Attributes:
    wass_loss_target (dict): Maps target keys to the sum of pairwise Wasserstine
      losses for the possible nuisance variables with the same target key.
      Note: The losses are the negative Wasserstein distances.
    grad_pen_target (dict): Maps target keys to the sum of pairwise gradient
      penalties for the possible nuisance variables with the same target key.
    w (dict): Mapping from keys to w tensorflow tensors going from inputs to
      features.
    b (dict): Mapping from keys to b tensorflow tensors going from inputs to
      features.
    ignore_disc (bool): If this is true, do not train a discriminator. This
      should be set to True when using a distance that does not need to be
      learned, e.g. the Wasserstein-2 distance.
  """

  def __init__(self,
               holder,
               feature_dim,
               batch_n,
               target_levels,
               nuisance_levels,
               layer_width=2,
               num_layers=2,
               lambda_mean=0.,
               lambda_cov=0.,
               power=1.,
               ignore_disc=False,
               mean_only=False):
    """Inits WassersteinNetwork.

    Args:
      holder (instance of DatasetHolder): Provides data to train model
      feature_dim (int): Dimension of feature space. Currently supports the same
        dimension as the input space only.
      batch_n (int): Number of points to use from each DataShuffler during
        training. Currently this is the same for each DataShuffler.
      target_levels (int or list of ints): Index or indices indicating which
        components of the input keys should have the property that for points
        sampled with the same values for these components the distributions
        should be indistinguishable.
      nuisance_levels (int or list of ints): Index or indices indicating which
        components of the input keys
      layer_width (int): number of neurons to use per layer for each function
        estimating a Wasserstein distance.
      num_layers (int): number of layers to use for estimating the Wasserstein
        loss (not including input).
      lambda_mean (float): penalty term for the mean term of the transformation.
      lambda_cov (float): penalty term for the cov term of the transformation.
      power (float): power of each pair-wise wasserstein distance to use.
      ignore_disc (bool): If this is true, do not train a discriminator. This
        should be set to True when using a distance that does not need to be
        learned, e.g. the Wasserstein-2 distance.
      mean_only (bool): Using the Wasserstein-2 distance, but only the mean
        component (i.e. no covariance dependence). This is for experimental
        purposes.

    Raises:
      ValueError: Keys must have the same input dimensionality.
    """

    super(WassersteinNetwork, self).__init__(holder, feature_dim, batch_n)
    self._layer_width = layer_width
    self._num_layers = num_layers
    self._target_levels = ensure_is_lst(target_levels)
    self._nuisance_levels = ensure_is_lst(nuisance_levels)
    self._lambda_mean = lambda_mean
    self._lambda_cov = lambda_cov
    shufflers = holder.data_shufflers

    ## The first component of each key indicates which categories have been
    ## used for sampling (i.e. the input keys).
    input_key_lengths = [len(key[INPUT_KEY_INDEX]) for key in shufflers]
    if not all(length == input_key_lengths[0] for length in input_key_lengths):
      raise ValueError("Keys must have the same input dimensionality.")

    ## Generate all possible keys using only the components of the target
    ## or nuisance variables.
    self._unique_targets = sorted(
        list(
            set([
                get_filtered_key(s[INPUT_KEY_INDEX], self._target_levels)
                for s in shufflers
            ])))
    self._unique_nuisances = sorted(
        list(
            set([
                get_filtered_key(s[INPUT_KEY_INDEX], self._nuisance_levels)
                for s in shufflers
            ])))

    ## Map from each possible target key to all input keys that generated it.
    self._keys_for_targets = collections.defaultdict(list)
    for target in self._unique_targets:
      for s in shufflers:
        if target == get_filtered_key(s[INPUT_KEY_INDEX], self._target_levels):
          self._keys_for_targets[target].append(s)

    ## Generate input placeholders.
    self._x_inputs, self._x_vals = self.get_input_vals()

    ## Make features.
    self.w, self.b, self._features = self.add_input_to_feature_layer()

    ## Generate pairwise Wasserstein losses.
    if ignore_disc:
      self.wass_loss_target = self.pairwise_wasserstein_2(mean_only=mean_only)
      self.grad_pen_target = None
      self.ignore_disc = True
    else:
      self.wass_loss_target, self.grad_pen_target = self.pairwise_wasserstein(
          power=power)
      self.ignore_disc = False

  def get_input_vals(self):
    """Obtain the input values from a given dataframe.

    There might be additional encodings which we wish to strip away.
    The additional encodings have dimension _encoding_dim and appear first.
    The input values have dimension _input_dim and appear second.

    Returns:
      x_vals (dict): Maps input keys to placeholders for the input values.
    """
    x_vals = {}
    x_inputs = {}
    for key in self.holder.data_shufflers:
      x_inputs[key] = tf.placeholder(
          tf.float32, [None, self._input_dim_with_encoder],
          name="x_" + INPUT_NAME)
      x_vals[key], _ = tf.split(x_inputs[key],
                                [self._input_dim, self._encoding_dim], 1)
    return x_inputs, x_vals

  def add_input_to_feature_layer(self):
    """Add a layer from the inputs to features.

    The transformation for inputs with the same nuisance variables is fixed.

    Returns:
      features (dict): Maps each data_shuffler key to the feature tensor.
    """
    w = {}
    b = {}
    features = {}
    for nuisance_key in self._unique_nuisances:
      w[nuisance_key], b[nuisance_key] = make_wb(
          self._input_dim, self._feature_dim, name=INPUT_NAME)
    for key in self.holder.data_shufflers:
      nuisance_key = get_filtered_key(key[INPUT_KEY_INDEX],
                                      self._nuisance_levels)
      features[key] = tf.matmul(self._x_vals[key],
                                w[nuisance_key] + b[nuisance_key])
    return w, b, features

  def pairwise_wasserstein(self, power=1.):
    """Generate pairwise Wasserstein losses.

    The pairs are the various nuisance variables (e.g. batch). This is done
    separately for each target variable (e.g. compound).

    Args:
      power (float): power of each pair-wise wasserstein distance to use.

    Returns:
      wass_loss_target (dict): Maps from target keys to sum of pairwise losses.
      grad_pen_target (dict): Maps from target keys to sum of pairwise gradient
        penalties.

    """
    wass_loss_target = {}
    grad_pen_target = {}

    ## Gradient reversal if using a discriminator.
    grad_rev_features = {
        key: reverse_gradient(self._features[key]) for key in self._features
    }

    for target in self._unique_targets:
      num_per_target = len(self._keys_for_targets[target])
      # When the target appears in multiple domains.
      if num_per_target > 1:
        normalization = num_per_target * (num_per_target - 1) / 2.
        wass_loss_target[target] = 0
        grad_pen_target[target] = 0
        ## Iterate through all pairs of nuisance for a given target
        for i in range(num_per_target):
          for j in range(i + 1, num_per_target):
            key_i = tuple(self._keys_for_targets[target][i])
            key_j = tuple(self._keys_for_targets[target][j])

            ## Generate W1 distance and gradient penalty.
            wass_dists, grad_pens, _ = wasserstein_distance(
                grad_rev_features[key_i],
                grad_rev_features[key_j],
                self._layer_width,
                self._num_layers,
                self._batch_n,
                seed=None)
            wass_loss_target[target] += tf.math.pow(wass_dists,
                                                    power) / normalization
            grad_pen_target[target] += grad_pens / normalization

    return wass_loss_target, grad_pen_target

  def pairwise_wasserstein_2(self, mean_only=False):
    """Generate pairwise Wasserstein-2 squared losses.

    This uses the closed-form solution assuming Guassian distributions.

    The pairs are the various nuisance variables (e.g. batch). This is done
    separately for each target variable (e.g. compound).

    Args:
      mean_only (bool): Restrict to use only the mean part.

    Returns:
      wass_loss_target (dict): Maps from target keys to sum of pairwise losses.
      grad_pen_target (dict): Maps from target keys to sum of pairwise gradient
        penalties.

    """
    wass_loss_target = {}

    for target in self._unique_targets:
      num_per_target = len(self._keys_for_targets[target])
      # When the target appears in multiple domains.
      if num_per_target > 1:
        wass_loss_target[target] = 0
        normalization = num_per_target * (num_per_target - 1) / 2.
        ## Iterate through all pairs of nuisance for a given target
        for i in range(num_per_target):
          for j in range(i + 1, num_per_target):
            key_i = tuple(self._keys_for_targets[target][i])
            key_j = tuple(self._keys_for_targets[target][j])

            ## Generate W2 distance and gradient penalty.
            wass_2_dists = wasserstein_2_distance(
                self._features[key_i],
                self._features[key_j],
                mean_only=mean_only) / normalization
            wass_loss_target[target] += wass_2_dists

    return wass_loss_target

  def penalty_term(self):
    """Penalty term on the affine transformation.

    This can be used to ensure that the affine transformation remains close to
    the identity transformation.

    Returns:
      loss_value (tensorflow float): Value of the penalty term to add.
    """
    loss_value = 0.
    identity = tf.eye(self._input_dim, dtype=tf.float32)
    for nuisance_key in self._unique_nuisances:
      w = self.w[nuisance_key]
      b = self.b[nuisance_key]
      loss_value += (
          self._lambda_mean * sum_of_square(b) +
          self._lambda_cov * sum_of_square(w - identity) / self._input_dim)
    return loss_value / len(self.w)

  def train(self,
            save_dir,
            num_steps_pretrain,
            num_steps,
            disc_steps_per_training_step,
            learning_rate,
            tf_optimizer=tf.train.RMSPropOptimizer,
            save_checkpoints=True,
            save_to_pickle=True):
    """Trains the network, saving checkpoints along the way.

    Args:
      save_dir (str): Directory to save pickle file
      num_steps_pretrain (int): Total steps for pre-training
      num_steps (int): Total step for main training loop.
      disc_steps_per_training_step (int): Number of training steps for the
        discriminator per training step for the features.
      learning_rate (float): Learning rate for the algorithm.
      tf_optimizer (tf Optimizer): Which optimizer to use
      save_checkpoints (bool): Whether or not to use tensorflow checkpoints.
      save_to_pickle (bool): Whether or not to save to a pickle file. This may
        be convenient for some purposes.
    """

    ## Generate paths and names for saving results.
    input_df_name = os.path.basename(FLAGS.input_df)
    params = (
        ("input_df", input_df_name),
        ("network_type", FLAGS.network_type),
        ("num_steps_pretrain", num_steps_pretrain),
        ("num_steps", num_steps),
        ("batch_n", self._batch_n),
        ("learning_rate", learning_rate),
        ("feature_dim", self._feature_dim),
        ("disc_steps_per_training_step", disc_steps_per_training_step),
        ("target_levels", tuple(FLAGS.target_levels)),
        ("nuisance_levels", tuple(FLAGS.nuisance_levels)),
        ("layer_width", self._layer_width),
        ("num_layers", self._num_layers),
        ("lambda_mean", self._lambda_mean),
        ("lambda_cov", self._lambda_cov),
        ("cov_fix", FLAGS.cov_fix),
    )

    folder_name = str(params)
    folder_path = os.path.join(save_dir, folder_name)
    ## Not writing to registry...
    # if FLAGS.reg_dir is None:
    #   reg_path = os.path.join(save_dir, REGISTRY_NAME + ".pkl")
    # else:
    #   reg_path = os.path.join(FLAGS.reg_dir, REGISTRY_NAME + ".pkl")
    pickle_path = os.path.join(folder_path, "data.pkl")
    checkpoint_path = os.path.join(folder_path, "checkpoints")

    for p in [save_dir, folder_path, checkpoint_path]:
      if not gfile.Exists(p):
        gfile.MkDir(p)

    ## Tensorflow items used for training
    global_step = tf.Variable(0, trainable=False, name="global_step")
    increment_global_step_op = tf.assign(global_step, global_step + 1)
    sorted_wass_values = [
        self.wass_loss_target[k] for k in self.wass_loss_target
    ]
    wasserstein_loss = tf.reduce_mean(sorted_wass_values)
    if self.ignore_disc:
      loss = wasserstein_loss
    else:
      loss = wasserstein_loss + self.penalty_term()
      grad_pen = tf.reduce_mean(list(self.grad_pen_target.values()))
    input_vars = [
        var for var in tf.trainable_variables() if INPUT_NAME in var.name
    ]
    critic_vars = [
        var for var in tf.trainable_variables() if CRITIC_LAYER_NAME in var.name
    ]
    optimizer = tf_optimizer(learning_rate)
    input_trainer = optimizer.minimize(
        loss, global_step=global_step, var_list=input_vars)
    if not self.ignore_disc:
      critic_trainer = optimizer.minimize(
          loss, global_step=global_step, var_list=critic_vars)

    tf.summary.scalar("loss", loss)
    if not self.ignore_disc:
      tf.summary.scalar("grad_pen", grad_pen)
    ## TODO(tabakg): Figure out why summary_op is not working.
    ## There currently seems to be an issue using summay_op, so it's set to None
    ## TODO(tabakg): There is also an issue with Supervisor when trying to load
    ## existing checkpoint files.
    _ = tf.summary.merge_all()  ## Should be summary_op
    sv = tf.train.Supervisor(logdir=checkpoint_path, summary_op=None)

    ## Used for saving history to pickle file
    loss_hist = []
    grad_pen_hist = []  ## still initialize when not using critic

    ## random history, used to monitor random seeds
    random_nums = []

    def do_loss_without_step():
      """Get losses and update loss_hist and gran_pen_hist, no training step."""

      feed_dict = {}
      for key, shuffler in self.holder.data_shufflers.items():
        input_mini, _ = shuffler.next_batch(self._batch_n)
        feed_dict[self._x_inputs[key]] = input_mini
      random_nums.append(self.holder.get_random_state().uniform(
          0, 1))  ## Testing random seed
      if self.ignore_disc:
        loss_val = sess.run([loss], feed_dict=feed_dict)[0]
        grad_pen_val = None
      else:
        loss_val, grad_pen_val = sess.run([loss, grad_pen], feed_dict=feed_dict)

      loss_hist.append(loss_val)
      grad_pen_hist.append(grad_pen_val)

    def do_train_step(trainer, increment_global_step_op, train=True):
      """A single training step.

      Side effects: Updates loss_hist, grad_pen_hist

      Args:
        trainer (Tf Operation): Specifies how to update variables in each step.
        increment_global_step_op (tf op): used to increment step.
        train (boolean): Whether or not to train. If train is false, only step
          and record loss values without actually training.

      Returns:
        step (int): Current timestep.
      """
      feed_dict = {}
      for key, shuffler in self.holder.data_shufflers.items():
        input_mini, _ = shuffler.next_batch(self._batch_n)
        feed_dict[self._x_inputs[key]] = input_mini
      if train:
        if self.ignore_disc:
          _, loss_val = sess.run([trainer, loss], feed_dict=feed_dict)
          grad_pen_val = None
        else:
          _, loss_val, grad_pen_val = sess.run([trainer, loss, grad_pen],
                                               feed_dict=feed_dict)
        step = tf.train.global_step(sess, global_step)  ## get updated step.
      else:
        if self.ignore_disc:
          loss_val = sess.run([loss], feed_dict=feed_dict)[0]
          grad_pen_val = None
        else:
          loss_val, grad_pen_val = sess.run([loss, grad_pen],
                                            feed_dict=feed_dict)
        ## if trainer is not ran, increment global step anyway.
        step = sess.run(increment_global_step_op)

      loss_hist.append(loss_val)
      grad_pen_hist.append(grad_pen_val)

      if (step % FLAGS.checkpoint_interval == 0 and
          step >= FLAGS.num_steps_pretrain):
        if save_checkpoints:
          sv.saver.save(sess, checkpoint_path, global_step)
        if save_to_pickle:
          w_val, b_val = sess.run([self.w, self.b], feed_dict=feed_dict)
          contents = {
              step: {
                  "loss_hist": loss_hist,
                  "grad_pen_hist": grad_pen_hist,
                  "w_val": w_val,
                  "b_val": b_val,
                  "random_nums": random_nums,
              },
              "params": dict(params)
          }
          add_contents_to_pickle(pickle_path, contents)
          # add_name_to_registry(reg_path, pickle_path)

      return step

    ## Tested as a potential cause for indeterminism.
    ## This did not cause indeterminism on the desktop.
    # config = tf.ConfigProto()
    # config.inter_op_parallelism_threads = 1
    # config.intra_op_parallelism_threads = 1
    # with sv.managed_session(config=config) as sess:
    with sv.managed_session() as sess:
      step = tf.train.global_step(sess, global_step)

      ## Get initial losses without stepping
      do_loss_without_step()

      ## Pre training to adjust Wasserstein function
      while step < num_steps_pretrain:
        if self.ignore_disc:
          break
        step = do_train_step(critic_trainer, increment_global_step_op)

      ## Main training part
      main_step = step - num_steps_pretrain
      while main_step < num_steps * (disc_steps_per_training_step + 1):
        ## Adjust critic disc_steps_per_training_step times
        if disc_steps_per_training_step != 0:
          while (main_step + 1) % disc_steps_per_training_step > 0:
            if self.ignore_disc:
              break
            step = do_train_step(critic_trainer, increment_global_step_op)
            main_step = step - num_steps_pretrain
        ## Train features if estimated distance is positive
        if self.ignore_disc:
          pos_dist = True  ## if ignoring discriminator, this is not an issue
        else:
          pos_dist = (-loss_hist[-1] > 0)
        step = do_train_step(
            input_trainer, increment_global_step_op, train=pos_dist)
        main_step = step - num_steps_pretrain
      sv.stop()


################################################################################
##### Functions for Saving to Pickle.
################################################################################


def read_pickle_helper(file_path, default=set):
  """Helper function to read data from a pickle file.

  If the file exists, load it. Otherwise, initialize the default type.

  Args:
    file_path (str): Path to pickle file.
    default (iterable): Python object stored in the pickle file.

  Returns:
    contents (iterable): The loaded contents or an empty default type.
  """
  if gfile.Exists(file_path):
    with gfile.GFile(file_path, mode="r") as f:
      contents = pickle.loads(f.read())
  else:
    contents = default()
  return contents


def write_pickle_helper(file_path, contents):
  """Helper function to write contents to a pickle file.

  Args:
    file_path (str): Path to pickle file.
    contents (iterable): Contents to save to pickle file_path.
  """
  with gfile.GFile(file_path, mode="w") as f:
    f.write(pickle.dumps(contents))


def add_name_to_registry(reg_path, pickle_path):
  """Adds the file_path to the set stored in the pickle file reg_path.

  The registry file contains a dictionary whose keys are possible datasets.
  The values for each dataset key is another dictionary, mapping from a
  trasformation name to a list of file paths. Each of these paths corresponds
  to a saved transformation file that was generated using different parameters.

  Args:
    reg_path (str): Registry path.
    pickle_path (str): Path to registered pickle file.
  """
  if FLAGS.network_type == WASSERSTEIN_NETWORK:
    transform_name = "wasserstein_transform"
  elif FLAGS.network_type == WASSERSTEIN_2_NETWORK:
    transform_name = "wasserstein_2_transform"
  elif FLAGS.network_type == WASSERSTEIN_SQRD_NETWORK:
    transform_name = "wasserstein_squared_transform"
  elif FLAGS.network_type == WASSERSTEIN_CUBED_NETWORK:
    transform_name = "wasserstein_cubed_transform"
  elif FLAGS.network_type == MEAN_ONLY_NETWORK:
    transform_name = "mean_only_transform"
  else:
    raise ValueError("Unknown network type, please add...")

  reg = read_pickle_helper(reg_path, dict)
  if FLAGS.input_df not in reg:
    reg[FLAGS.input_df] = {}
  if transform_name not in reg[FLAGS.input_df]:
    reg[FLAGS.input_df][transform_name] = []
  if pickle_path not in reg[FLAGS.input_df][transform_name]:
    reg[FLAGS.input_df][transform_name].append(pickle_path)
  write_pickle_helper(reg_path, reg)


def add_contents_to_pickle(file_path, contents):
  """Adds contents to a pickle file.

  Args:
    file_path (str): Full path to the pickle file.
    contents (dict): Maps a timestep to parameters at that time.
  """
  old_contents = read_pickle_helper(file_path, dict)
  contents.update(old_contents)
  write_pickle_helper(file_path, contents)


def main(argv):
  del argv  # Unused.

  if FLAGS.network_type not in POSSIBLE_NETWORKS:
    raise ValueError("Unknown network type.")
  tf.reset_default_graph()
  with tf.Graph().as_default():
    tf.set_random_seed(FLAGS.tf_seed)

    ## Load embedding dataframe
    df = io_utils.read_dataframe_from_hdf5(FLAGS.input_df)

    ## TODO(tabakg): Add training routine for this network as well.
    ## This network did not seem to work as well, so this is secondary priority.

    using_w = WASSERSTEIN_NETWORK in FLAGS.network_type
    using_w_sqrd = WASSERSTEIN_SQRD_NETWORK in FLAGS.network_type
    using_w_cubed = WASSERSTEIN_CUBED_NETWORK in FLAGS.network_type
    using_w_1 = using_w or using_w_sqrd or using_w_cubed
    using_w_2 = WASSERSTEIN_2_NETWORK in FLAGS.network_type
    using_mean_only = MEAN_ONLY_NETWORK in FLAGS.network_type

    ## Script for initializing WassersteinNetwork
    if using_w_1 or using_w_2 or using_mean_only:
      if using_w:
        logging.info("using WASSERSTEIN_NETWORK.")
        power = 1.
      elif using_w_sqrd:
        logging.info("using WASSERSTEIN_SQRD_NETWORK.")
        power = 2.
      elif using_w_cubed:
        logging.info("using WASSERSTEIN_CUBED_NETWORK.")
        power = 3.

      ## TODO(tabakg): Possible bug when input_category_level=None.
      holder = DatasetHolder(
          df,
          input_category_level=FLAGS.nuisance_levels,
          batch_input_info="one_hot",
          random_state=np.random.RandomState(seed=FLAGS.seed))

      ## Old
      holder.add_shufflers(FLAGS.nuisance_levels + FLAGS.target_levels, None)

      # holder.add_shufflers(FLAGS.nuisance_levels, None)

      nuisance_levels = list(range(len(FLAGS.nuisance_levels)))
      target_levels = list(
          range(
              len(FLAGS.nuisance_levels),
              len(FLAGS.nuisance_levels) + len(FLAGS.target_levels)))

      if using_w_1:
        network = WassersteinNetwork(
            holder,
            FLAGS.feature_dim,
            FLAGS.batch_n,
            target_levels,
            nuisance_levels,
            layer_width=FLAGS.layer_width,
            num_layers=FLAGS.num_layers,
            lambda_mean=FLAGS.lambda_mean,
            lambda_cov=FLAGS.lambda_cov,
            power=power)

      else:  # using_w_2 or using_mean_only

        if using_mean_only:
          logging.info("using MEAN_ONLY_NETWORK.")
          mean_only = True
        else:
          logging.info("using WASSERSTEIN_2_NETWORK.")
          mean_only = False

        network = WassersteinNetwork(
            holder,
            FLAGS.feature_dim,
            FLAGS.batch_n,
            target_levels,
            nuisance_levels,
            layer_width=FLAGS.layer_width,
            num_layers=FLAGS.num_layers,
            lambda_mean=FLAGS.lambda_mean,
            lambda_cov=FLAGS.lambda_cov,
            ignore_disc=True,
            mean_only=mean_only)

      network.train(FLAGS.save_dir, FLAGS.num_steps_pretrain, FLAGS.num_steps,
                    FLAGS.disc_steps_per_training_step, FLAGS.learning_rate)


if __name__ == "__main__":
  app.run(main)
