# Path: disentanglement_lib/evaluation/metrics/irs.py
# def scalable_disentanglement_score(gen_factors, latents, diff_quantile=0.99):
#   """Computes IRS scores of a dataset.
# 
#   Assumes no noise in X and crossed generative factors (i.e. one sample per
#   combination of gen_factors). Assumes each g_i is an equally probable
#   realization of g_i and all g_i are independent.
# 
#   Args:
#     gen_factors: Numpy array of shape (num samples, num generative factors),
#       matrix of ground truth generative factors.
#     latents: Numpy array of shape (num samples, num latent dimensions), matrix
#       of latent variables.
#     diff_quantile: Float value between 0 and 1 to decide what quantile of diffs
#       to select (use 1.0 for the version in the paper).
# 
#   Returns:
#     Dictionary with IRS scores.
#   """
#   num_gen = gen_factors.shape[1]
#   num_lat = latents.shape[1]
# 
#   # Compute normalizer.
#   max_deviations = np.max(np.abs(latents - latents.mean(axis=0)), axis=0)
#   cum_deviations = np.zeros([num_lat, num_gen])
#   for i in range(num_gen):
#     unique_factors = np.unique(gen_factors[:, i], axis=0)
#     assert unique_factors.ndim == 1
#     num_distinct_factors = unique_factors.shape[0]
#     for k in range(num_distinct_factors):
#       # Compute E[Z | g_i].
#       match = gen_factors[:, i] == unique_factors[k]
#       e_loc = np.mean(latents[match, :], axis=0)
# 
#       # Difference of each value within that group of constant g_i to its mean.
#       diffs = np.abs(latents[match, :] - e_loc)
#       max_diffs = np.percentile(diffs, q=diff_quantile*100, axis=0)
#       cum_deviations[:, i] += max_diffs
#     cum_deviations[:, i] /= num_distinct_factors
#   # Normalize value of each latent dimension with its maximal deviation.
#   normalized_deviations = cum_deviations / max_deviations[:, np.newaxis]
#   irs_matrix = 1.0 - normalized_deviations
#   disentanglement_scores = irs_matrix.max(axis=1)
#   if np.sum(max_deviations) > 0.0:
#     avg_score = np.average(disentanglement_scores, weights=max_deviations)
#   else:
#     avg_score = np.mean(disentanglement_scores)
# 
#   parents = irs_matrix.argmax(axis=1)
#   score_dict = {}
#   score_dict["disentanglement_scores"] = disentanglement_scores
#   score_dict["avg_score"] = avg_score
#   score_dict["parents"] = parents
#   score_dict["IRS_matrix"] = irs_matrix
#   score_dict["max_deviations"] = max_deviations
#   return score_dict


import os
import matplotlib
import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
import numpy as np
import tensorflow.compat.v1 as tf
from disentanglement_lib.evaluation.metrics.irs import scalable_disentanglement_score
# coding=utf-8
# Copyright 2018 The DisentanglementLib Authors.  All rights reserved.
#
# 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.

"""Visualization code for Interventional Robustness Score.

Based on the paper https://arxiv.org/abs/1811.00007.
"""

matplotlib.use("Agg")  # Set headless-friendly backend.


def vis_all_interventional_effects(gen_factors, latents, output_dir):
  """Compute Matrix of all interventional effects."""

# Path: disentanglement_lib/evaluation/metrics/irs.py
# def scalable_disentanglement_score(gen_factors, latents, diff_quantile=0.99):
#   """Computes IRS scores of a dataset.
# 
#   Assumes no noise in X and crossed generative factors (i.e. one sample per
#   combination of gen_factors). Assumes each g_i is an equally probable
#   realization of g_i and all g_i are independent.
# 
#   Args:
#     gen_factors: Numpy array of shape (num samples, num generative factors),
#       matrix of ground truth generative factors.
#     latents: Numpy array of shape (num samples, num latent dimensions), matrix
#       of latent variables.
#     diff_quantile: Float value between 0 and 1 to decide what quantile of diffs
#       to select (use 1.0 for the version in the paper).
# 
#   Returns:
#     Dictionary with IRS scores.
#   """
#   num_gen = gen_factors.shape[1]
#   num_lat = latents.shape[1]
# 
#   # Compute normalizer.
#   max_deviations = np.max(np.abs(latents - latents.mean(axis=0)), axis=0)
#   cum_deviations = np.zeros([num_lat, num_gen])
#   for i in range(num_gen):
#     unique_factors = np.unique(gen_factors[:, i], axis=0)
#     assert unique_factors.ndim == 1
#     num_distinct_factors = unique_factors.shape[0]
#     for k in range(num_distinct_factors):
#       # Compute E[Z | g_i].
#       match = gen_factors[:, i] == unique_factors[k]
#       e_loc = np.mean(latents[match, :], axis=0)
# 
#       # Difference of each value within that group of constant g_i to its mean.
#       diffs = np.abs(latents[match, :] - e_loc)
#       max_diffs = np.percentile(diffs, q=diff_quantile*100, axis=0)
#       cum_deviations[:, i] += max_diffs
#     cum_deviations[:, i] /= num_distinct_factors
#   # Normalize value of each latent dimension with its maximal deviation.
#   normalized_deviations = cum_deviations / max_deviations[:, np.newaxis]
#   irs_matrix = 1.0 - normalized_deviations
#   disentanglement_scores = irs_matrix.max(axis=1)
#   if np.sum(max_deviations) > 0.0:
#     avg_score = np.average(disentanglement_scores, weights=max_deviations)
#   else:
#     avg_score = np.mean(disentanglement_scores)
# 
#   parents = irs_matrix.argmax(axis=1)
#   score_dict = {}
#   score_dict["disentanglement_scores"] = disentanglement_scores
#   score_dict["avg_score"] = avg_score
#   score_dict["parents"] = parents
#   score_dict["IRS_matrix"] = irs_matrix
#   score_dict["max_deviations"] = max_deviations
#   return score_dict


# Path: disentanglement_lib/visualize/visualize_irs.py
import os
import matplotlib
import matplotlib.pyplot as plt  # pylint: disable=g-import-not-at-top
import numpy as np
import tensorflow.compat.v1 as tf
from disentanglement_lib.evaluation.metrics.irs import scalable_disentanglement_score

# coding=utf-8
# Copyright 2018 The DisentanglementLib Authors.  All rights reserved.
#
# 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.

"""Visualization code for Interventional Robustness Score.

Based on the paper https://arxiv.org/abs/1811.00007.
"""

matplotlib.use("Agg")  # Set headless-friendly backend.


def vis_all_interventional_effects(gen_factors, latents, output_dir):
  """Compute Matrix of all interventional effects."""

  res = scalable_disentanglement_score(gen_factors, latents)