# Lint as: python2, python3
# Copyright 2018 The TensorFlow 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.
# ==============================================================================
"""Utilities for generating image summaries using matplotlib."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import functools
import traceback

import REDACTED.transformer_lingvo.lingvo.compat as tf
from REDACTED.transformer_lingvo.lingvo.core import py_utils
from matplotlib.backends import backend_agg
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np
import six
from six.moves import range
from six.moves import zip


def ToUnicode(text):
  if not isinstance(text, six.text_type):
    text = six.ensure_text(text, 'utf-8')
  return text


def AddPlot(unused_fig,
            axes,
            data,
            title=u'',
            xlabel=u'',
            ylabel=u'',
            fontsize='small',
            xlim=None,
            ylim=None,
            suppress_xticks=False,
            suppress_yticks=False):
  """Convenience function to add a plot."""
  axes.plot(data)
  axes.set_title(ToUnicode(title), size=fontsize)
  axes.set_xlabel(ToUnicode(xlabel), size=fontsize)
  axes.set_ylabel(ToUnicode(ylabel), size=fontsize)
  if xlim:
    axes.set_xlim(xlim)
  if ylim:
    axes.set_ylim(ylim)
  if suppress_xticks:
    axes.set_xticks([])
  if suppress_yticks:
    axes.set_yticks([])


def AddImage(fig,
             axes,
             data,
             cmap='bone_r',
             clim=None,
             show_colorbar=True,
             title=u'',
             xlabel=u'',
             ylabel=u'',
             fontsize='small',
             origin='lower',
             suppress_xticks=False,
             suppress_yticks=False,
             aspect='auto',
             vmin=None,
             vmax=None):
  """Convenience function to plot data as an image on the given axes."""
  image = axes.imshow(
      data,
      cmap=cmap,
      origin=origin,
      aspect=aspect,
      interpolation='nearest',
      vmin=vmin,
      vmax=vmax)
  if show_colorbar:
    fig.colorbar(image)
  if clim is not None:
    image.set_clim(clim)
  axes.set_title(ToUnicode(title), size=fontsize)
  axes.set_xlabel(ToUnicode(xlabel), size=fontsize)
  axes.set_ylabel(ToUnicode(ylabel), size=fontsize)
  if suppress_xticks:
    axes.set_xticks([])
  if suppress_yticks:
    axes.set_yticks([])


def AddScatterPlot(unused_fig,
                   axes,
                   xs,
                   ys,
                   title=u'',
                   xlabel=u'',
                   ylabel=u'',
                   fontsize='small',
                   xlim=None,
                   ylim=None,
                   suppress_xticks=False,
                   suppress_yticks=False,