import numpy as np
import matplotlib.pyplot as plt

#
#   The purpose of this script is to get a feeling what the transformation from spatial to fourier domain does.
#
#   The user can set white or black pixels in the spatial domain and the fourier domain will be updated accordingly.
#   The left mouse button creates white pixels (value = 1), the right mouse button black pixels (value = 0).
#       Both one-time clicks work as well as dragging (moving while button is pressed down) the mouse.
#   The drawing radius can be changed with the + and - keys, also during dragging.
#   The "i"-key inverts the spatial data (0s become 1s and 1s become 0s), the "r"-key resets it to 0s.
#


n = 100
drawradius = 5  # pixels the user edits with one click, the clicked cell is in the middle
                # drawradius 1 means 1 pixel, 2 means 9 pixels (wrapping the clicked cell) and so on


def spatial_to_fourier():
    return np.real(np.fft.fftshift(np.fft.fft2(data_spatial)))


def fourier_to_spatial(fourier_data):  # TODO
    return  # allow pixel-manipulation in the fourier domain and change the spatial domain accordingly
            # but what editing operations make sense?


fig_spatial, ax_spatial = plt.subplots()
ax_spatial.set(title='Spatial domain')
data_spatial = np.zeros((n, n))  # np.random.random((n, n))
im_spatial = ax_spatial.imshow(data_spatial, cmap='gray', origin='lower')

fig_fourier, ax_fourier = plt.subplots()
ax_fourier.set(title='Fourier domain')
data_fourier = spatial_to_fourier()
im_fourier = ax_fourier.imshow(data_fourier, origin='lower')

row, col = 0, 0


def edit_values(value):
    for r in range(row - drawradius + 1, row + drawradius):
        for c in range(col - drawradius + 1, col + drawradius):
            if 0 <= r < n and 0 <= c < n:
                data_spatial[r, c] = value


def update_figures():
    # update spatial
    im_spatial.set_data(data_spatial)
    im_spatial.autoscale()
    fig_spatial.canvas.draw()
    # update fourier
    im_fourier.set_data(spatial_to_fourier())
    im_fourier.autoscale()  # adjust vmin/vmax automatically to new set of values
    fig_fourier.canvas.draw()


def handle_mouse_event(event):
    global row, col
    if (event.inaxes is None) or (event.button is not 1 and event.button is not 3):  # 1: left, 3: right
        return
    current_col = int(event.xdata + 0.5)
    current_row = int(event.ydata + 0.5)
    if current_col == col and current_row == row:  # still in same cell
        return
    col = current_col
    row = current_row
    edit_values(1 if event.button == 1 else 0)
    update_figures()


def handle_key_released(event):
    global data_spatial, drawradius
    if event.key == 'i':  # invert
        for r in range(0, n):
            for c in range(0, n):
                data_spatial[r, c] = 1 - data_spatial[r, c]  # toggle 1s and 0s
        update_figures()
        print 'inverted spatial data'
    if event.key == 'r':  # reset
        data_spatial = np.zeros((n, n))
        update_figures()
        print 'reset spatial data'
    if event.key == '+':
        drawradius += 1
        print 'drawradius: ' + str(drawradius)
    if event.key == '-':
        if drawradius > 1:
            drawradius -= 1
            print 'drawradius: ' + str(drawradius)


fig_spatial.canvas.mpl_connect('key_release_event', handle_key_released)
# have mouse release events as well as mouse motion events trigger handle_mouse_event(), to enable editing by both click and dragging
fig_spatial.canvas.mpl_connect('button_release_event', handle_mouse_event)
fig_spatial.canvas.mpl_connect('motion_notify_event', handle_mouse_event)

plt.show()