import sys
import numpy as np

# Only ask users to install matplotlib if they actually need it
try:
    import matplotlib.pyplot as plt
except:
    print('To display the environment in a window, please install matplotlib, eg:')
    print('pip3 install --user matplotlib')
    sys.exit(-1)

class Window:
    """
    Window to draw a gridworld instance using Matplotlib
    """

    def __init__(self, title, figsize=(3, 3)):
        self.fig = None

        self.imshow_obj = None

        # Create the figure and axes
        self.fig, self.ax = plt.subplots(
            # figsize=(10, 5),
            figsize=figsize,
        )

        # Show the env name in the window title
        self.fig.canvas.set_window_title(title)

        # Turn off x/y axis numbering/ticks
        self.ax.xaxis.set_ticks_position('none')
        self.ax.yaxis.set_ticks_position('none')
        _ = self.ax.set_xticklabels([])
        _ = self.ax.set_yticklabels([])

        # list of text handles
        self.txt_handles = []

        # Flag indicating the window was closed
        self.closed = False

        def close_handler(evt):
            self.closed = True

        self.fig.canvas.mpl_connect('close_event', close_handler)

    def show_img(self, img):
        """
        Show an image or update the image being shown
        """

        # Show the first image of the environment
        if self.imshow_obj is None:
            self.imshow_obj = self.ax.imshow(img, interpolation='bilinear')

        self.imshow_obj.set_data(img)
        self.fig.canvas.draw()

        # Let matplotlib process UI events
        # This is needed for interactive mode to work properly
        # plt.pause(0.001)

    def set_caption(self, text, relevant_set=None):
        """
        Set/update the caption text below the image
        """

        # plt.xlabel(text)
        # text = "All utterances:\n\n"+text
        lines = text.split("\n")

        if len(lines) > 8:
            lines = ["..."]+lines[-8:]

        text = "\n".join(lines)

        if hasattr(self, "caption"):
            self.caption.set_text(text)
        else:
            # self.caption = plt.text(400, 250, text, ha="left",wrap=True)
            self.caption = plt.text(330, 250, text, ha="left", wrap=True)

        if relevant_set is not None:
            # if a line in the text has one of these strings it will be put in the relevant set

            relevant_lines = ["Relevant utterances:\n"] + [
                l for l in text.rsplit("\n") if any([r in l for r in relevant_set])
            ] + ["\n"]
            relevant_text = "\n".join(relevant_lines)


            if hasattr(self, "relevant_caption"):
                self.relevant_caption.set_text(relevant_text)
            else:
                self.relevant_caption = plt.text(-200, 250, relevant_text, ha="left")


    def reg_key_handler(self, key_handler):
        """
        Register a keyboard event handler
        """

        # Keyboard handler
        self.fig.canvas.mpl_connect('key_press_event', key_handler)

    def show(self, block=True):
        """
        Show the window, and start an event loop
        """

        # If not blocking, trigger interactive mode
        if not block:
            plt.ion()

        # Show the plot
        # In non-interative mode, this enters the matplotlib event loop
        # In interactive mode, this call does not block
        plt.show()

    def close(self):
        """
        Close the window
        """

        plt.close()

    def add_text(self, *args, **kwargs):

        kwargs['transform'] = self.ax.transAxes
        self.txt_handles.append(self.ax.text(*args, **kwargs))

    def clear_text(self):

        if len(self.txt_handles) > 0:
            while len(self.txt_handles) > 0:
                self.txt_handles.pop().remove()