# Author: Yuke Zhu
# Modified By: Kaichun Mo

import pyglet
import scipy.misc

class SimpleImageViewer(object):

  def __init__(self, display=None):
    self.window = None
    self.isopen = False
    self.display = display

  def imshow(self, arr):
    if self.window is None:
      height, width, channels = arr.shape
      if height > 512 or width > 512:
          height = 512
          width = width * 512 / height
          arr = scipy.misc.imresize(arr, [height, width])

      self.window = pyglet.window.Window(width=width, height=height, display=self.display, caption="Dataset Viewer")
      self.width = width
      self.height = height
      self.isopen = True

    height, width, channels = arr.shape
    if height > 512 or width > 512:
        arr = scipy.misc.imresize(arr, [self.height, self.width])

    assert arr.shape == (self.height, self.width, 3), "You passed in an image with the wrong number shape"
    image = pyglet.image.ImageData(self.width, self.height, 'RGB', arr.tobytes(), pitch=self.width * -3)
    self.window.clear()
    self.window.switch_to()
    self.window.dispatch_events()
    image.blit(0,0)
    self.window.flip()

  def close(self):
    if self.isopen:
      self.window.close()
      self.isopen = False

  def __del__(self):
    self.close()
