moulder/interactive.py

from __future__ import division, absolute_import
from future.builtins import super, zip

try:
    import cPickle as pickle
except ImportError:
    import pickle

import os
import sys
import numpy
import matplotlib
from matplotlib import pyplot, widgets, patches
from matplotlib.lines import Line2D
from matplotlib.figure import Figure
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QSizePolicy, QMainWindow, QApplication, QAction
from PyQt5.QtWidgets import QMenu, QWidget, QVBoxLayout, QMessageBox
from PyQt5.QtWidgets import QSlider, QHBoxLayout, QLabel, QDialog
from PyQt5.QtWidgets import QDialogButtonBox

from fatiando import utils
from fatiando.gravmag import talwani
from fatiando.mesher import Polygon

LINE_ARGS = dict(
    linewidth=2, linestyle='-', color='k', marker='o',
    markerfacecolor='k', markersize=5, animated=False, alpha=0.6)


class Moulder(FigureCanvasQTAgg):

    # The tolerance range for mouse clicks on vertices. In pixels.
    epsilon = 5
    # App instructions printed in the figure suptitle
    instructions = ' | '.join([
        'n: New polygon', 'd: delete', 'click: select/move', 'a: add vertex',
        'r: reset view', 'esc: cancel'])

    def __init__(self, parent, x, z, min_depth, max_depth,
                 density_range=[-2000, 2000], width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)

        self.min_depth, self.max_depth = min_depth, max_depth
        self._x, self._z = x, z
        self.density_range = density_range
        self._predicted = numpy.zeros_like(x)
        self._data = None
        self.cmap = pyplot.cm.RdBu_r
        self.canvas = self.fig.canvas

        self.polygons = []
        self.lines = []
        self.densities = []

        # Set arbitrary density and error values (only for first implementations)
        # It will be determined by sliders/entries in MoulderApp
        self._density = 0
        self._error = 0

        # Data min and max (only for first implementations)
        # They will be determined when data is imported
        self.dmin, self.dmax = 0, 0

        self._figure_setup()
        self._init_markers()
        self._connect()

    @property
    def x(self):
        return self._x

    @x.setter
    def x(self, new_value):
        self._x = numpy.asarray(new_value)

    @property
    def z(self):
        return self._z

    @z.setter
    def z(self, new_value):
        self._z = numpy.asarray(new_value)

    @property
    def data(self):
        return self._data

    @data.setter
    def data(self, new_data):
        self._data = new_data

    @property
    def density(self):
        return self._density

    @density.setter
    def density(self, value):
        """
        Callback when density slider is edited
        """
        self._density = value
        if self._ipoly is not None:
            self.densities[self._ipoly] = value
            self.polygons[self._ipoly].set_color(self._density2color(value))
            # self._update_data()
            self._update_data_plot()
            self.canvas.draw()

    @property
    def error(self):
        return self._error

    @error.setter
    def error(self, value):
        """
        Callback when error slider is edited
        """
        self._error = value
        self._update_data_plot()

    @property
    def predicted(self):
        self._predicted = talwani.gz(self.x, self.z, self.model)
        if self.error > 0:
            self._predicted = utils.contaminate(self._predicted, self.error)
        return self._predicted

    @property
    def model(self):
        """
        The polygon model drawn as :class:`fatiando.mesher.Polygon` objects.
        """
        m = [Polygon(p.xy, {'density': d}, force_clockwise=True)
             for p, d in zip(self.polygons, self.densities)]
        return m

    def set_meassurement_points(self, x, z):
        self.x = x
        self.z = z
        self.modelax.set_xlim(self.x.min(), self.x.max())
        self.modelax.set_ylim(self.max_depth, self.min_depth) # y inverted axe
        self.predicted_line.remove()
        self.predicted_line, = self.dataax.plot(self.x, self.predicted, '-r')
        self._update_data_plot()

    def _figure_setup(self):
        self.dataax, self.modelax = self.fig.subplots(2, 1, sharex=True)
        self.dataax.set_ylabel("Gravity Anomaly [mGal]")
        self.dataax.set_ylim((-200, 200))
        self.dataax.grid(True)
        self.modelax.set_xlabel("x [m]")
        self.modelax.set_ylabel("z [m]")
        self.modelax.set_xlim(self.x.min(), self.x.max())
        self.modelax.set_ylim(self.min_depth, self.max_depth)
        self.modelax.grid(True)
        self.modelax.invert_yaxis()
        self.predicted_line, = self.dataax.plot(self.x, self.predicted, '-r')
        self.canvas.draw()

    def _init_markers(self):
        self._ivert = None
        self._ipoly = None
        self._lastevent = None
        self._drawing = False
        self._add_vertex = False
        self._xy = []
        self._drawing_plot = None
        self.background = None

    def _connect(self):
        """
        Connect the matplotlib events to their callback methods.
        """
        # Make the proper callback connections
        self.canvas.mpl_connect('button_press_event',
                                self._button_press_callback)
        self.canvas.mpl_connect('button_release_event',
                                self._button_release_callback)
        self.canvas.mpl_connect('motion_notify_event',
                                self._mouse_move_callback)

    def _density2color(self, density):
        """
        Map density values to colors using the given *cmap* attribute.

        Parameters:

        * density : 1d-array
            The density values of the model polygons

        Returns

        * colors : 1d-array
            The colors mapped to each density value (returned by a matplotlib
            colormap object.

        """
        dmin, dmax = self.density_range
        return self.cmap((density - dmin)/(dmax - dmin))

    def _make_polygon(self, vertices, density):
        """
        Create a polygon for drawing.

        Polygons are matplitlib.patches.Polygon objects for the fill and
        matplotlib.lines.Line2D for the contour.

        Parameters:

        * vertices : list of [x, z]
            List of the [x, z]  coordinate pairs of each vertex of the polygon
        * density : float
            The density of the polygon (used to set the color)

        Returns:

        * polygon, line
            The matplotlib Polygon and Line2D objects

        """
        poly = patches.Polygon(vertices, animated=False, alpha=0.9,
                               color=self._density2color(density))
        x, y = list(zip(*poly.xy))
        line = Line2D(x, y, **LINE_ARGS)
        return poly, line

    def _update_data_plot(self):
        """
        Update the predicted data plot in the *dataax*.

        Adjusts the xlim of the axes to fit the data.
        """
        predicted = self.predicted
        self.predicted_line.set_ydata(predicted)
        vmin = 1.2*min(predicted.min(), self.dmin)
        vmax = 1.2*max(predicted.max(), self.dmax)
        self.dataax.set_ylim(vmin, vmax)
        self.dataax.grid(True)
        self.canvas.draw()

    def _get_polygon_vertice_id(self, event):
        """
        Find out which vertex of which polygon the event happened in.

        If the click was inside a polygon (not on a vertex), identify that
        polygon.

        Returns:

        * p, v : int, int
            p: the index of the polygon the event happened in or None if
            outside all polygons.
            v: the index of the polygon vertex that was clicked or None if the
            click was not on a vertex.

        """
        distances = []
        indices = []
        for poly in self.polygons:
            x, y = poly.get_transform().transform(poly.xy).T
            d = numpy.sqrt((x - event.x)**2 + (y - event.y)**2)
            distances.append(d.min())
            indices.append(numpy.argmin(d))
        p = numpy.argmin(distances)
        if distances[p] >= self.epsilon:
            # Check if the event was inside a polygon
            x, y = event.x, event.y
            p, v = None, None
            for i, poly in enumerate(self.polygons):
                if poly.contains_point([x, y]):
                    p = i
                    break
        else:
            v = indices[p]
            last = len(self.polygons[p].xy) - 1
            if v == 0 or v == last:
                v = [0, last]
        return p, v

    def _add_new_vertex(self, event):
        """
        Add new vertex to polygon
        """
        vertices = self.polygons[self._ipoly].get_xy()
        x, y = vertices[:, 0], vertices[:, 1]
        # Compute the angle between the vectors to each pair of
        # vertices corresponding to each line segment of the polygon
        x1, y1 = x[:-1], y[:-1]
        x2, y2 = numpy.roll(x1, -1), numpy.roll(y1, -1)
        u = numpy.vstack((x1 - event.xdata, y1 - event.ydata)).T
        v = numpy.vstack((x2 - event.xdata, y2 - event.ydata)).T
        angle = numpy.arccos(numpy.sum(u*v, 1) /
                             numpy.sqrt(numpy.sum(u**2, 1)) /
                             numpy.sqrt(numpy.sum(v**2, 1)))
        position = angle.argmax() + 1
        x = numpy.hstack((x[:position], event.xdata, x[position:]))
        y = numpy.hstack((y[:position], event.ydata, y[position:]))
        new_vertices = numpy.vstack((x, y)).T
        return new_vertices

    def _button_press_callback(self, event):
        """
        What actions to perform when a mouse button is clicked
        """
        if event.inaxes != self.modelax:
            return
        if event.button == 1 and not self._drawing and self.polygons:
            self._lastevent = event
            if not self._add_vertex:
                for line, poly in zip(self.lines, self.polygons):
                    poly.set_animated(False)
                    line.set_animated(False)
                    line.set_color([0, 0, 0, 0])
                self.canvas.draw()
                # Find out if a click happened on a vertice
                # and which vertice of which polygon
                self._ipoly, self._ivert = self._get_polygon_vertice_id(event)
                if self._ipoly is not None:
                    # self.density_slider.set_val(self.densities[self._ipoly])
                    self.polygons[self._ipoly].set_animated(True)
                    self.lines[self._ipoly].set_animated(True)
                    self.lines[self._ipoly].set_color([0, 1, 0, 0])
                    self.canvas.draw()
                    self.background = self.canvas.copy_from_bbox(
                        self.modelax.bbox)
                    self.modelax.draw_artist(self.polygons[self._ipoly])
                    self.modelax.draw_artist(self.lines[self._ipoly])
                    self.canvas.blit(self.modelax.bbox)
            else:
                # If a polygon is selected, we will add a new vertex by
                # removing the polygon and inserting a new one with the extra
                # vertex.
                if self._ipoly is not None:
                    vertices = self._add_new_vertex(event)
                    density = self.densities[self._ipoly]
                    polygon, line = self._make_polygon(vertices, density)
                    self.polygons[self._ipoly].remove()
                    self.lines[self._ipoly].remove()
                    self.polygons.pop(self._ipoly)
                    self.lines.pop(self._ipoly)
                    self.polygons.insert(self._ipoly, polygon)
                    self.lines.insert(self._ipoly, line)
                    self.modelax.add_patch(polygon)
                    self.modelax.add_line(line)
                    self.lines[self._ipoly].set_color([0, 1, 0, 0])
                    self.canvas.draw()
                    # self._update_data()
                    self._update_data_plot()
        elif self._drawing:
            if event.button == 1:
                self._xy.append([event.xdata, event.ydata])
                self._drawing_plot.set_data(list(zip(*self._xy)))
                self.canvas.restore_region(self.background)
                self.modelax.draw_artist(self._drawing_plot)
                self.canvas.blit(self.modelax.bbox)
            elif event.button == 3:
                if len(self._xy) >= 3:
                    poly, line = self._make_polygon(self._xy, self.density)
                    self.polygons.append(poly)
                    self.lines.append(line)
                    self.densities.append(self.density)
                    self.modelax.add_patch(poly)
                    self.modelax.add_line(line)
                    self._drawing_plot.remove()
                    self._drawing_plot = None
                    self._xy = None
                    self._drawing = False
                    self._ipoly = len(self.polygons) - 1
                    self.lines[self._ipoly].set_color([0, 1, 0, 0])
                    self.dataax.set_title(self.instructions)
                    self.canvas.draw()
                    # self._update_data()
                    self._update_data_plot()

    def _button_release_callback(self, event):
        """
        Reset place markers on mouse button release
        """
        if event.inaxes != self.modelax:
            return
        if event.button != 1:
            return
        if self._add_vertex:
            self._add_vertex = False
        if self._ivert is None and self._ipoly is None:
            return
        self.background = None
        for line, poly in zip(self.lines, self.polygons):
            poly.set_animated(False)
            line.set_animated(False)
        self.canvas.draw()
        self._ivert = None
        # self._ipoly is only released when clicking outside
        # the polygons
        self._lastevent = None
        # self._update_data()
        self._update_data_plot()

    def _key_press_callback(self, event_key):
        """
        What to do when a key is pressed on the keyboard.
        """
        if event_key == 'd':
            if self._drawing and self._xy:
                self._xy.pop()
                if self._xy:
                    self._drawing_plot.set_data(list(zip(*self._xy)))
                else:
                    self._drawing_plot.set_data([], [])
                self.canvas.restore_region(self.background)
                self.modelax.draw_artist(self._drawing_plot)
                self.canvas.blit(self.modelax.bbox)
            elif self._ivert is not None:
                poly = self.polygons[self._ipoly]
                line = self.lines[self._ipoly]
                if len(poly.xy) > 4:
                    verts = numpy.atleast_1d(self._ivert)
                    poly.xy = numpy.array([xy for i, xy in enumerate(poly.xy)
                                           if i not in verts])
                    line.set_data(list(zip(*poly.xy)))
                    # self._update_data()
                    self._update_data_plot()
                    self.canvas.restore_region(self.background)
                    self.modelax.draw_artist(poly)
                    self.modelax.draw_artist(line)
                    self.canvas.blit(self.modelax.bbox)
                    self._ivert = None
            elif self._ipoly is not None:
                self.polygons[self._ipoly].remove()
                self.lines[self._ipoly].remove()
                self.polygons.pop(self._ipoly)
                self.lines.pop(self._ipoly)
                self.densities.pop(self._ipoly)
                self._ipoly = None
                self.canvas.draw()
                # self._update_data()
                self._update_data_plot()
        elif event_key == 'n':
            self._ivert = None
            self._ipoly = None
            for line, poly in zip(self.lines, self.polygons):
                poly.set_animated(False)
                line.set_animated(False)
                line.set_color([0, 0, 0, 0])
            self.canvas.draw()
            self.background = self.canvas.copy_from_bbox(self.modelax.bbox)
            self._drawing = True
            self._xy = []
            self._drawing_plot = Line2D([], [], **LINE_ARGS)
            self._drawing_plot.set_animated(True)
            self.modelax.add_line(self._drawing_plot)
            self.dataax.set_title(' | '.join([
                'left click: set vertice', 'right click: finish',
                'esc: cancel']))
            self.canvas.draw()
        elif event_key == 'escape':
            if self._add_vertex:
                self._add_vertex = False
            else:
                self._drawing = False
                self._xy = []
                if self._drawing_plot is not None:
                    self._drawing_plot.remove()
                    self._drawing_plot = None
                for line, poly in zip(self.lines, self.polygons):
                    poly.set_animated(False)
                    line.set_animated(False)
                    line.set_color([0, 0, 0, 0])
            self.canvas.draw()
        elif event_key == 'r':
            self.modelax.set_xlim(self.x.min(), self.x.max())
            self.modelax.set_ylim(self.min_depth, self.max_depth)
            self._update_data_plot()
        elif event_key == 'a':
            self._add_vertex = not self._add_vertex

    def _mouse_move_callback(self, event):
        """
        Handle things when the mouse move.
        """
        if event.inaxes != self.modelax:
            return
        if event.button != 1:
            return
        if self._ivert is None and self._ipoly is None:
            return
        if self._add_vertex:
            return
        x, y = event.xdata, event.ydata
        p = self._ipoly
        v = self._ivert
        if self._ivert is not None:
            self.polygons[p].xy[v] = x, y
        else:
            dx = x - self._lastevent.xdata
            dy = y - self._lastevent.ydata
            self.polygons[p].xy[:, 0] += dx
            self.polygons[p].xy[:, 1] += dy
        self.lines[p].set_data(list(zip(*self.polygons[p].xy)))
        self._lastevent = event
        self.canvas.restore_region(self.background)
        self.modelax.draw_artist(self.polygons[p])
        self.modelax.draw_artist(self.lines[p])
        self.canvas.blit(self.modelax.bbox)