import numpy as np
import matplotlib.pyplot as plt

####### dir ############
SALIENCY_DIR = f'./../train_output/train_n0.001_run4_t2_b50_b400_b800_b1400_adc/saliency.npy'

########## load saliency map ############
heatmap = np.load(SALIENCY_DIR)
heatmap = np.squeeze(heatmap)
#take one image out
heatmap = np.squeeze(abs(heatmap))
max_value = np.amax(heatmap)
min_value = np.amin(heatmap)


class IndexTracker:
    def __init__(self, ax, X):
        self.ax = ax
        ax.set_title('use scroll wheel to navigate images')

        self.X = X
        rows, cols, self.slices = X.shape
        self.ind = self.slices//2

        self.im = ax.imshow(self.X[:, :, self.ind], cmap='jet',vmin=min_value, vmax=max_value)
        self.update()

    def on_scroll(self, event):
        print("%s %s" % (event.button, event.step))
        if event.button == 'up':
            self.ind = (self.ind + 1) % self.slices
        else:
            self.ind = (self.ind - 1) % self.slices
        self.update()

    def update(self):
        self.im.set_data(self.X[:, :, self.ind])
        self.ax.set_ylabel('slice %s' % self.ind)
        self.im.axes.figure.canvas.draw()

# plt.figure()
# fig, ax = plt.subplots(1, 1)
# tracker = IndexTracker(ax, heatmap[:,:,:,0])
# fig.canvas.mpl_connect('scroll_event', tracker.on_scroll)
# plt.show()

plt.figure()
fig, ax = plt.subplots(1, 1)
tracker = IndexTracker(ax, heatmap[:,:,:,4])
fig.canvas.mpl_connect('scroll_event', tracker.on_scroll)
plt.show()