from __future__ import division

from matplotlib import colors
import matplotlib.pyplot as plt
import numpy as np

from utils import xy_to_id, id_to_xy

class GUI(object):
    Graphics user interface for drawing the outputs a topology optimization
    def __init__(self, nelx, nely, title=""):
        """Initialize plot and plot the initial design"""
        plt.ion()  # Ensure that redrawing is possible
        self.fig, = plt.subplots() =, nely)).T, cmap='gray',
            interpolation='none', norm=colors.Normalize(vmin=-1, vmax=0))
        # self.fig.tight_layout()
        self.nelx, self.nely = nelx, nely

    def plot_force_arrows(self, f):
        """Add arrows to the plot for each force."""
        arrowprops = {"arrowstyle": "->", "connectionstyle": "arc3",
            "lw": "2", "color": 0}
        cmap ="hsv", f.shape[1] + 1)
        for load_i in range(f.shape[1]):
            nz = np.nonzero(f[:, load_i])
            arrowprops["color"] = cmap(load_i)
            for i in range(nz[0].shape[0]):
                x, y = id_to_xy(nz[0][i] // 2, self.nelx, self.nely)
                x = max(min(x, self.nelx - 1), 0)
                y = max(min(y, self.nely - 1), 0)
                z = int(nz[0][i] % 2)
                mag = -50 * f[nz[0][i], load_i]
      "", xy=(x, y), xycoords="data",
                    xytext = (0 if z else mag, mag if z else 0),
                    textcoords="offset points", arrowprops=arrowprops)

    def update(self, xPhys, title=None):
        """Plot to screen""", self.nely)).T)
        if title is not None: