import sys

if len(sys.argv) == 2:

    if sys.argv[1] == '-h':

        print("python "+sys.argv[0]+" <path to dir of scans>")

        print("python "+sys.argv[0]+" <path to dir of scans> <path to save dir>")

        exit(1)

import matplotlib.pyplot as plt

from matplotlib.widgets import Button

from procedures.attack_pipeline import *

from utils.equalizer import *

import matplotlib.animation as animation

import time

class GUI(object):

    # If load_path is to a *.dcm or *.mhd file then only this scan is loaded

    # If load_path is to a directory, then all scans are loaded. It is assumed that each scan is in its own subdirectory.

    # save_path is the directory to save the tampered scans (as dicom)

    def __init__(self, load_path, save_path=None):

        # init manipulator

        self.savepath = save_path

        self.filepaths = self._load_paths(load_path)  # load all scans filepaths in path

        self.fileindex = 0

        self.manipulator = scan_manipulator()

        self.manipulator.load_target_scan(self.filepaths[self.fileindex])

        self.hist_state = True

        self.inject_coords = []

        self.remove_coords = []

        # init plot

        self.eq = histEq(self.manipulator.scan)

        self.slices, self.cols, self.rows = self.manipulator.scan.shape

        self.ind = self.slices // 2

        self.pause_start = 0

        self.fig, self.ax = plt.subplots(1, 1, dpi=100)

        self.fig.suptitle('CT-GAN: Malicious Tampering of 3D Medical Imagery using Deep Learning\nTool by Yisroel Mirsky', fontsize=14, fontweight='bold')

        plt.subplots_adjust(bottom=0.2)

        self.ani_direction = 'down'

        self.animation = None

        self.animation_state = True

        self.plot()

        self.ax.set_title(os.path.split(self.filepaths[self.fileindex])[-1]) #filename

        # register click/scroll events

        self.action_state = 'inject' #default state

        self.fig.canvas.mpl_connect('button_press_event', self.onclick)

        self.fig.canvas.mpl_connect('scroll_event', self.onscroll)

        # register buttons

        axanim = plt.axes([0.1, 0.21, 0.2, 0.075])

        self.banim = Button(axanim, 'Toggle Animation')

        self.banim.on_clicked(self.toggle_animation)

        axinj = plt.axes([0.1, 0.05, 0.1, 0.075])

        axrem = plt.axes([0.21, 0.05, 0.1, 0.075])

        self.binj = Button(axinj, 'Inject')

        self.binj.on_clicked(self.inj_on)

        self.brem = Button(axrem, 'Remove')

        self.brem.on_clicked(self.rem_on)

        axhist = plt.axes([0.35, 0.05, 0.2, 0.075])

        self.bhist = Button(axhist, 'Toggle HistEQ')

        self.bhist.on_clicked(self.hist)

        axprev = plt.axes([0.59, 0.05, 0.1, 0.075])

        axsave = plt.axes([0.7, 0.05, 0.1, 0.075])

        axnext = plt.axes([0.81, 0.05, 0.1, 0.075])

        self.bnext = Button(axnext, 'Next')

        self.bnext.on_clicked(self.next)

        self.bprev = Button(axprev, 'Previous')

        self.bprev.on_clicked(self.prev)

        self.bsave = Button(axsave, 'Save')

        self.bsave.on_clicked(self.save)

        self.maximize_window()

        self.update()

        plt.show()

    def _load_paths(self,path):

        filepaths = []

        # load single scan?

        if (path.split('.')[-1] == "dcm") or (path.split('.')[-1] == "mhd"):

            filepaths.append(path)

            return filepaths

        # try load directory of scans...

        files = os.listdir(path)

        for file in files:

            if os.path.isdir(file):

                subdir = os.path.join(path,file)

                subdir_files = os.listdir(subdir)

                if subdir_files[0].split('.')[-1] == "dcm": #folder contains dicom

                    filepaths.append(os.path.join(path,subdir))

                elif (subdir_files[0].split('.')[-1] == "mhd") or (subdir_files[0].split('.')[-1] == "raw"): # MHD

                    filepaths.append(os.path.join(path,subdir,subdir_files[0]))

            elif file.split('.')[-1] == "mhd":

                filepaths.append(os.path.join(path,file))

        return filepaths

    def onclick(self, event):

        # print('%s click: button=%d, x=%d, y=%d, xdata=%f, ydata=%f' %

        #       ('double' if event.dblclick else 'single', event.button,

        #        event.x, event.y, event.xdata, event.ydata))

        if event.xdata is not None:

            coord = np.array([self.ind,event.ydata,event.xdata],dtype=int)

            if coord[1] > 0 and coord[2] > 0:

                self.pause_start = np.Inf #pause while working

                if self.action_state == 'inject':

                    self.ax.set_title("Injecting...")

                    self.im.axes.figure.canvas.draw()

                    self.manipulator.tamper(coord, action='inject', isVox=True)

                    self.inject_coords.append(coord)

                else:

                    self.ax.set_title("Removing...")

                    self.im.axes.figure.canvas.draw()

                    self.manipulator.tamper(coord, action='remove', isVox=True)

                    self.remove_coords.append(coord)

                self.pause_start = time.time() #pause few secs to see result before continue

                self.ax.set_title(os.path.split(self.filepaths[self.fileindex])[-1])  # filename

                self.update()

    def onscroll(self, event):

        if event.button == 'up':

            self.ind = (self.ind + 1) % self.slices

        else:

            self.ind = (self.ind - 1) % self.slices

        self.update()

    def toggle_animation(self, event):

        self.animation_state = not self.animation_state

        if self.animation_state:

            self.pause_start = 0

        else:

            self.pause_start = np.Inf

    def inj_on(self, event):

        self.action_state = 'inject'

    def rem_on(self, event):

        self.action_state = 'remove'

    def hist(self, event):

        self.hist_state = not self.hist_state

        self.plot()

        self.update()

    def next(self, event):

        self.fileindex = (self.fileindex + 1) % len(self.filepaths)

        self.loadscan(self.fileindex)

    def prev(self, event):

        self.fileindex = (self.fileindex - 1) % len(self.filepaths)

        self.loadscan(self.fileindex)

    def save(self, event):

        if self.savepath is not None:

            self.ax.set_title("Saving...")

            self.im.axes.figure.canvas.draw()

            uuid = os.path.split(self.filepaths[self.fileindex])[-1][:-4]

            #save scan

            self.manipulator.save_tampered_scan(os.path.join(self.savepath,uuid),output_type='dicom')

            #save coords

            file_exists = False

            if os.path.exists(os.path.join(self.savepath,"tamper_coordinates.csv")):

                file_exists = True

            f = open(os.path.join(self.savepath,"tamper_coordinates.csv"),"a+")

            load_filename = os.path.split(self.filepaths[self.fileindex])[-1]

            if not file_exists:

                f.write("filename, x, y, z, tamper_type\n") #header

            for coord in self.inject_coords:

                f.write(load_filename+", "+str(coord[2])+", "+str(coord[1])+", "+str(coord[0])+", "+"inject\n")

            for coord in self.remove_coords:

                f.write(load_filename+", "+str(coord[2])+", "+str(coord[1])+", "+str(coord[0])+", "+"remove\n")

            f.close()

            self.ax.set_title(load_filename)  # filename

            self.im.axes.figure.canvas.draw()

    def update(self):

        if self.hist_state:

            self.im.set_data(self.eq.equalize(self.manipulator.scan[self.ind,:,:]))

        else:

            self.im.set_data(self.manipulator.scan[self.ind,:,:])

        self.ax.set_ylabel('slice %s' % self.ind)

        self.im.axes.figure.canvas.draw()

    def loadscan(self,fileindex):

        #load screen

        self.im.set_data(np.ones((self.cols,self.rows))*-1000)

        self.ax.set_title("Loading...")

        self.im.axes.figure.canvas.draw()

        self.remove_coords.clear()

        self.inject_coords.clear()

        #load scan

        self.manipulator.load_target_scan(self.filepaths[fileindex])

        self.slices, self.cols, self.rows = self.manipulator.scan.shape

        self.ind = self.slices//2

        self.ax.clear()

        self.eq = histEq(self.manipulator.scan)

        self.plot()

        self.ax.set_title(os.path.split(self.filepaths[fileindex])[-1]) #filename

        self.ax.set_ylabel('slice %s' % self.ind)

        self.im.axes.figure.canvas.draw()

    def plot(self):

        self.ax.clear()

        if self.hist_state:

            self.im = self.ax.imshow(self.eq.equalize(self.manipulator.scan[self.ind,:,:]),cmap="bone")#, cmap="bone", vmin=-1000, vmax=1750)

        else:

            self.im = self.ax.imshow(self.manipulator.scan[self.ind,:,:], cmap="bone", vmin=-1000, vmax=1750)

        self.animation = animation.FuncAnimation(self.fig, self.animate, interval=100)

    def animate(self,i):

        if self.animation_state:

            if time.time() - self.pause_start > 1:

                if self.ind == self.slices-1:

                    self.ani_direction = 'up'

                elif self.ind == 0:

                    self.ani_direction = 'down'

                if self.ani_direction == 'up':

                    self.ind-=1

                else:

                    self.ind+=1

                self.update()

    def maximize_window(self):

        try: #'QT4Agg'

            figManager = plt.get_current_fig_manager()

            figManager.window.showMaximized()

        except:

            try: #'TkAgg'

                mng = plt.get_current_fig_manager()

                mng.window.state('zoomed')

            except:

                try: #'wxAgg'

                    mng = plt.get_current_fig_manager()

                    mng.frame.Maximize(True)

                except:

                    print("Could not maximize window")

if (len(sys.argv) == 1) or (len(sys.argv) > 3):

    loadpath = "data\\healthy_scans"

    savepath = "data\\tampered_scans"

if len(sys.argv) == 2:

    loadpath = sys.argv[1]

    savepath = "data\\tampered_scans"

if len(sys.argv) == 3:

    loadpath = sys.argv[1]

    savepath = sys.argv[2]

gui = GUI(load_path=loadpath,save_path=savepath)