code: # 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.
  
  """Implementation of DDPG."""
  
  import typing
  
  from dm_env import specs as dm_env_specs
  import numpy as np
  import tensorflow as tf
  from tf_agents.specs.tensor_spec import TensorSpec
  
  from representation_batch_rl.batch_rl import critic
  from representation_batch_rl.batch_rl.encoders import ConvStack
  from representation_batch_rl.batch_rl.encoders import ImageEncoder
  from representation_batch_rl.batch_rl.encoders import make_impala_cnn_network
  from representation_batch_rl.representation_batch_rl import tf_utils
  
  
  class CQL(object):
    """Class performing CQL training."""
  
    def __init__(self,
                 observation_spec,
                 action_spec,
                 actor_lr = 1e-4,
                 critic_lr = 3e-4,
                 discount = 0.99,
                 tau = 0.005,
                 target_entropy = 0.0,
                 reg = 0.0,
                 num_cql_actions = 10,
                 bc_pretraining_steps = 40_000,
                 min_q_weight = 10.0,
                 num_augmentations = 1,
                 rep_learn_keywords = 'outer',
                 batch_size = 256):
      """Creates networks.
  
      Args:
        observation_spec: environment observation spec.
        action_spec: Action spec.
        actor_lr: Actor learning rate.
        critic_lr: Critic learning rate.
        discount: MDP discount.
        tau: Soft target update parameter.
        target_entropy: Target entropy.
        reg: Coefficient for out of distribution regularization.
        num_cql_actions: Number of actions to sample for CQL loss.
        bc_pretraining_steps: Use BC loss instead of CQL loss for N steps.
        min_q_weight: CQL alpha.
        num_augmentations: Num of random crops
        rep_learn_keywords: Representation learning loss to add.
        batch_size: Batch size
      """
      self.num_augmentations = num_augmentations
      self.batch_size = batch_size
      self.rep_learn_keywords = rep_learn_keywords.split('__')
  
      critic_kwargs = {}
  
      if observation_spec.shape == (64, 64, 3):
        # IMPALA for Procgen
        def conv_stack():
          return make_impala_cnn_network(
              depths=[16, 32, 32], use_batch_norm=False, dropout_rate=0.)
  
        state_dim = 256
      else:
        # Reduced architecture for DMC
        def conv_stack():
          return ConvStack(observation_spec.shape)
        state_dim = 50
  
      conv_stack_critic = conv_stack()
      conv_target_stack_critic = conv_stack()
  
      if observation_spec.shape == (64, 64, 3):
        conv_stack_critic.output_size = state_dim
        conv_target_stack_critic.output_size = state_dim
      # Combine and stop_grad some of the above conv stacks
      critic_kwargs['encoder'] = ImageEncoder(
          conv_stack_critic, feature_dim=state_dim, bprop_conv_stack=True)
      # Note: the target critic does not share any weights.
      critic_kwargs['encoder_target'] = ImageEncoder(
          conv_target_stack_critic, feature_dim=state_dim, bprop_conv_stack=True)
  
      if self.num_augmentations == 0:
        dummy_state = tf.constant(
            np.zeros(shape=[1] + list(observation_spec.shape)))
      else:  # account for padding of +4 everywhere and then cropping out 68
        dummy_state = tf.constant(np.zeros(shape=[1, 68, 68, 3]))
  
      @tf.function
      def init_models():
        critic_kwargs['encoder'](dummy_state)
        critic_kwargs['encoder_target'](dummy_state)
  
      init_models()
  
      hidden_dims = (256, 256)
      # self.actor = policies.CategoricalPolicy(state_dim, action_spec,
      #               hidden_dims=hidden_dims, encoder=actor_kwargs['encoder'])
      action_dim = action_spec.maximum.item() + 1
  
      self.action_dim = action_dim
  
      self.log_alpha = tf.Variable(tf.math.log(1.0), trainable=True)
      self.log_cql_alpha = self.log_alpha
      self.alpha_optimizer = tf.keras.optimizers.Adam(learning_rate=actor_lr)
  
      self.critic = critic.Critic(
          state_dim,
          action_dim,
          hidden_dims=hidden_dims,
          encoder=critic_kwargs['encoder'],
          discrete_actions=True,
          linear='linear_Q' in self.rep_learn_keywords)
      self.critic_target = critic.Critic(
          state_dim,
          action_dim,
          hidden_dims=hidden_dims,
          encoder=critic_kwargs['encoder_target'],
          discrete_actions=True,
          linear='linear_Q' in self.rep_learn_keywords)
  
      @tf.function
      def init_models2():
        """This function initializes all auxiliary networks (state and action encoders) with dummy input (Procgen-specific, 68x68x3, 15 actions).
        """
        dummy_state = tf.zeros((1, 68, 68, 3), dtype=tf.float32)
        phi_s = self.critic.encoder(dummy_state)
        phi_a = tf.eye(15, dtype=tf.float32)
        if 'linear_Q' in self.rep_learn_keywords:
          _ = self.critic.critic1.state_encoder(phi_s)
          _ = self.critic.critic2.state_encoder(phi_s)
          _ = self.critic.critic1.action_encoder(phi_a)
          _ = self.critic.critic2.action_encoder(phi_a)
          _ = self.critic_target.critic1.state_encoder(phi_s)
          _ = self.critic_target.critic2.state_encoder(phi_s)
          _ = self.critic_target.critic1.action_encoder(phi_a)
          _ = self.critic_target.critic2.action_encoder(phi_a)
  
      init_models2()
  
      critic.soft_update(self.critic, self.critic_target, tau=1.0)
      self.critic_optimizer = tf.keras.optimizers.Adam(learning_rate=critic_lr)
      self.tau = tau
  
      self.reg = reg
      self.target_entropy = target_entropy
      self.discount = discount
  
      self.num_cql_actions = num_cql_actions
      self.bc_pretraining_steps = bc_pretraining_steps
      self.min_q_weight = min_q_weight
  
      self.bc = None
  
      self.model_dict = {
          'critic': self.critic,
          'critic_target': self.critic_target,
          'critic_optimizer': self.critic_optimizer,
          'alpha_optimizer': self.alpha_optimizer
      }
  
    @property
    def alpha(self):
      return tf.constant(0.)
  
    @property
    def cql_alpha(self):
      return tf.exp(self.log_cql_alpha)
  
    def fit_critic(self, states, actions,
                   next_states, next_actions, rewards,
                   discounts):
      """Updates critic parameters.
  
      Args:
        states: Batch of states.
        actions: Batch of actions.
        next_states: Batch of next states.
        next_actions: Batch of next actions from training policy.
        rewards: Batch of rewards.
        discounts: Batch of masks indicating the end of the episodes.
  
      Returns:
        Dictionary with information to track.
      """
      action_indices = tf.stack(
          [tf.range(tf.shape(actions)[0], dtype=tf.int64), actions], axis=-1)
      next_action_indices = tf.stack(
          [tf.range(tf.shape(next_actions)[0], dtype=tf.int64), next_actions],
          axis=-1)
  
      if self.num_augmentations > 1:
        target_q = 0.
        for i in range(self.num_augmentations):
          next_q1_i, next_q2_i = self.critic_target(next_states[i], actions=None)
          target_q_i = tf.expand_dims(
              rewards, 1) + self.discount * tf.expand_dims(
                  discounts, 1) * tf.minimum(next_q1_i, next_q2_i)
          target_q += target_q_i
        target_q /= self.num_augmentations
      elif self.num_augmentations == 1:
        next_q1, next_q2 = self.critic_target(
            next_states[0], actions=None, stop_grad_features=False)
        target_q = tf.expand_dims(
            rewards, 1) + self.discount * tf.expand_dims(
                discounts, 1) * tf.minimum(next_q1, next_q2)
      else:
        next_q1, next_q2 = self.critic_target(next_states, actions=None)
        target_q = tf.expand_dims(rewards, 1) + self.discount * tf.expand_dims(
            discounts, 1) * tf.minimum(next_q1, next_q2)
  
      target_q = tf.gather_nd(target_q, indices=next_action_indices)
  
      with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(self.critic.trainable_variables)
  
        if self.num_augmentations > 1:
          critic_loss = 0.
          q1 = 0.
          q2 = 0.
          for i in range(self.num_augmentations):
            q1_i, q2_i = self.critic(states[i], actions=None)
            critic_loss_i = (
                tf.losses.mean_squared_error(
                    target_q, tf.gather_nd(q1_i, indices=action_indices)) +
                tf.losses.mean_squared_error(
                    target_q, tf.gather_nd(q2_i, indices=action_indices)))
            q1 += q1_i
            q2 += q2_i
            critic_loss += critic_loss_i
          q1 /= self.num_augmentations
          q2 /= self.num_augmentations
          critic_loss /= self.num_augmentations
        elif self.num_augmentations == 1:
          q1, q2 = self.critic(states[0], actions=None)
          critic_loss = (
              tf.losses.mean_squared_error(
                  target_q, tf.gather_nd(q1, indices=action_indices)) +
              tf.losses.mean_squared_error(
                  target_q, tf.gather_nd(q2, indices=action_indices)))
        else:
          # Ensure num_augmentations is non-negative
          assert self.num_augmentations == 0
          q1, q2 = self.critic(states, actions=None)
          critic_loss = (
              tf.losses.mean_squared_error(
                  target_q, tf.gather_nd(q1, indices=action_indices)) +
              tf.losses.mean_squared_error(
                  target_q, tf.gather_nd(q2, indices=action_indices)))
        q = tf.minimum(q1, q2)
        cql_logsumexp = tf.reduce_logsumexp(q, 1)
        cql_loss = tf.reduce_mean(cql_logsumexp -
                                  tf.gather_nd(q, indices=action_indices))
  
        critic_loss += (self.reg * cql_loss)
  
      critic_grads = tape.gradient(critic_loss, self.critic.trainable_variables)
  
      self.critic_optimizer.apply_gradients(
          zip(critic_grads, self.critic.trainable_variables))
  
      critic.soft_update(self.critic, self.critic_target, tau=self.tau)
  
      return {
          'q1': tf.reduce_mean(q1),
          'q2': tf.reduce_mean(q2),
          'critic_loss': critic_loss,
          'cql_loss': cql_loss
      }
  
    @tf.function
    def update_step(self,
                    replay_buffer_iter,
                    train_target='both'):
      """Performs a single training step for critic and embedding.
  
      Args:
        replay_buffer_iter: A tensorflow graph iteratable object.
        train_target: string specifying whether update RL and or representation
  
      Returns:
        Dictionary with losses to track.
      """
      del train_target
      transition = next(replay_buffer_iter)
      numpy_dataset = isinstance(replay_buffer_iter, np.ndarray)
      # observation: n_batch x n_timesteps x 1 x H*W*3*n_frames x 1 ->
      # n_batch x H x W x 3*n_frames
      if not numpy_dataset:
        states = transition.observation[:, 0]
        next_states = transition.observation[:, 1]
        actions = transition.action[:, 0]
        rewards = transition.reward[:, 0]
        discounts = transition.discount[:, 0]
  
        if transition.observation.dtype == tf.uint8:
          states = tf.cast(states, tf.float32) / 255.
          next_states = tf.cast(next_states, tf.float32) / 255.
      else:
        states, actions, rewards, next_states, discounts = transition
  
      if self.num_augmentations > 0:
        states, next_states = tf_utils.image_aug(
            states,
            next_states,
            img_pad=4,
            num_augmentations=self.num_augmentations,
            obs_dim=64,
            channels=3,
            cropped_shape=[self.batch_size, 68, 68, 3])
  
      next_actions = self.act(next_states, data_aug=True)
  
      critic_dict = self.fit_critic(states, actions, next_states, next_actions,
                                    rewards, discounts)
  
      return critic_dict
  
    @tf.function
    def act(self, states, data_aug=False):
      """Act with batch of states.
  
      Args:
        states: tf.tensor n_batch x 64 x 64 x 3
        data_aug: bool, whether to use stochastic data aug (else deterministic)
  
      Returns:
        action: tf.tensor
      """
      if data_aug and self.num_augmentations > 0:
        states = states[0]
      if self.num_augmentations > 0:
        # use pad of 2 to bump 64 to 68 with 2 + 64 + 2 on each side
        img_pad = 2
        paddings = tf.constant(
            [[0, 0], [img_pad, img_pad], [img_pad, img_pad], [0, 0]],
            dtype=tf.int32)
        states = tf.cast(
            tf.pad(tf.cast(states * 255., tf.int32), paddings, 'SYMMETRIC'),
            tf.float32) / 255.
  
      q1, q2 = self.critic(states, actions=None)
      q = tf.minimum(q1, q2)
      actions = tf.argmax(q, -1)
      return actions
   # 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.
  
  """Utils for debiasing ML models."""
  
  import math
  
  import numpy as np
  
  
  class RandomizedThreshold:
    """Threshold optimizer (RTO) to debias models via postprocessing.
  
    See: https://arxiv.org/abs/2106.12887.
  
    This is a solver to the following optimiation problem:
      minimize   gamma/2 ||x||^2 - y^Tx
      s.t.       x satisfies DP constraint with tolerance eps and parameter rho.
  
    There are no assumptions about y in this code but, in general, y should be the
    predictions of the original classifier.
    """
  
    def __init__(self, gamma=1.0, eps=0.0, rho=None):
      """Instantiate object.
  
      Args:
        gamma: The regularization parameter gamma (for randomization). Set this to
          1 if the goal is to minmize changes to the original scores.
        eps: Tolerance parameter for bias between 0 and 1 inclusive.
        rho: The rho parameter in the post-hoc rule. If None, rho = E[y].
      """
      if eps < 0:
        raise ValueError('eps must be non-negative.')
  
      if gamma <= 0:
        raise ValueError('gamma must be a strictly positive number.')
  
      if rho is not None and rho <= 0:
        raise ValueError('rho must be either None or a strictly positive number.')
  
      self.num_groups = 1
      self.gamma = gamma
      self.eps = eps
      self.rho = rho
      self.avrg_y_score = 0
  
      # model paramters (Lagrange dual variables)
      self.lambdas = []
      self.mus = []
  
    def fit(self, y_orig, group_feature, sgd_steps,
            full_gradient_epochs=1_000, verbose=True, batch_size=256,
            ignore_warnings=False):
      """Debias predictions w.r.t. the sensitive class in each demographic group.
  
      This procedure takes as input a vector y=y_orig and solves the optimization
      problem subject to the statistical parity constraint.
        minimize_x   gamma/2 ||x||^2 - y^Tx
        s.t.      x satisfies DP constraints with tolerance eps and parameter rho.
  
      IMPORTANT: If this is used for postprocessing a classifier,
        the scores y_orig need to be rescaled linearly to [-1, +1].
  
      Training proceeds in two rounds. First is SGD. Second is full gradient
      descent. Full gradient descent is recommended when debiasing deep neural
      nets because the scores are concentrated around the extremes
      so high preciseion might be needed. Because the loss is smooth, the lr
      in full gradient method does not need tuning. It can be set to gamma / 2.0.
  
      Args:
        y_orig: A vector of the original probability scores. If this is used for
          debiasing binary classifiers, y_orig = 2 * p(y=1) -1.
        group_feature: An array containing the group id of each instance starting
          from group 0 to group K-1.
        sgd_steps: Number of minibatch steps in SGD.
        full_gradient_epochs: Number of epochs in full gradient descent phase.
        verbose: Set to True to display progress.
        batch_size: Size of minibatches in SGD.
        ignore_warnings: Set to True to suppress warnings.
  
      Returns:
        None.
      """
      if min(y_orig) >= 0:
        self.yscale = 'positive'
      else:
        self.yscale = 'negative'
  
      y_orig = np.array(y_orig)
      num_groups = len(set(group_feature))  # number of demographic groups
  
      if (min(y_orig) < -1 or max(y_orig) > 1) and not ignore_warnings:
        print('Warning: the scores y_orig are not in the range [-1, +1].'
              'To suppress this message, set ignore_warnings=True.')
  
      if self.yscale == 'positive' and not ignore_warnings:
        print('Warning: if this is for postprocessing a binary classifier, '
              'the scores need to be rescaled to [-1, +1]. To suppress this '
              'message, set ignore_warnings=True.')
      if min(group_feature) != 0 or (max(group_feature) != num_groups - 1):
        raise ValueError('group_feature should be in {0, 1, .. K-1} where '
                         'K is the nubmer of groups. Some groups are missing.')
  
      self.num_groups = num_groups
      eps0 = self.eps / 2.0
      gamma = self.gamma
  
      # Store group membership ids in a dictionary.
      xk_groups = {}
      for k in range(num_groups):
        xk_groups[k] = []
      for i in range(len(group_feature)):
        xk_groups[group_feature[i]].append(i)
  
      for k in xk_groups:
        assert xk_groups[k]  # All groups must be non-empty.
  
      self.avrg_y_score = float(sum(y_orig))/len(y_orig)
      if self.rho is None:
        if self.yscale == 'positive':
          self.rho = self.avrg_y_score
        else:
          self.rho = self.avrg_y_score / 2.0 + 0.5
  
      # The parameters we optimize in the algorithm are lambdas and mus.
      # lambdas_final and mus_final are running averages (final output).
      lambdas = np.zeros((num_groups,))
      mus = np.zeros((num_groups,))
      lambdas_final = np.zeros((num_groups,))  # running averages
      mus_final = np.zeros((num_groups,))  # running averages
  
      # SGD is carried out in each group separately due to decomposition of the
      # optimization problem.
      num_samples_sgd = sgd_steps * batch_size
      lr = gamma * math.sqrt(1.0 / num_samples_sgd)
  
      # Begin the projected SGD phase.
      if verbose:
        print('SGD phase started:')
      for k in range(num_groups):
        if verbose:
          print('Group %d.\t\t%02d%%'%(k, int(100*k/num_groups)), end='\r')
  
        idx = np.array(list(xk_groups[k]))  # instance IDs in group k
        group_size = len(idx)
        for _ in range(sgd_steps):
          # Using random.randint is 10x faster than random.choice.
          batch_ids = np.random.randint(0, group_size, batch_size)
          batch_ids = idx[batch_ids]
  
          # The code below is a faster implementation of:
          # xi_arg = y_orig[batch_ids] - (lambdas[k] - mus[k])
          # xi_gradient = xi_arg/gamma
          # xi_gradient = np.maximum(xi_gradient, 0.)
          # xi_gradient = np.minimum(xi_gradient, 1.)
  
          lambda_minus_mu = lambdas[k] - mus[k]
          xi_arg = np.maximum(y_orig[batch_ids], lambda_minus_mu)
          xi_arg = np.minimum(xi_arg, gamma + lambda_minus_mu)
          mean_xi = (np.mean(xi_arg) - lambda_minus_mu) / gamma
  
          lambda_gradient = eps0 + self.rho - mean_xi
          mu_gradient = eps0 - self.rho + mean_xi
  
          # stochastic gradient descent
          if eps0 > 1e-3:
            lambdas[k] = max(0, lambdas[k] - lr * batch_size * lambda_gradient)
            mus[k] = max(0, mus[k] - lr * batch_size * mu_gradient)
          else:
            # If self.eps=0, we can drop mus and optimize lambdas only but
            # lambdas will not be constrained to be non-negative in this case.
            lambdas[k] = lambdas[k] - lr * batch_size * lambda_gradient
  
          # lambdas_final and mus_final are running averages.
          lambdas_final[k] += lambdas[k] / sgd_steps
          mus_final[k] += mus[k] / sgd_steps
  
      # Now switch to full gradient descent.
      # Because the objective is smooth, lr=gamma/2 works.
      if verbose and full_gradient_epochs:
        print('\nFull gradient descent phase started:')
      for k in range(num_groups):
        if verbose:
          print('Group {}.'.format(k))
  
        idx = np.array(list(xk_groups[k]))
        for _ in range(full_gradient_epochs):
          lambda_minus_mu = lambdas_final[k] - mus_final[k]
          xi_arg = np.maximum(y_orig[idx], lambda_minus_mu)
          xi_arg = np.minimum(xi_arg, gamma + lambda_minus_mu)
          mean_xi = (np.mean(xi_arg) - lambda_minus_mu) / gamma
  
          full_grad_lambda = eps0 + self.rho - mean_xi
          full_grad_mu = eps0 - self.rho + mean_xi
  
          if eps0 > 1e-3:
            lambdas_final[k] = max(0,
                                   lambdas_final[k] - 0.5*gamma*full_grad_lambda)
            mus_final[k] = max(0, mus_final[k] - 0.5*gamma*full_grad_mu)
          else:
            lambdas_final[k] = lambdas_final[k] - 0.5*gamma*full_grad_lambda
  
      self.lambdas = lambdas_final
      self.mus = mus_final
  
    def predict(self, y_orig, group_feature, ignore_warnings=False):
      """Debiases the predictions.
  
      Given the original scores y, post-process them according to the learned
      model such that the predictions satisfy the desired fairness criteria.
  
      Args:
        y_orig: Original classifier scores. If this is for postprocessing binary
          classifiers, y_orig = 2 * p(y=1) -1.
        group_feature: An array containing the group id of each instance starting
          from group 0 to group K-1.
        ignore_warnings: Set to True to suppress warnings.
  
      Returns:
        y_new_prob: y_new_prob[i] is the probability of predicting the positive
          class for the instance i.
      """
      if (((min(y_orig) >= 0 and self.yscale == 'negative') or
           (min(y_orig) < 0 and self.yscale == 'positive')) and
          not ignore_warnings):
        print('Warning: the scores seem to have a difference scale from the '
              'training data. '
              'If the data is scaled in [0, 1], e.g. for preprocessing, or '
              'in [-1, +1], e.g. for postprocessing, make sure the test labels '
              'are scaled similarly.')
  
      num_examples = len(y_orig)  # number of training examples
      gamma = self.gamma
      lambdas = self.lambdas
      mus = self.mus
  
      y_new_prob = np.zeros((num_examples,))
      for i in range(num_examples):
        k = group_feature[i]
        if y_orig[i] < (lambdas[k]-mus[k]):
          y_new_prob[i] = 0
        elif y_orig[i] < (lambdas[k]-mus[k]) + gamma:
          y_new_prob[i] = (1.0/gamma)*(y_orig[i]-(lambdas[k]-mus[k]))
        else:
          y_new_prob[i] = 1.0
  
      return y_new_prob
   # 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.
  
  # Copyright 2021 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.
  
  """Script for analyzing the annotations of GoEmotions.
  
  The analysis includes calculating high-level statistics as well as correlation
  among emotion labels.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import os
  from absl import app
  from absl import flags
  import matplotlib.pyplot as plt
  import numpy as np
  import pandas as pd
  from scipy.cluster.hierarchy import dendrogram
  from scipy.cluster.hierarchy import linkage
  from scipy.spatial.distance import pdist
  import seaborn as sns
  
  FLAGS = flags.FLAGS
  
  flags.DEFINE_string("data", "data/full_dataset",
                      "Directory containing full dataset.")
  
  flags.DEFINE_string("plot_dir", "plots",
                      "Directory for saving plots and analyses.")
  
  flags.DEFINE_string("emotion_file", "data/emotions.txt",
                      "File containing list of emotions.")
  flags.DEFINE_string("sentiment_dict", "data/sentiment_dict.json",
                      "Sentiment dictionary.")
  
  
  def CheckAgreement(ex, min_agreement, all_emotions, max_agreement=100):
    """Return the labels that at least min_agreement raters agree on."""
    sum_ratings = ex[all_emotions].sum(axis=0)
    agreement = ((sum_ratings >= min_agreement) & (sum_ratings <= max_agreement))
    return ",".join(sum_ratings.index[agreement].tolist())
  
  
  def CountLabels(labels):
    if (not isinstance(labels, float)) and labels:
      return len(labels.split(","))
    return 0
  
  
  def main(_):
    print("Loading data...")
    dfs = []
    for filename in os.listdir(FLAGS.data):
      if filename.endswith(".csv"):
        dfs.append(
            pd.read_csv(os.path.join(FLAGS.data, filename), encoding="utf-8"))
    data = pd.concat(dfs)
    print("%d Examples" % (len(set(data["id"]))))
    print("%d Annotations" % len(data))
  
    if not os.path.isdir(FLAGS.plot_dir):
      os.makedirs(FLAGS.plot_dir)
  
    with open(FLAGS.emotion_file, "r") as f:
      all_emotions = f.read().splitlines()
    all_emotions_neutral = all_emotions + ["neutral"]
    print("%d emotion Categories" % len(all_emotions))
  
    print("%d unique raters" % len(data["rater_id"].unique()))
    print("%.3f marked unclear" %
          (data["example_very_unclear"].sum() / len(data)))
  
    # Since the ones marked as difficult have no labels, exclude those
    data = data[data[all_emotions_neutral].sum(axis=1) != 0]
  
    print("Distribution of number of labels per example:")
    print(data[all_emotions_neutral].sum(axis=1).value_counts() / len(data))
    print("%.2f with more than 3 labels" %
          ((data[all_emotions_neutral].sum(axis=1) > 3).sum() /
           len(data)))  # more than 3 labels
  
    print("Label distributions:")
    print((data[all_emotions_neutral].sum(axis=0).sort_values(ascending=False) /
           len(data) * 100).round(2))
  
    print("Plotting label correlations...")
    ratings = data.groupby("id")[all_emotions].mean()
  
    # Compute the correlation matrix
    corr = ratings.corr()
  
    # Generate a mask for the upper triangle
    mask = np.zeros_like(corr, dtype=np.bool)
    mask[np.triu_indices_from(mask)] = True
  
    # Set up the matplotlib figure
    fig, _ = plt.subplots(figsize=(11, 9))
  
    # Generate a custom diverging colormap
    cmap = sns.diverging_palette(220, 10, as_cmap=True)
  
    # Draw the heatmap with the mask and correct aspect ratio
    sns.heatmap(
        corr,
        mask=mask,
        cmap=cmap,
        vmax=.3,
        center=0,
        square=True,
        linewidths=.5,
        cbar_kws={"shrink": .5})
    fig.savefig(
        FLAGS.plot_dir + "/correlations.pdf",
        dpi=500,
        format="pdf",
        bbox_inches="tight")
  
    print("Plotting hierarchical relations...")
    z = linkage(
        pdist(ratings.T, metric="correlation"),
        method="ward",
        optimal_ordering=True)
    fig = plt.figure(figsize=(11, 4), dpi=400)
    plt.xlabel("")
    plt.ylabel("")
    dendrogram(
        z,
        labels=ratings.columns,
        leaf_rotation=90.,  # rotates the x axis labels
        leaf_font_size=12,  # font size for the x axis labels
        color_threshold=1.05,
    )
    fig.savefig(
        FLAGS.plot_dir + "/hierarchical_clustering.pdf",
        dpi=600,
        format="pdf",
        bbox_inches="tight")
  
    sent_color_map = {
        "positive": "#BEECAF",
        "negative": "#94bff5",
        "ambiguous": "#FFFC9E"
    }
    with open(FLAGS.sentiment_dict) as f:
      sent_dict = json.loads(f.read())
    sent_colors = {}
    for e in all_emotions:
      if e in sent_dict["positive"]:
        sent_colors[e] = sent_color_map["positive"]
      elif e in sent_dict["negative"]:
        sent_colors[e] = sent_color_map["negative"]
      else:
        sent_colors[e] = sent_color_map["ambiguous"]
  
    # Generate a mask for the upper triangle
    mask = np.zeros_like(corr, dtype=np.bool)
    mask[np.diag_indices(mask.shape[0])] = True
  
    # Generate a custom diverging colormap
    cmap = sns.diverging_palette(220, 10, as_cmap=True)
  
    row_colors = pd.Series(
        corr.columns, index=corr.columns, name="sentiment").map(sent_colors)
  
    # Draw the heatmap with the mask and correct aspect ratio
    g = sns.clustermap(
        corr,
        mask=mask,
        cmap=cmap,
        vmax=.3,
        vmin=-0.3,
        center=0,
        row_linkage=z,
        col_linkage=z,
        col_colors=row_colors,
        linewidths=.1,
        cbar_kws={
            "ticks": [-.3, -.15, 0, .15, .3],
            "use_gridspec": False,
            "orientation": "horizontal",
        },
        figsize=(10, 10))
  
    g.ax_row_dendrogram.set_visible(False)
    g.cax.set_position([.34, -0.05, .5, .03])
  
    for label in sent_color_map:
      g.ax_col_dendrogram.bar(
          0, 0, color=sent_color_map[label], label=label, linewidth=0)
  
    g.ax_col_dendrogram.legend(
        title="Sentiment", loc="center", bbox_to_anchor=(1.1, .5))
  
    g.savefig(FLAGS.plot_dir + "/hierarchical_corr.pdf", dpi=600, format="pdf")
  
    print("Calculating agreements...")
    unique_labels = data.groupby("id").apply(CheckAgreement, 1,
                                             all_emotions_neutral).to_dict()
    data["unique_labels"] = data["id"].map(unique_labels)
    agree_dict_2 = data.groupby("id").apply(CheckAgreement, 2,
                                            all_emotions_neutral).to_dict()
    data["agree_2"] = data["id"].map(agree_dict_2)
    agree_dict = data.groupby("id").apply(CheckAgreement, 3,
                                          all_emotions_neutral).to_dict()
    data["agree_3"] = data["id"].map(agree_dict)
    agree_dict = data.groupby("id").apply(CheckAgreement, 1, all_emotions_neutral,
                                          1).to_dict()
    data["no_agree"] = data["id"].map(agree_dict)
  
    filtered_2 = data[data["agree_2"].str.len() > 0]
    print(
        "%d (%d%%) of the examples have 2+ raters agreeing on at least one emotion label"
        % (len(filtered_2["id"].unique()), (len(filtered_2) / len(data) * 100)))
  
    filtered_3 = data[data["agree_3"].str.len() > 0]
    print(
        "%d (%d%%) of the examples have 3+ raters agreeing on at least one emotion label"
        % (len(filtered_3["id"].unique()), (len(filtered_3) / len(data) * 100)))
  
    print("Plotting number of labels...")
    data["num_unique_prefilter"] = data["unique_labels"].apply(CountLabels)
    data["num_unique_postfilter"] = data["agree_2"].apply(CountLabels)
    unique_ex = data.drop_duplicates("id")
    df = pd.DataFrame({
        "count":
            unique_ex["num_unique_prefilter"].tolist() +
            unique_ex["num_unique_postfilter"].tolist(),
        "type": ["pre-filter"] * len(unique_ex) + ["post-filter"] * len(unique_ex)
    })
  
    fig = plt.figure(dpi=600)
    ax = sns.countplot(
        data=df, x="count", hue="type", palette=["skyblue", "navy"])
    plt.xlim(-.5, 7.5)
    plt.legend(loc="center right", fontsize="x-large")
    plt.ylabel("Number of Examples", fontsize="x-large")
    plt.xlabel("Number of Labels", fontsize="x-large")
    plt.draw()
    labels = [item.get_text() for item in ax.get_yticklabels()]
    ax.set_yticklabels(["%dk" % (int(int(label) / 1000)) for label in labels])
    plt.tight_layout()
  
    fig.savefig(
        FLAGS.plot_dir + "/number_of_labels.pdf",
        dpi=600,
        format="pdf",
        bbox_inches="tight")
  
    print("Proportion of agreement per label:")
    print(
        filtered_2[all_emotions_neutral].sum(axis=0).sort_values(ascending=False)
        / len(data))
  
  
  if __name__ == "__main__":
    app.run(main)
  
apis: tensorflow.reduce_logsumexp tensorflow.keras.optimizers.Adam tensorflow.math.log tensorflow.minimum tensorflow.zeros numpy.zeros tensorflow.gather_nd tensorflow.shape tensorflow.eye tensorflow.reduce_mean tensorflow.expand_dims tensorflow.exp tensorflow.GradientTape tensorflow.cast tensorflow.argmax tensorflow.constant
apis: numpy.zeros numpy.maximum numpy.array numpy.random.randint numpy.mean numpy.minimum
apis: numpy.zeros_like scipy.spatial.distance.pdist matplotlib.pyplot.legend matplotlib.pyplot.draw scipy.cluster.hierarchy.dendrogram numpy.triu_indices_from pandas.Series matplotlib.pyplot.figure matplotlib.pyplot.tight_layout matplotlib.pyplot.subplots matplotlib.pyplot.xlim matplotlib.pyplot.ylabel pandas.concat numpy.diag_indices matplotlib.pyplot.xlabel