text: # Copyright 2017 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.
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from collections import deque
  from math import floor, sqrt
  
  import numpy as np
  import tensorflow as tf
  
  import umap
  from tensorboard.plugins.uMAP import im_util
  from tensorboard.plugins.uMAP.shared_config import SECTION_HEIGHT,\
    IMAGE_WIDTH, DEFAULT_CONFIG, SECTION_INFO_FILENAME
  from tensorboard.plugins.uMAP.file_system_tools import write_pickle
  
  MIN_SQUARE_SIZE = 3
  
  
  class Visualizer(object):
  
    def __init__(self, logdir):
      self.logdir = logdir
      self.sections_over_time = deque([], DEFAULT_CONFIG['window_size'])
      self.config = dict(DEFAULT_CONFIG)
      self.old_config = dict(DEFAULT_CONFIG)
  
  
    def _reshape_conv_array(self, array, section_height, image_width):
      '''Reshape a rank 4 array to be rank 2, where each column of block_width is
      a filter, and each row of block height is an input channel. For example:
  
      [[[[ 11,  21,  31,  41],
         [ 51,  61,  71,  81],
         [ 91, 101, 111, 121]],
        [[ 12,  22,  32,  42],
         [ 52,  62,  72,  82],
         [ 92, 102, 112, 122]],
        [[ 13,  23,  33,  43],
         [ 53,  63,  73,  83],
         [ 93, 103, 113, 123]]],
       [[[ 14,  24,  34,  44],
         [ 54,  64,  74,  84],
         [ 94, 104, 114, 124]],
        [[ 15,  25,  35,  45],
         [ 55,  65,  75,  85],
         [ 95, 105, 115, 125]],
        [[ 16,  26,  36,  46],
         [ 56,  66,  76,  86],
         [ 96, 106, 116, 126]]],
       [[[ 17,  27,  37,  47],
         [ 57,  67,  77,  87],
         [ 97, 107, 117, 127]],
        [[ 18,  28,  38,  48],
         [ 58,  68,  78,  88],
         [ 98, 108, 118, 128]],
        [[ 19,  29,  39,  49],
         [ 59,  69,  79,  89],
         [ 99, 109, 119, 129]]]]
  
         should be reshaped to:
  
         [[ 11,  12,  13,  21,  22,  23,  31,  32,  33,  41,  42,  43],
          [ 14,  15,  16,  24,  25,  26,  34,  35,  36,  44,  45,  46],
          [ 17,  18,  19,  27,  28,  29,  37,  38,  39,  47,  48,  49],
          [ 51,  52,  53,  61,  62,  63,  71,  72,  73,  81,  82,  83],
          [ 54,  55,  56,  64,  65,  66,  74,  75,  76,  84,  85,  86],
          [ 57,  58,  59,  67,  68,  69,  77,  78,  79,  87,  88,  89],
          [ 91,  92,  93, 101, 102, 103, 111, 112, 113, 121, 122, 123],
          [ 94,  95,  96, 104, 105, 106, 114, 115, 116, 124, 125, 126],
          [ 97,  98,  99, 107, 108, 109, 117, 118, 119, 127, 128, 129]]
      '''
  
      # E.g. [100, 24, 24, 10]: this shouldn't be reshaped like normal.
      if array.shape[1] == array.shape[2] and array.shape[0] != array.shape[1]:
        array = np.rollaxis(np.rollaxis(array, 2), 2)
  
      block_height, block_width, in_channels = array.shape[:3]
      rows = []
  
      max_element_count = section_height * int(image_width / MIN_SQUARE_SIZE)
      element_count = 0
  
      for i in range(in_channels):
        rows.append(array[:, :, i, :].reshape(block_height, -1, order='F'))
  
        # This line should be left in this position. Gives it one extra row.
        if element_count >= max_element_count and not self.config['show_all']:
          break
  
        element_count += block_height * in_channels * block_width
  
      return np.vstack(rows)
  
  
    def _reshape_irregular_array(self, array, section_height, image_width):
      '''Reshapes arrays of ranks not in {1, 2, 4}
      '''
      section_area = section_height * image_width
      flattened_array = np.ravel(array)
  
      if not self.config['show_all']:
        flattened_array = flattened_array[:int(section_area/MIN_SQUARE_SIZE)]
  
      cell_count = np.prod(flattened_array.shape)
      cell_area = section_area / cell_count
  
      cell_side_length = max(1, floor(sqrt(cell_area)))
      row_count = max(1, int(section_height / cell_side_length))
      col_count = int(cell_count / row_count)
  
      # Reshape the truncated array so that it has the same aspect ratio as
      # the section.
  
      # Truncate whatever remaining values there are that don't fit. Hopefully
      # it doesn't matter that the last few (< section count) aren't there.
      section = np.reshape(flattened_array[:row_count * col_count],
                           (row_count, col_count))
  
      return section
  
  
    def _determine_image_width(self, arrays, show_all):
      final_width = IMAGE_WIDTH
  
      if show_all:
        for array in arrays:
          rank = len(array.shape)
  
          if rank == 1:
            width = len(array)
          elif rank == 2:
            width = array.shape[1]
          elif rank == 4:
            width = array.shape[1] * array.shape[3]
          else:
            width = IMAGE_WIDTH
  
          if width > final_width:
            final_width = width
  
      return final_width
  
  
    def _determine_section_height(self, array, show_all):
      rank = len(array.shape)
      height = SECTION_HEIGHT
  
      if show_all:
        if rank == 1:
          height = SECTION_HEIGHT
        if rank == 2:
          height = max(SECTION_HEIGHT, array.shape[0])
        elif rank == 4:
          height = max(SECTION_HEIGHT, array.shape[0] * array.shape[2])
        else:
          height = max(SECTION_HEIGHT, np.prod(array.shape) // IMAGE_WIDTH)
  
      return height
  
  
    def _arrays_to_sections(self, arrays):
      '''
      input: unprocessed numpy arrays.
      returns: columns of the size that they will appear in the image, not scaled
               for display. That needs to wait until after variance is computed.
      '''
      sections = []
      sections_to_resize_later = {}
      show_all = self.config['show_all']
      image_width = self._determine_image_width(arrays, show_all)
  
      for array_number, array in enumerate(arrays):
        rank = len(array.shape)
        section_height = self._determine_section_height(array, show_all)
  
        if rank == 1:
          section = np.atleast_2d(array)
        elif rank == 2:
          section = array
        elif rank == 4:
          section = self._reshape_conv_array(array, section_height, image_width)
        else:
          section = self._reshape_irregular_array(array,
                                                  section_height,
                                                  image_width)
        # Only calculate variance for what we have to. In some cases (biases),
        # the section is larger than the array, so we don't want to calculate
        # variance for the same value over and over - better to resize later.
        # About a 6-7x speedup for a big network with a big variance window.
        section_size = section_height * image_width
        array_size = np.prod(array.shape)
  
        if section_size > array_size:
          sections.append(section)
          sections_to_resize_later[array_number] = section_height
        else:
          sections.append(im_util.resize(section, section_height, image_width))
  
      self.sections_over_time.append(sections)
  
      if self.config['mode'] == 'variance':
        sections = self._sections_to_variance_sections(self.sections_over_time)
  
      for array_number, height in sections_to_resize_later.items():
        sections[array_number] = im_util.resize(sections[array_number],
                                                height,
                                                image_width)
      return sections
  
  
    def _sections_to_variance_sections(self, sections_over_time):
      '''Computes the variance of corresponding sections over time.
  
      Returns:
        a list of np arrays.
      '''
      variance_sections = []
  
      for i in range(len(sections_over_time[0])):
        time_sections = [sections[i] for sections in sections_over_time]
        variance = np.var(time_sections, axis=0)
        variance_sections.append(variance)
  
      return variance_sections
  
  
    def _sections_to_image(self, sections):
      padding_size = 5
  
      sections = im_util.scale_sections(sections, self.config['scaling'])
  
      final_stack = [sections[0]]
      padding = np.zeros((padding_size, sections[0].shape[1]))
  
      for section in sections[1:]:
        final_stack.append(padding)
        final_stack.append(section)
  
      return np.vstack(final_stack).astype(np.uint8)
  
  
    def _maybe_clear_deque(self):
      '''Clears the deque if certain parts of the config have changed.'''
  
      for config_item in ['values', 'mode', 'show_all']:
        if self.config[config_item] != self.old_config[config_item]:
          self.sections_over_time.clear()
          break
  
      self.old_config = self.config
  
      window_size = self.config['window_size']
      if window_size != self.sections_over_time.maxlen:
        self.sections_over_time = deque(self.sections_over_time, window_size)
  
  
    def _save_section_info(self, arrays, sections):
      infos = []
  
      if self.config['values'] == 'trainable_variables':
        names = [x.name for x in tf.trainable_variables()]
      else:
        names = range(len(arrays))
  
      for array, section, name in zip(arrays, sections, names):
        info = {}
  
        info['name'] = name
        info['shape'] = str(array.shape)
        info['min'] = '{:.3e}'.format(section.min())
        info['mean'] = '{:.3e}'.format(section.mean())
        info['max'] = '{:.3e}'.format(section.max())
        info['range'] = '{:.3e}'.format(section.max() - section.min())
        info['height'] = section.shape[0]
  
        infos.append(info)
  
      write_pickle(infos, '{}/{}'.format(self.logdir, SECTION_INFO_FILENAME))
  
  
    def build_frame(self, arrays):
      self._maybe_clear_deque()
  
      arrays = arrays if isinstance(arrays, list) else [arrays]
      reducer = umap.UMAP(random_state=42)
      arrays = [reducer.fit_transform(a) for a in arrays] if isinstance(arrays, list) else [reducer.fit_transform(arrays)]
      sections = self._arrays_to_sections(arrays)
      self._save_section_info(arrays, sections)
      final_image = self._sections_to_image(sections)
      final_image = im_util.apply_colormap(final_image, self.config['colormap'])
  
      return final_image
  
    def update(self, config):
      self.config = config
  